From 0e3f52a4c0b071e8f297decce690a6ba2b6615ff Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 13 Aug 2026 19:18:22 -0500 Subject: [PATCH 01/43] feat(model_prices): add gemini-3.1-flash-lite-image Register Nano Banana 2 Lite on the unprefixed, gemini/, and vertex_ai/ keys so completion_cost and pass-through spend tracking no longer treat the model as unmapped --- ...odel_prices_and_context_window_backup.json | 95 +++++++ model_prices_and_context_window.json | 95 +++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 1 + ...ini_3_1_flash_lite_image_model_metadata.py | 242 ++++++++++++++++++ tests/test_litellm/test_utils.py | 2 + 5 files changed, 435 insertions(+) create mode 100644 tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1eb72c887b5..ca67d5d6844 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1eb72c887b5..ca67d5d6844 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -18141,6 +18141,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -19949,6 +19987,42 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_vision": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -38760,6 +38834,27 @@ "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_vision": true, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, 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 3aa41e18f1e..36fc98a1f09 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 @@ -1546,6 +1546,7 @@ def test_service_tier_fallback_pricing(): [ "gemini-3-pro-image-preview", "gemini-3.1-flash-image-preview", + "gemini-3.1-flash-lite-image", ], ) def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py new file mode 100644 index 00000000000..aa6f03a47ff --- /dev/null +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py @@ -0,0 +1,242 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm import completion_cost +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +UNPREFIXED = "gemini-3.1-flash-lite-image" +GEMINI = "gemini/gemini-3.1-flash-lite-image" +VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" +ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) + +INPUT_COST = 2.5e-07 +INPUT_COST_BATCHES = 1.25e-07 +OUTPUT_TEXT_COST = 1.5e-06 +OUTPUT_TEXT_COST_BATCHES = 7.5e-07 +OUTPUT_IMAGE_TOKEN_COST = 3e-05 +OUTPUT_COST_PER_1K_IMAGE = 0.0336 +INPUT_COST_PER_IMAGE = 0.00028 +CACHE_READ_COST = 2.5e-08 +MAX_INPUT_TOKENS = 65536 +MAX_OUTPUT_TOKENS = 4096 +TOKENS_PER_1K_IMAGE = 1120 + + +def _load(path: Path) -> dict: + with open(path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_gemini_3_1_flash_lite_image_is_registered(model: str): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["mode"] == "image_generation" + assert info["input_cost_per_token"] == INPUT_COST + assert info["input_cost_per_token_batches"] == INPUT_COST_BATCHES + assert info["output_cost_per_token"] == OUTPUT_TEXT_COST + assert info["output_cost_per_token_batches"] == OUTPUT_TEXT_COST_BATCHES + assert info["output_cost_per_image"] == OUTPUT_COST_PER_1K_IMAGE + assert info["output_cost_per_image_token"] == OUTPUT_IMAGE_TOKEN_COST + assert info["max_input_tokens"] == MAX_INPUT_TOKENS + assert info["max_output_tokens"] == MAX_OUTPUT_TOKENS + assert info["max_tokens"] == MAX_OUTPUT_TOKENS + assert info["supports_reasoning"] is False + assert info["supports_response_schema"] is False + assert info["supports_vision"] is True + for field in ("supports_web_search", "search_context_cost_per_query", "web_search_billing_unit"): + assert field not in info + + +def test_gemini_3_1_flash_lite_image_provider_specific_fields(): + cost_map = _load(MAIN_PATH) + + unprefixed = cost_map[UNPREFIXED] + assert unprefixed["litellm_provider"] == "vertex_ai-language-models" + assert unprefixed["cache_read_input_token_cost"] == CACHE_READ_COST + assert unprefixed["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert unprefixed["supports_function_calling"] is False + assert unprefixed["supports_prompt_caching"] is True + assert unprefixed["supports_pdf_input"] is True + assert unprefixed["supports_video_input"] is True + assert unprefixed["supported_modalities"] == ["text", "image", "video"] + + gemini = cost_map[GEMINI] + assert gemini["litellm_provider"] == "gemini" + assert gemini["supports_function_calling"] is True + assert gemini["supports_prompt_caching"] is False + assert "cache_read_input_token_cost" not in gemini + assert gemini["supported_modalities"] == ["text", "image"] + assert gemini["supported_output_modalities"] == ["text", "image"] + assert gemini["rpm"] == 1000 + assert gemini["tpm"] == 4000000 + assert gemini["input_cost_per_image"] == INPUT_COST_PER_IMAGE + + vertex = cost_map[VERTEX] + assert vertex["litellm_provider"] == "vertex_ai-language-models" + assert vertex["cache_read_input_token_cost"] == CACHE_READ_COST + assert vertex["input_cost_per_image"] == INPUT_COST_PER_IMAGE + assert vertex["supports_function_calling"] is False + assert vertex["supports_prompt_caching"] is True + + +def test_one_k_image_price_matches_official_token_math(): + assert TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST == OUTPUT_COST_PER_1K_IMAGE + assert TOKENS_PER_1K_IMAGE * INPUT_COST == INPUT_COST_PER_IMAGE + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_backup_matches_main(model: str): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_gemini_prefix_routes_to_gemini(): + routed_model, provider, _, _ = get_llm_provider(model=GEMINI) + assert routed_model == UNPREFIXED + assert provider == "gemini" + + +def test_vertex_prefix_routes_to_vertex(): + routed_model, provider, _, _ = get_llm_provider(model=VERTEX) + assert routed_model == UNPREFIXED + assert provider == "vertex_ai" + + +def test_text_token_cost(local_model_cost_map): + prompt_cost, text_completion_cost = cost_per_token(model=GEMINI, prompt_tokens=1000, completion_tokens=500) + assert prompt_cost == pytest.approx(1000 * INPUT_COST) + assert text_completion_cost == pytest.approx(500 * OUTPUT_TEXT_COST) + + +def test_completion_cost_bills_one_k_image(local_model_cost_map): + response = ModelResponse() + response.model = UNPREFIXED + response.usage = Usage( + prompt_tokens=7, + completion_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=7 + TOKENS_PER_1K_IMAGE, + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=TOKENS_PER_1K_IMAGE, text_tokens=0), + ) + billed = completion_cost( + completion_response=response, + model=UNPREFIXED, + custom_llm_provider="vertex_ai", + ) + expected = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 7 * INPUT_COST + assert billed == pytest.approx(expected) + + +def test_image_tokens_are_not_billed_as_text(local_model_cost_map): + usage = Usage( + completion_tokens=1345, + prompt_tokens=10, + total_tokens=1355, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=225, + rejected_prediction_tokens=None, + text_tokens=0, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None + ), + ) + + _, image_completion_cost = generic_cost_per_token( + model=UNPREFIXED, + usage=usage, + custom_llm_provider="vertex_ai", + ) + + expected_completion_cost = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 225 * OUTPUT_TEXT_COST + bugged_text_only_cost = 1345 * OUTPUT_TEXT_COST + assert image_completion_cost > bugged_text_only_cost * 2 + assert image_completion_cost == pytest.approx(expected_completion_cost) + + +def test_gemini_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = gemini_image_generation_cost_calculator(model=GEMINI, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + assert cost != OUTPUT_COST_PER_1K_IMAGE + + +def test_vertex_image_generation_uses_token_pricing(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), + ) + + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + + +def test_vertex_image_generation_falls_back_to_flat_image_price(local_model_cost_map): + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) + cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) + assert cost == pytest.approx(2 * OUTPUT_COST_PER_1K_IMAGE) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e9e6167fb9..ada053ee38c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4312,11 +4312,13 @@ class TestVertexEmbeddingEncodingFormat: "vertex_ai/gemini-3-pro-image-preview", "vertex_ai/gemini-3.1-flash-image", "vertex_ai/gemini-3.1-flash-image-preview", + "vertex_ai/gemini-3.1-flash-lite-image", "gemini/gemini-2.5-flash-image", "gemini/gemini-3-pro-image", "gemini/gemini-3-pro-image-preview", "gemini/gemini-3.1-flash-image", "gemini/gemini-3.1-flash-image-preview", + "gemini/gemini-3.1-flash-lite-image", ], ) def test_gemini_image_models_do_not_support_reasoning( From 34e692c903c9d75b52065b4093c7d80b7eb2e00b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:16:22 -0700 Subject: [PATCH 02/43] fix(batches): decode model-encoded output file id so completed batches book spend Adds e2e coverage for batches terminal state and cost write-back, failure paths, per-backend file content downloads, and two-gateway routing (LIT-5730). --- litellm/batches/batch_utils.py | 40 +- tests/e2e/batches/COVERAGE.md | 87 +++- tests/e2e/batches/batch_client.py | 29 +- tests/e2e/batches/capabilities.py | 10 + tests/e2e/batches/test_batches_e2e.py | 471 +++++++++++++++++- .../llm_nonconversational.yaml | 10 + .../test_litellm/batches/test_batch_utils.py | 37 ++ 7 files changed, 646 insertions(+), 38 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 0cf22d82ca6..6eb13d2cba7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -296,6 +296,32 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +def _provider_output_file_id(output_file_id: str) -> str: + """ + Resolve the file id the provider actually knows: unified ids yield their embedded + llm_output_file_id, model-encoded ids decode to the raw provider id, raw ids pass through. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_original_file_id, + ) + + unified_file_id: Final = _is_base64_encoded_unified_file_id(output_file_id) + if not unified_file_id: + return get_original_file_id(output_file_id) + try: + extracted: Final = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", + output_file_id, + e, + ) + return output_file_id + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", extracted) + return extracted + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", @@ -311,23 +337,11 @@ async def _fetch_batch_output_file_content( Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + file_id: Final = _provider_output_file_id(batch.output_file_id) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index ca48204962a..f02d4eb4fe4 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -1,9 +1,11 @@ # Batches Test Coverage Matrix Live e2e coverage of the Batches API over a real proxy, real provider keys, and -real cost. Synchronous tier only: a batch's completion window is 24h, so these -tests never wait for `completed`. They assert the proxy accepts, routes, retrieves, -cancels, and lists a batch; everything created is deleted on teardown. +real cost. Mostly synchronous tier: a batch's completion window is 24h, so the +lifecycle matrix never waits for `completed`. It asserts the proxy accepts, routes, +retrieves, cancels, and lists a batch; everything created is deleted on teardown. +The exception is `TestBatchTerminalState`, which covers the completed state and +cost write-back via a cross-run marker baton (design below). ## Provider x operation @@ -12,19 +14,26 @@ row per supported (provider, scenario) pair, so there are no skipped cells in th parametrized run. The batches suite never skips: missing provider creds or upstream failures are hard test failures (see `tests/e2e/CLAUDE.md`). -| Provider | create | retrieve | cancel | list | file backing | -|-----------|--------|----------|--------|------|--------------| -| OpenAI | yes | yes | yes | yes | OpenAI Files | -| Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | -| Bedrock | yes (unified only) | yes | no (limited upstream) | no | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | +| Provider | create | retrieve | cancel | list | content download | file backing | +|-----------|--------|----------|--------|------|------------------|--------------| +| OpenAI | yes | yes | yes | yes | yes (lifecycle + terminal output) | OpenAI Files | +| Azure | yes | yes | yes | yes | yes (byte-verbatim) | Azure Files | +| Vertex AI | yes | yes | yes | yes | yes (provider-transformed) | GCS (`gcs_bucket_name` / `GCS_BUCKET_NAME` on model) | +| Bedrock | yes (unified only) | yes | no (limited upstream) | no | yes (provider-transformed) | S3 (`s3_bucket_name` + `aws_*` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off -(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. +(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix; +flipping those gates is tracked in LIT-4774 and deliberately not part of this suite. Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only); `model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no model-less passthrough path. +`GET /v1/files/{id}/content` is exercised for the unified upload path per backend in +`test_unified_file_content_downloads`. Azure stores the JSONL verbatim, so its download +is asserted byte-equal to the upload. Vertex (GCS) and Bedrock (S3) transform lines at +upload time, so those assert a 200 with non-empty parseable JSON lines instead. Gemini +(non-Vertex) raises `NotImplementedError` for file content and has no cell here. + ## Routing scenarios (per `litellm/proxy/batches_endpoints/endpoints.py`) Each create-capable provider runs all four. The test asserts the returned file id @@ -71,11 +80,59 @@ File delete asserts `object=="file"` and `deleted==True`. | `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared ProxyClient; runtime batch model registration via /model/new; denial helpers | | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | -| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | +| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, per-backend content download, failure paths, second-hop routing, terminal state + cost | + +## Failure paths + +`TestBatchFailurePaths` pins the customer-facing error contracts. A malformed input +file is a 400 at upload naming the bad content. A JSONL line whose url contradicts +the batch endpoint passes create (providers validate asynchronously) and drives the +batch to `failed` with structured `errors.data` (code/line/message), a null +`output_file_id`, and a $0 spend row keyed `{batch_id}_batch_cost` (LIT-4852: a +failed batch books $0 instead of crashing cost tracking). Cancelling that failed +batch is a 409 naming the terminal status. A file id encoded for one deployment wins +over a conflicting `model` param on create: the batch routes and re-encodes by the +file's embedded model (foreign-id precedence). + +## Second hop (two chained gateways) + +`TestBatchSecondHop` registers a `litellm_proxy/` deployment pointing at +the proxy's own base URL with a freshly minted virtual key, so unified upload and +create traverse gateway -> gateway -> OpenAI (LIT-5347, PR #36240). The pin: +`target_model_names` is rewritten to the inner deployment on the second hop and the +nested managed ids round-trip retrieve. This self-chaining only needs the proxy to +reach its own `PROXY_BASE_URL`, which holds both locally and on the e2e stage. + +## Terminal state + cost write-back (cross-run marker baton) + +The 24h completion window rules out submit-and-wait inside one run, so +`TestBatchTerminalState` amortizes across runs. Each run submits a 1-line marker +batch (stable metadata key/value plus a per-run field) and deliberately never +cancels or deletes it or its input file: the marker is the baton the next run picks +up (OpenAI files expire on their own after ~30 days). Polling is list-only, up to 5 +minutes, because retrieving a non-terminal batch books a $0 spend row whose +request_id then blocks the later real-cost row (`skip_duplicates`); the single +retrieve happens only once a completed marker exists. The assertion target is the +newest completed marker from ANY run: run-scoped deployment names mean the list +re-encodes prior-run batches under new encoded ids, so their spend keys are fresh +and a prior-run marker is billable by this run. On the 6h stage cadence the full +assertions are therefore deterministic from run 2 onward. On a cold start (no +completed marker within the poll budget) the test passes on the submission +assertions alone: a documented vacuous pass, not a skip. Markers aged past the 24h +window (25h-73h band, within the newest 100-item list page) must be terminal. + +The cost assertion is the LIT-5730 headline: retrieving a completed model-encoded +batch must write a positive spend row with call_type `aretrieve_batch` and token +usage. Before the fix in `litellm/batches/batch_utils.py`, the retrieve endpoint +re-encoded the response's `output_file_id` in place before the queued logging +worker ran, the worker sent that encoded id to OpenAI, got a 404, and the spend row +never landed. ## Out of scope (intentionally) -Driving a batch to `completed`, cost tracking on completion, and the DB write-back -are not covered here; the 24h window makes them unfit for a synchronous gate. That -logic belongs in a DI-stubbed proxy integration test under `tests/test_litellm/proxy/` -where the provider client is injected to return `completed` deterministically. +Unified (managed) batch cost is owned by the hourly `CheckBatchCost` poller, and a +terminal DB status short-circuits retrieve for those ids, so the terminal-state cell +uses the encoded path; poller timing does not fit an e2e gate and belongs in a +DI-stubbed proxy integration test under `tests/test_litellm/proxy/`. Bedrock +cancel/list stay gated pending LIT-4774. Gemini (non-Vertex) file content raises +`NotImplementedError` upstream and is not a coverage cell. diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..21a56f3398f 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -44,6 +44,17 @@ class FileList(BaseModel): data: list[FileObject] = [] +class BatchErrorItem(BaseModel): + code: str | None = None + line: int | None = None + message: str | None = None + + +class BatchErrorList(BaseModel): + object: str | None = None + data: list[BatchErrorItem] = [] + + class BatchObject(BaseModel): id: str object: str | None = None @@ -51,6 +62,9 @@ class BatchObject(BaseModel): endpoint: str | None = None input_file_id: str | None = None output_file_id: str | None = None + error_file_id: str | None = None + errors: BatchErrorList | None = None + metadata: dict[str, str] | None = None completion_window: str | None = None created_at: int | None = None model: str | None = None @@ -72,12 +86,18 @@ class BatchCreateBody(BaseModel): endpoint: str = "/v1/chat/completions" completion_window: str = "24h" model: str | None = None + metadata: dict[str, str] | None = None class ModelQuery(BaseModel): model: str | None = None +class BatchListQuery(BaseModel): + model: str | None = None + limit: int | None = None + + def is_model_access_denied(resp: StreamingResponse) -> bool: """True if the proxy rejected the call because the key may not access the model.""" return resp.status_code == 403 and "key_model_access_denied" in resp.body @@ -168,12 +188,17 @@ class BatchClient: ) def list_batches( - self, *, key: str, provider: str | None = None + self, + *, + key: str, + provider: str | None = None, + model: str | None = None, + limit: int | None = None, ) -> Result[BatchList]: return self.proxy.transport.get( _batches_path(provider), headers=self.proxy.transport.bearer(key), - params=NoBody(), + params=BatchListQuery(model=model, limit=limit), response_type=BatchList, ) diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 3988fb5e7e1..ce1f68184a7 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -210,6 +210,16 @@ def is_model_encoded_id(id_str: str) -> bool: return False +def decoded_model_from_id(id_str: str) -> str | None: + """Deployment name embedded in a model-encoded file/batch id, or None.""" + for prefix in ("file-", "batch_"): + if id_str.startswith(prefix): + decoded = _b64_decode(id_str[len(prefix) :]) + if decoded.startswith("litellm:") and ";model," in decoded: + return decoded.split(";model,", 1)[1].split(";")[0] + return None + + def matches_id_shape(shape: IdShape, id_str: str) -> bool: if shape == "managed": return is_managed_id(id_str) diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..75b3a6cd758 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -1,11 +1,12 @@ """Live e2e for the Batches API across every provider LiteLLM supports. -Synchronous tier only: a batch's completion window is 24h, so these never wait for -"completed". Each case uploads a tiny JSONL, creates the batch through one of the -four routing scenarios, asserts it was accepted (non-terminal status) and routed to -the right provider, then retrieves / cancels / lists where the provider supports it. -Everything created is deleted on teardown. Completion + cost tracking are out of -scope here (see COVERAGE.md). +Mostly synchronous tier: a batch's completion window is 24h, so the lifecycle +matrix never waits for "completed". Each case uploads a tiny JSONL, creates the +batch through one of the four routing scenarios, asserts it was accepted +(non-terminal status) and routed to the right provider, then retrieves / cancels / +lists where the provider supports it. Everything created is deleted on teardown. +The exception is TestBatchTerminalState, which carries completed-state + cost +write-back coverage via a cross-run marker baton (design in COVERAGE.md). Routing signal: for provider_fallback the raw batch id discriminates the provider; for the encoded/unified/model_param scenarios the proxy re-encodes the id, so the @@ -22,8 +23,9 @@ from datetime import datetime, timedelta, timezone from typing import Callable import pytest +from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import PROXY_BASE_URL, unique_marker from batch_client import ( UPLOAD_FILENAME, @@ -40,9 +42,12 @@ from capabilities import ( CAPABILITIES, FILE_ID_SHAPE, OPENAI_BATCH_MODEL, + PROVIDERS, Capability, + Provider, batch_model_name, coverage_cells_for_lifecycle, + decoded_model_from_id, is_managed_id, matches_id_shape, raw_id_matches_provider, @@ -475,9 +480,22 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( OPENAI_FILE_CONTENT_BACKEND = "gpt-4o-mini" +FILE_CONTENT_CELLS = { + "azure": "llm.files.azure_openai.content.nonstream.works", + "vertex_ai": "llm.files.vertex.content.nonstream.works", + "bedrock": "llm.files.bedrock.content.nonstream.works", +} +BYTE_FIDELITY_CONTENT_PROVIDERS = frozenset({"azure"}) + class TestBatchFileContent: - """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes.""" + """GET /v1/files/{id}/content returns the uploaded batch JSONL bytes. + + Azure stores the upload verbatim, so its download is asserted byte-equal. + Vertex (GCS) and Bedrock (S3) transform each JSONL line into the provider's + request format at upload time, so their downloads assert 200 plus non-empty + parseable JSON lines instead of byte equality. + """ @pytest.mark.covers( "llm.files.openai.content.nonstream.works", @@ -521,6 +539,62 @@ class TestBatchFileContent: "downloaded file content must match the uploaded JSONL bytes" ) + @pytest.mark.parametrize( + "provider", + [ + pytest.param( + p, + id=p.name, + marks=pytest.mark.covers( + FILE_CONTENT_CELLS[p.name], exercised_on=["files"] + ), + ) + for p in PROVIDERS + if p.name in FILE_CONTENT_CELLS + ], + ) + def test_unified_file_content_downloads( + self, + provider: Provider, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + payload = render_jsonl(provider.raw_model) + file = unwrap( + client.upload_file( + content=payload, + form=FileUploadForm(purpose="batch", target_model_names=provider.model), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert_file_object(file, provider=provider.name) + assert is_managed_id(file.id), ( + f"{provider.name}: unified upload must return a managed file id, got {file.id!r}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"{provider.name}: file content must be 200, " + f"got {downloaded.status_code}: {downloaded.body[:300]}" + ) + body = downloaded.body.strip() + assert body, f"{provider.name}: file content download returned an empty body" + if provider.name in BYTE_FIDELITY_CONTENT_PROVIDERS: + assert body == payload.decode().strip(), ( + f"{provider.name}: downloaded content must match the uploaded JSONL bytes" + ) + else: + for line in body.splitlines(): + assert json.loads(line), ( + f"{provider.name}: content line is not JSON: {line[:200]}" + ) + class TestOpenAIFiles: """GET /v1/files (list) and GET /v1/files/{id} (retrieve) over the OpenAI route. @@ -866,3 +940,384 @@ class TestHostedVllmBatch: f"hosted_vllm batch has non-transitional status {batch.status!r}" ) assert_batch_object(batch) + + +BATCH_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"}) +FAILED_BATCH_POLL_SECONDS = 120.0 +FAILED_BATCH_POLL_INTERVAL_SECONDS = 5.0 + +AZURE_BATCH_RAW_MODEL = next(p.raw_model for p in PROVIDERS if p.name == "azure") + + +def _mismatched_endpoint_jsonl(model: str) -> bytes: + line = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": model, "input": "ping"}, + } + return (json.dumps(line) + "\n").encode() + + +def _poll_until_terminal(client: BatchClient, batch_id: str, key: str) -> BatchObject: + deadline = time.monotonic() + FAILED_BATCH_POLL_SECONDS + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + while fetched.status not in BATCH_TERMINAL_STATUSES and time.monotonic() < deadline: + time.sleep(FAILED_BATCH_POLL_INTERVAL_SECONDS) + fetched = retrieve_batch(client, batch_id, key=key, provider=None) + return fetched + + +class TestBatchFailurePaths: + """Customer-facing failure contracts for /v1/batches. + + A malformed input file is rejected at upload with a 400 naming the bad + content. A JSONL line whose url contradicts the batch endpoint is accepted + at create (providers validate asynchronously) and drives the batch to + "failed" with structured per-line errors, a null output_file_id, and a + zero-cost spend row (LIT-4852: a failed batch must book $0, not crash cost + tracking). Cancelling that already-failed batch returns a 409 naming the + terminal status. A file id encoded for one deployment wins over a + conflicting model param on create: the batch routes (and re-encodes) by the + file's embedded model, pinning that precedence. + """ + + @pytest.mark.covers( + "llm.batches.openai.malformed_jsonl.nonstream.works", + exercised_on=["files"], + ) + def test_malformed_jsonl_upload_rejected( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + result = client.upload_file( + content=b"this is not json\n", + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=resources.key(), + ) + match result: + case UnknownApiError(status_code=400, body=body): + assert "json" in body.lower(), ( + f"400 must name the malformed JSONL so users can fix the file, got: {body[:300]}" + ) + case _: + pytest.fail(f"malformed JSONL upload must be rejected with a 400, got: {result}") + + @pytest.mark.covers( + "llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works", + "llm.batches.openai.cancel_terminal.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_endpoint_mismatch_fails_batch_and_cancel_conflicts( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=_mismatched_endpoint_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + + fetched = _poll_until_terminal(client, batch.id, key) + assert fetched.status == "failed", ( + f"endpoint-mismatched batch must fail, got {fetched.status!r}" + ) + assert fetched.output_file_id is None, ( + f"failed batch must have no output file, got {fetched.output_file_id!r}" + ) + assert fetched.errors is not None and fetched.errors.data, ( + "failed batch must surface structured errors so users can fix the JSONL" + ) + first_error = fetched.errors.data[0] + assert first_error.message, "batch error item has no message" + assert first_error.code, "batch error item has no code" + + rows = client.proxy.poll_logs_for_request_id(f"{fetched.id}_batch_cost") + assert rows, ( + f"failed batch {fetched.id} wrote no spend row; retrieve must book $0 (LIT-4852)" + ) + assert all((row.spend or 0) == 0 for row in rows), ( + f"failed batch must cost $0, got {[(r.request_id, r.spend) for r in rows]}" + ) + assert rows[0].call_type == "aretrieve_batch", ( + f"batch cost row call_type={rows[0].call_type!r}" + ) + + conflict = client.cancel_batch(batch.id, key=key) + match conflict: + case UnknownApiError(status_code=409, body=body): + assert "failed" in body.lower(), ( + f"409 must name the terminal status blocking the cancel, got: {body[:300]}" + ) + case _: + pytest.fail(f"cancel of a failed batch must return a 409 conflict, got: {conflict}") + + @pytest.mark.covers( + "llm.batches.openai.foreign_file_id.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_foreign_encoded_file_id_routes_by_file_model( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl(AZURE_BATCH_RAW_MODEL), + form=FileUploadForm(purpose="batch"), + model=AZURE_BATCH_MODEL, + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert decoded_model_from_id(file.id) == AZURE_BATCH_MODEL, ( + f"upload did not encode the azure deployment into the file id: {file.id!r}" + ) + + created = client.create_batch( + body=BatchCreateBody(input_file_id=file.id, model=OPENAI_BATCH_MODEL), key=key + ) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert decoded_model_from_id(batch.id) == AZURE_BATCH_MODEL, ( + "create with a foreign encoded file id must route by the file's embedded model, " + f"but the batch id encodes {decoded_model_from_id(batch.id)!r} " + f"(model param was {OPENAI_BATCH_MODEL!r})" + ) + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "retrieved foreign-file batch has no status" + + +class TestBatchSecondHop: + """Two-proxy batch routing: a litellm_proxy deployment chained to the gateway + itself (LIT-5347, PR #36240). + + The hop deployment's litellm_params point litellm_proxy/ at this + gateway's own base URL with a freshly minted virtual key, so the unified + upload and batch create traverse gateway -> gateway -> OpenAI. The regression + this pins: target_model_names must be rewritten to the inner deployment on + the second hop and the nested managed ids must round-trip retrieve. + """ + + @pytest.mark.covers( + "llm.batches.openai.second_hop.nonstream.works", + exercised_on=["batches", "files"], + ) + def test_unified_create_and_retrieve_via_chained_gateway( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + hop_name = batch_model_name("openai-batch-hop") + model_id = client.create_model( + hop_name, + LiteLLMParamsBody( + model=f"litellm_proxy/{OPENAI_BATCH_MODEL}", + api_base=PROXY_BASE_URL, + api_key=key, + ), + ) + resources.defer(lambda: client.delete_model(model_id)) + + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch", target_model_names=hop_name), + key=key, + ) + ) + resources.defer(quietly(lambda: client.delete_file(file.id, key=key))) + assert is_managed_id(file.id), ( + f"second-hop unified upload must return a managed file id, got {file.id!r}" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(quietly(lambda: client.cancel_batch(batch.id, key=key))) + + assert is_managed_id(batch.id), ( + f"second-hop create must return a managed batch id, got {batch.id!r}" + ) + assert batch.status in CREATED_BATCH_STATUSES, ( + f"second-hop batch has non-transitional status {batch.status!r}" + ) + assert_batch_object(batch) + + fetched = retrieve_batch(client, batch.id, key=key, provider=None) + assert fetched.id == batch.id + assert fetched.status, "second-hop retrieve returned no status" + + +class BatchOutputBody(BaseModel): + choices: list[object] = [] + + +class BatchOutputResponse(BaseModel): + status_code: int | None = None + body: BatchOutputBody | None = None + + +class BatchOutputLine(BaseModel): + response: BatchOutputResponse + + +TERMINAL_MARKER_KEY = "litellm_e2e_suite" +TERMINAL_MARKER_VALUE = "batches-terminal-baton" +TERMINAL_POLL_SECONDS = 300.0 +TERMINAL_POLL_INTERVAL_SECONDS = 10.0 +TERMINAL_LIST_LIMIT = 100 +TERMINAL_BAND_MIN_AGE_SECONDS = 25 * 3600 +TERMINAL_BAND_MAX_AGE_SECONDS = 73 * 3600 + + +def _marker_batches(client: BatchClient, key: str) -> list[BatchObject]: + listed = unwrap( + client.list_batches(key=key, model=OPENAI_BATCH_MODEL, limit=TERMINAL_LIST_LIMIT) + ) + return [ + b + for b in listed.data + if (b.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE + ] + + +def _await_completed_marker( + client: BatchClient, key: str +) -> tuple[BatchObject | None, list[BatchObject]]: + deadline = time.monotonic() + TERMINAL_POLL_SECONDS + while True: + markers = _marker_batches(client, key) + completed = max( + (b for b in markers if b.status == "completed"), + key=lambda b: b.created_at or 0, + default=None, + ) + if completed is not None or time.monotonic() >= deadline: + return completed, markers + time.sleep(TERMINAL_POLL_INTERVAL_SECONDS) + + +def _assert_aged_markers_terminal(markers: list[BatchObject]) -> None: + now = time.time() + stuck = [ + b + for b in markers + if b.created_at is not None + and TERMINAL_BAND_MIN_AGE_SECONDS <= now - b.created_at <= TERMINAL_BAND_MAX_AGE_SECONDS + and b.status not in BATCH_TERMINAL_STATUSES + ] + assert not stuck, ( + "marker batches past their 24h completion window must be terminal; stuck: " + f"{[(b.id, b.status, b.created_at) for b in stuck]}" + ) + + +class TestBatchTerminalState: + """Terminal state + cost write-back via a cross-run marker baton. + + Each run submits a 1-line marker batch (stable metadata key/value plus a + per-run field) and never cancels or deletes it: the marker is the baton the + next run picks up. Polling is list-only for up to 5 minutes because a + retrieve of a non-terminal batch books a $0 spend row whose request_id then + blocks the real-cost row (skip_duplicates); the single retrieve happens only + once a completed marker exists. The assertion target is the newest completed + marker from ANY run, so on the 6h stage cadence the full assertions are + deterministic from run 2 onward. On a cold start (no marker has ever + completed within the poll budget) the test passes on the submission + assertions alone: that is a documented vacuous pass, not a skip, and this + run's marker becomes the next run's target. Markers aged past OpenAI's 24h + completion window (25h-73h band, within the newest list page) must be + terminal. The cost assertion is the LIT-5730 headline: retrieving a + completed model-encoded batch must write a positive spend row keyed + {batch_id}_batch_cost; before the fix the logging worker fetched the + re-encoded output_file_id, 404d, and the row never landed. + """ + + @pytest.mark.covers( + "llm.batches.openai.terminal_state.nonstream.works", + "llm.batches.openai.terminal_state.nonstream.cost_logged", + exercised_on=["batches", "files"], + ) + def test_completed_batch_downloads_output_and_books_cost( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=render_jsonl("gpt-4o-mini"), + form=FileUploadForm(purpose="batch"), + model=OPENAI_BATCH_MODEL, + key=key, + ) + ) + created = client.create_batch( + body=BatchCreateBody( + input_file_id=file.id, + metadata={ + TERMINAL_MARKER_KEY: TERMINAL_MARKER_VALUE, + "run": unique_marker(), + }, + ), + key=key, + ) + require_successful_call(created) + submitted = BatchObject.model_validate_json(created.body) + assert submitted.status in CREATED_BATCH_STATUSES, ( + f"marker batch has non-transitional status {submitted.status!r}" + ) + assert (submitted.metadata or {}).get(TERMINAL_MARKER_KEY) == TERMINAL_MARKER_VALUE, ( + f"create dropped the marker metadata: {submitted.metadata!r}" + ) + + completed, markers = _await_completed_marker(client, key) + _assert_aged_markers_terminal(markers) + if completed is None: + return + + fetched = retrieve_batch(client, completed.id, key=key, provider=None) + assert fetched.status == "completed", ( + f"listed-completed marker retrieved as {fetched.status!r}" + ) + assert fetched.output_file_id, "completed batch has no output_file_id" + + downloaded = client.proxy.transport.download( + f"/v1/files/{fetched.output_file_id}/content", + headers=client.proxy.transport.bearer(key), + ) + assert downloaded.status_code == 200, ( + f"output content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + first_line = BatchOutputLine.model_validate_json(downloaded.body.strip().splitlines()[0]) + assert first_line.response.status_code == 200, ( + f"batch output line reports failure: {downloaded.body[:400]}" + ) + assert first_line.response.body is not None and first_line.response.body.choices, ( + "batch output line has no choices" + ) + + rows = client.proxy.poll_logs_for_request_id( + f"{fetched.id}_batch_cost", + predicate=lambda found: any((row.spend or 0) > 0 for row in found), + ) + priced = [row for row in rows if (row.spend or 0) > 0] + assert priced, ( + f"completed batch {fetched.id} wrote no positive-cost spend row under " + f"request_id {fetched.id}_batch_cost; cost write-back is broken (LIT-5730)" + ) + cost_row = priced[0] + assert cost_row.call_type == "aretrieve_batch", ( + f"batch cost row call_type={cost_row.call_type!r}" + ) + assert (cost_row.total_tokens or 0) > 0, ( + f"batch cost row has no token usage: {cost_row.total_tokens!r}" + ) diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bb7169509eb..3de462683d2 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -24,6 +24,13 @@ - {id: llm.batches.hosted_vllm.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible batch create"} - {id: llm.batches.openai.key_model_access_denied.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Key model restriction 403 on upload/create"} - {id: llm.batches.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.18 / LIT-4778", rationale: "Missing input_file_id and invalid batch id rejected"} +- {id: llm.batches.openai.terminal_state.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "A batch actually reaches completed and its output file downloads through GET /v1/files/{id}/content with per-line provider responses"} +- {id: llm.batches.openai.terminal_state.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_batches_e2e.py / LIT-5730", fail_before_fix: proven, rationale: "Retrieving a completed model-encoded batch writes a positive spend row keyed {batch_id}_batch_cost (pins LIT-4852/LIT-5666; before the fix the logging worker 404d fetching the re-encoded output_file_id and the row was never written)"} +- {id: llm.batches.openai.malformed_jsonl.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Uploading a non-JSON batch file is rejected with a 400 naming the bad line"} +- {id: llm.batches.openai.jsonl_endpoint_mismatch.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "JSONL line url that contradicts the batch endpoint drives the batch to failed with structured errors, retrieve stays clean, and the terminal retrieve books a zero-cost spend row (LIT-4852)"} +- {id: llm.batches.openai.cancel_terminal.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Cancelling an already-terminal batch returns a 409 conflict naming the terminal status"} +- {id: llm.batches.openai.foreign_file_id.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "Create with one deployment's encoded file id and a conflicting model param routes by the file's embedded model; the returned batch id pins that precedence"} +- {id: llm.batches.openai.second_hop.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5347", rationale: "A litellm_proxy deployment chained to the gateway itself preserves target_model_names through nested unified ids; upload, create, and retrieve work over the two-hop chain (PR #36240)"} - {id: llm.files.openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "openai_files_endpoints/files_endpoints.py:46", rationale: "File upload returns OpenAIFileObject"} - {id: llm.files.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.16 / LIT-4778", rationale: "File upload without purpose rejected"} - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} @@ -36,6 +43,9 @@ - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} +- {id: llm.files.azure_openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on an Azure unified file returns the uploaded JSONL bytes verbatim"} +- {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} +- {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index ebe093c591c..08cdf945b80 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -800,6 +800,43 @@ async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monke assert captured["custom_llm_provider"] == "vertex_ai" +@pytest.mark.asyncio +async def test_output_file_content_model_encoded_file_id_decoded_to_provider_id(monkeypatch): + import litellm.files.main as files_main + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + encoded_id = encode_file_id_with_model("file-Y3FHrMpi7uCkDpY6fgWGeR", "my-batch-model") + + await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="openai") + + assert captured["file_id"] == "file-Y3FHrMpi7uCkDpY6fgWGeR" + assert captured["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_output_file_content_raw_openai_file_id_passes_through(monkeypatch): + import litellm.files.main as files_main + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + + await bu._fetch_batch_output_file_content(_batch("file-abc123"), custom_llm_provider="openai") + + assert captured["file_id"] == "file-abc123" + + def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens): return { "request": { From 86efa2bcfde555e519140f6d9721e9570aada2d2 Mon Sep 17 00:00:00 2001 From: longwind48 Date: Fri, 21 Aug 2026 17:27:28 +0800 Subject: [PATCH 03/43] feat(bedrock): serve gpt-5.6 cross-region inference profiles on bedrock runtime GPT-5.6 Sol, Terra and Luna reached the bedrock-runtime data plane on 2026-08-17, separately from the existing bedrock-mantle path. On runtime they are served only through cross-region inference profiles, so bedrock/us.openai.gpt-5.6-* had no cost map entry and fell through to the Invoke route, which rewrites the token cap to max_tokens and is rejected as unsupported_parameter on both /v1/chat/completions and /v1/responses. Register the Geo and Global profiles as bedrock_converse so routing reaches Converse, which AWS documents and serves for these models, and price each profile from its own published rate table. No bare key: the control plane reports inferenceTypesSupported INFERENCE_PROFILE with no on-demand throughput, so a bare id is not invocable. Declare the published cache-read and cache-write rates. Bedrock rejects an explicit cachePoint block for these models, so supports_prompt_caching stays off, but it caches long prefixes implicitly and reports the cache tokens in usage either way. Without the cost fields a cache-read turn bills only its uncached tokens: measured against live Bedrock, a 15609-token cached prefix came to $0.000176 instead of $0.00876095. Clients that resend a long prefix every turn are the worst affected. Reasoning stays unadvertised. Converse rejects the Anthropic-shaped thinking block LiteLLM sends for reasoning_effort; the shape these models accept is additionalModelRequestFields {"reasoning": {"effort": ...}}, which needs a transform change tracked by #34105. Advertising it without that change is what made the earlier attempt in #37307 fail. --- ...odel_prices_and_context_window_backup.json | 150 ++++++++++ model_prices_and_context_window.json | 150 ++++++++++ ..._cross_region_inference_profile_mapping.py | 258 +++++++++++++++++- 3 files changed, 557 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 91c10d13e8e..be8e6da5f59 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48545,6 +48545,156 @@ "supports_tool_choice": true, "supports_vision": true }, + "us.openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-sol": { + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-terra": { + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-terra": { + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-luna": { + "input_cost_per_token": 2.2e-07, + "input_cost_per_token_above_272k_tokens": 4.4e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-07, + "cache_read_input_token_cost": 2.2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, + "output_cost_per_token": 1.32e-06, + "output_cost_per_token_above_272k_tokens": 1.98e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-luna": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 91c10d13e8e..be8e6da5f59 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48545,6 +48545,156 @@ "supports_tool_choice": true, "supports_vision": true }, + "us.openai.gpt-5.6-sol": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.375e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-sol": { + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-terra": { + "input_cost_per_token": 2.2e-06, + "input_cost_per_token_above_272k_tokens": 4.4e-06, + "cache_creation_input_token_cost": 2.75e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, + "output_cost_per_token": 1.32e-05, + "output_cost_per_token_above_272k_tokens": 1.98e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-terra": { + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "us.openai.gpt-5.6-luna": { + "input_cost_per_token": 2.2e-07, + "input_cost_per_token_above_272k_tokens": 4.4e-07, + "cache_creation_input_token_cost": 2.75e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.5e-07, + "cache_read_input_token_cost": 2.2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, + "output_cost_per_token": 1.32e-06, + "output_cost_per_token_above_272k_tokens": 1.98e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "global.openai.gpt-5.6-luna": { + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "cache_read_input_token_cost": 5.5e-07, diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 3a27f3ed002..22aba59fb5d 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -1,13 +1,132 @@ """Test Bedrock cross-region inference profile model mapping""" +import json import os import sys +from functools import lru_cache +from pathlib import Path +from typing import NamedTuple + +import pytest sys.path.insert(0, os.path.abspath("../../../..")) +import litellm +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.utils import _get_model_info_helper from litellm.cost_calculator import completion_cost -from litellm.types.utils import ModelResponse, Usage, Choices, Message +from litellm.types.utils import ( + Choices, + Message, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Resolve models against this checkout's cost map instead of the network-fetched + ``main`` copy, which lags this branch until merge.""" + original_converse_models = set(litellm.bedrock_converse_models) + 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() + try: + litellm.bedrock_converse_models.update( + key + for key, value in litellm.model_cost.items() + if isinstance(value, dict) + and value.get("litellm_provider") == "bedrock_converse" + ) + yield + finally: + litellm.bedrock_converse_models.clear() + litellm.bedrock_converse_models.update(original_converse_models) + litellm.get_model_info.cache_clear() + + +class GptProfile(NamedTuple): + model_id: str + input_cost: float + input_cost_above_272k: float + cache_write: float + cache_write_above_272k: float + cache_read: float + cache_read_above_272k: float + output_cost: float + output_cost_above_272k: float + + +GPT_5_6_PROFILES = [ + GptProfile( + model_id="us.openai.gpt-5.6-sol", + input_cost=5.5e-06, input_cost_above_272k=1.1e-05, + cache_write=6.875e-06, cache_write_above_272k=1.375e-05, + cache_read=5.5e-07, cache_read_above_272k=1.1e-06, + output_cost=3.3e-05, output_cost_above_272k=4.95e-05, + ), + GptProfile( + model_id="global.openai.gpt-5.6-sol", + input_cost=5e-06, input_cost_above_272k=1e-05, + cache_write=6.25e-06, cache_write_above_272k=1.25e-05, + cache_read=5e-07, cache_read_above_272k=1e-06, + output_cost=3e-05, output_cost_above_272k=4.5e-05, + ), + GptProfile( + model_id="us.openai.gpt-5.6-terra", + input_cost=2.2e-06, input_cost_above_272k=4.4e-06, + cache_write=2.75e-06, cache_write_above_272k=5.5e-06, + cache_read=2.2e-07, cache_read_above_272k=4.4e-07, + output_cost=1.32e-05, output_cost_above_272k=1.98e-05, + ), + GptProfile( + model_id="global.openai.gpt-5.6-terra", + input_cost=2e-06, input_cost_above_272k=4e-06, + cache_write=2.5e-06, cache_write_above_272k=5e-06, + cache_read=2e-07, cache_read_above_272k=4e-07, + output_cost=1.2e-05, output_cost_above_272k=1.8e-05, + ), + GptProfile( + model_id="us.openai.gpt-5.6-luna", + input_cost=2.2e-07, input_cost_above_272k=4.4e-07, + cache_write=2.75e-07, cache_write_above_272k=5.5e-07, + cache_read=2.2e-08, cache_read_above_272k=4.4e-08, + output_cost=1.32e-06, output_cost_above_272k=1.98e-06, + ), + GptProfile( + model_id="global.openai.gpt-5.6-luna", + input_cost=2e-07, input_cost_above_272k=4e-07, + cache_write=2.5e-07, cache_write_above_272k=5e-07, + cache_read=2e-08, cache_read_above_272k=4e-08, + output_cost=1.2e-06, output_cost_above_272k=1.8e-06, + ), +] + + +@lru_cache(maxsize=1) +def _packaged_cost_map(): + """The map litellm actually resolves against, for fields ModelInfoBase drops.""" + path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json" + return json.loads(path.read_text()) + + +def _bedrock_response(model, usage): + return ModelResponse( + id="test", + created=1234567890, + model=model, + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="OK", role="assistant"), + ) + ], + usage=usage, + ) def test_bedrock_cross_region_inference_profile_mapping(): @@ -52,3 +171,140 @@ def test_proxy_cost_calculation_scenario(): ) expected_cost = (100 * 8e-07) + (50 * 4e-06) assert cost == expected_cost + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_map): + """GPT-5.6 is served by Converse on bedrock-runtime, never by Invoke.""" + assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse" + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map): + """Geo and Global profiles carry their own published rates, per context tier.""" + model_info = _get_model_info_helper( + model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" + ) + + assert model_info["litellm_provider"] == "bedrock_converse" + assert model_info["mode"] == "chat" + assert model_info["max_input_tokens"] == 1000000 + assert model_info["input_cost_per_token"] == profile.input_cost + assert ( + model_info["input_cost_per_token_above_272k_tokens"] + == profile.input_cost_above_272k + ) + assert model_info["output_cost_per_token"] == profile.output_cost + assert ( + model_info["output_cost_per_token_above_272k_tokens"] + == profile.output_cost_above_272k + ) + assert model_info["cache_creation_input_token_cost"] == profile.cache_write + assert ( + model_info["cache_creation_input_token_cost_above_272k_tokens"] + == profile.cache_write_above_272k + ) + assert model_info["cache_read_input_token_cost"] == profile.cache_read + assert ( + model_info["cache_read_input_token_cost_above_272k_tokens"] + == profile.cache_read_above_272k + ) + + +def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map): + """A prompt over 272K tokens is billed at the long-context rate, not the base rate.""" + response = _bedrock_response( + "bedrock/us.openai.gpt-5.6-sol", + Usage(prompt_tokens=300000, completion_tokens=1000, total_tokens=301000), + ) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx((300000 * 1.1e-05) + (1000 * 4.95e-05), rel=1e-9) + + +def test_bedrock_gpt_5_6_bills_cache_read_tokens(local_model_cost_map): + """Bedrock caches long prefixes implicitly and reports them, so a cache-read turn + must be billed at the cache rate rather than dropped to zero.""" + usage = Usage( + prompt_tokens=15611, + completion_tokens=5, + total_tokens=15616, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=15609), + ) + response = _bedrock_response("bedrock/us.openai.gpt-5.6-sol", usage) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + expected = (2 * 5.5e-06) + (15609 * 5.5e-07) + (5 * 3.3e-05) + assert cost == pytest.approx(expected, rel=1e-9) + # Without cache_read_input_token_cost the cached prefix bills at zero. + assert cost > (15611 * 5.5e-06) * 0.1 + + +def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map): + """The write side of the same cache cycle is billed at the 30m cache-write rate.""" + usage = Usage( + prompt_tokens=15611, + completion_tokens=5, + total_tokens=15616, + cache_creation_input_tokens=15609, + ) + response = _bedrock_response("bedrock/us.openai.gpt-5.6-sol", usage) + + cost = completion_cost( + completion_response=response, + model="bedrock/us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock", + ) + + expected = (2 * 5.5e-06) + (15609 * 6.875e-06) + (5 * 3.3e-05) + assert cost == pytest.approx(expected, rel=1e-9) + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_advertises_only_converse_supported_features( + profile, local_model_cost_map +): + model_info = _get_model_info_helper( + model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" + ) + + assert model_info["supports_function_calling"] is True + assert model_info["supports_tool_choice"] is True + assert model_info["supports_vision"] is True + + # Bedrock rejects an explicit cachePoint block for these models, so the flag that + # offers caller-driven caching stays off even though the cache rates are declared. + assert not model_info.get("supports_prompt_caching") + + # ModelInfoBase drops these two, so they are read from the map litellm resolves. + raw = _packaged_cost_map()[profile.model_id] + assert raw["supported_modalities"] == ["text", "image"] + assert raw["supported_output_modalities"] == ["text"] + # No bedrock_converse entry declares supported_endpoints; these models are reachable + # on chat completions and on the Responses API without it. + assert "supported_endpoints" not in raw + + +@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) +def test_bedrock_gpt_5_6_offers_tools_but_not_reasoning(profile, local_model_cost_map): + """Converse rejects the Anthropic-shaped thinking block LiteLLM emits for + reasoning_effort, so neither reasoning param may be offered yet, while the tool + params these models do accept must be.""" + supported = AmazonConverseConfig().get_supported_openai_params( + model=f"bedrock/{profile.model_id}" + ) + + assert "tools" in supported + assert "tool_choice" in supported + assert "reasoning_effort" not in supported + assert "thinking" not in supported From 13d4074492aa03b4d35a62fc8ffb8de2ef40e8dc Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 21 Aug 2026 04:41:21 -0700 Subject: [PATCH 04/43] test(mcp): retire the last file of the dead tests/litellm mirror tests/litellm/ was a second mirror beside tests/test_litellm/ that no workflow, Makefile target, or CircleCI job ever named. Its other 33 files were reconciled during August 2026; this one stayed behind under a ci-coverage-allowlist entry asking a later pass to decide which of its five orphan behaviours still hold. They no longer hold as written: 25 of its 32 cases fail against today's code, because the file froze on the day it stopped being collected and the endpoints kept moving. Three of the five are already covered by the live twin, and better. test_get_request_base_url_xff_trust_gate parametrizes the trust gate in both directions, including the exact untrusted-caller case the orphan asserted, and the standard and legacy protected-resource shapes are both exercised through use_standard_pattern. The other two were the only tests anywhere for validate_trusted_redirect_uri under that same gate, so they are ported rather than dropped, rebuilt on the live file's request-mock conventions. Both directions are load-bearing: forcing is_request_from_trusted_proxy to True fails the untrusted case, forcing it to False fails the trusted one. 313 tests pass in the live file, up from 311. Dropping the dead file clears one zero-assert TQ001 violation, so its ceiling ratchets down with it. --- .github/ci-coverage-allowlist.yml | 10 - test-quality-budget.json | 2 +- .../mcp_server/test_discoverable_endpoints.py | 1268 ----------------- .../mcp_server/test_discoverable_endpoints.py | 49 + 4 files changed, 50 insertions(+), 1279 deletions(-) delete mode 100644 tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index ff8fa864d4a..918589f84d1 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -48,16 +48,6 @@ test_paths: choice it informed is settled paths: - tests/code_coverage_tests/test_aio_http_image_conversion.py - - reason: >- - The last file of a second mirror that sat beside tests/test_litellm and ran nowhere. Its - other 33 files landed in the real mirror during August 2026, 30 as moves and 3 by merging - their bodies into the live file of the same name. This one cannot follow either route yet: - its live twin was rewritten from 1268 lines to 9434, and of the 19 tests here 5 have no - counterpart while 25 assertions fail against today's code, so what survives that rewrite - is a judgement about the endpoints, not a merge. Revisit by deciding which of the five - behaviours still hold - paths: - - tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py - reason: >- No job invokes this suite and its files mix pure transformation tests with ones driving live vendor vector stores, so assigning them needs a per-file decision diff --git a/test-quality-budget.json b/test-quality-budget.json index 1613c8c75cb..91ae881c83a 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -1,6 +1,6 @@ { "TQ001": { - "limit": 750 + "limit": 746 }, "TQ002": { "limit": 742 diff --git a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py deleted file mode 100644 index 2a8768df722..00000000000 --- a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ /dev/null @@ -1,1268 +0,0 @@ -"""Tests for MCP OAuth discoverable endpoints""" - -import pytest -from fastapi import HTTPException -from unittest.mock import AsyncMock, MagicMock, patch - -TRUSTED_PROXY_IP = "10.0.0.5" -TRUSTED_PROXY_RANGES = ["10.0.0.0/8"] - - -def set_request_from_trusted_proxy(mock_request): - mock_request.client = MagicMock() - mock_request.client.host = TRUSTED_PROXY_IP - - -@pytest.fixture -def trusted_proxy_origin_headers(): - with ( - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), - patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), - ): - yield - - -@pytest.mark.asyncio -async def test_authorize_endpoint_includes_response_type(): - """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Mock the encryption functions to avoid needing a signing key - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify response is a redirect - assert response.status_code == 307 # FastAPI RedirectResponse default - - # Verify response_type is in the redirect URL - assert "response_type=code" in response.headers["location"] - assert "https://provider.com/oauth/authorize" in response.headers["location"] - assert "client_id=test_client_id" in response.headers["location"] - assert "scope=read+write" in response.headers["location"] - - -@pytest.mark.asyncio -async def test_authorize_endpoint_forwards_pkce_parameters(): - """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server (simulating Google OAuth) - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="669428968603-test.apps.googleusercontent.com", - client_secret="GOCSPX-test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm-proxy.example.com/" - mock_request.headers = {} - - # Mock the encryption function - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state_with_pkce" - - # Call authorize endpoint with PKCE parameters - response = await authorize( - request=mock_request, - client_id="669428968603-test.apps.googleusercontent.com", - mcp_server_name="google_mcp", - redirect_uri="http://localhost:60108/callback", - state="test_client_state", - code_challenge="x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk", - code_challenge_method="S256", - ) - - # Verify response is a redirect - assert response.status_code == 307 - - # Verify PKCE parameters are included in the redirect URL - location = response.headers["location"] - assert "https://accounts.google.com/o/oauth2/v2/auth" in location - assert "code_challenge=x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk" in location - assert "code_challenge_method=S256" in location - assert "client_id=669428968603-test.apps.googleusercontent.com" in location - assert "response_type=code" in location - - -@pytest.mark.asyncio -async def test_token_endpoint_forwards_code_verifier(): - """Test that token endpoint forwards code_verifier for PKCE flow""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="669428968603-test.apps.googleusercontent.com", - client_secret="GOCSPX-test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm-proxy.example.com/" - mock_request.headers = {} - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "ya29.test_access_token", - "token_type": "Bearer", - "expires_in": 3599, - "scope": "openid email https://www.googleapis.com/auth/drive", - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client with AsyncMock for async methods - from unittest.mock import AsyncMock - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_async_client = MagicMock() - # Use AsyncMock for the async post method - mock_async_client.post = AsyncMock(return_value=mock_response) - mock_get_client.return_value = mock_async_client - - # Call token endpoint with code_verifier - response = await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="4/test_authorization_code", - redirect_uri="http://localhost:60108/callback", - client_id="669428968603-test.apps.googleusercontent.com", - mcp_server_name="google_mcp", - client_secret="GOCSPX-test_secret", - code_verifier="test_code_verifier_from_client", - ) - - # Verify that the token endpoint was called with code_verifier - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - - # Check the data parameter includes code_verifier - assert call_args[1]["data"]["code_verifier"] == "test_code_verifier_from_client" - assert call_args[1]["data"]["code"] == "4/test_authorization_code" - assert ( - call_args[1]["data"]["client_id"] - == "669428968603-test.apps.googleusercontent.com" - ) - assert call_args[1]["data"]["client_secret"] == "GOCSPX-test_secret" - assert call_args[1]["data"]["grant_type"] == "authorization_code" - - # Verify response - response_data = response.body - import json - - token_data = json.loads(response_data) - assert token_data["access_token"] == "ya29.test_access_token" - assert token_data["token_type"] == "Bearer" - - -@pytest.mark.asyncio -async def test_register_client_without_mcp_server_name_returns_dummy(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client(request=mock_request) - - assert result == { - "client_id": "dummy_client", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_register_client_returns_existing_server_credentials(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - oauth2_server = MCPServer( - server_id="stored_server", - name="stored_server", - server_name="stored_server", - alias="stored_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="existing-client", - client_secret="existing-secret", - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - - try: - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) - finally: - global_mcp_server_manager.registry.clear() - - assert result == { - "client_id": "stored_server", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_register_client_remote_registration_success(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - oauth2_server = MCPServer( - server_id="remote_server", - name="remote_server", - server_name="remote_server", - alias="remote_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id=None, - client_secret=None, - authorization_url="https://provider.example/oauth/authorize", - token_url="https://provider.example/oauth/token", - registration_url="https://provider.example/oauth/register", - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://proxy.litellm.example/" - mock_request.headers = {} - - request_payload = { - "client_name": "Litellm Proxy", - "grant_types": ["authorization_code", "refresh_token"], - "response_types": ["code"], - "token_endpoint_auth_method": "client_secret_post", - } - - mock_response = MagicMock() - mock_response.json.return_value = { - "client_id": "generated-client", - "client_secret": "generated-secret", - } - mock_response.raise_for_status = MagicMock() - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - try: - with ( - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value=request_payload), - ), - patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", - return_value=mock_async_client, - ), - ): - response = await register_client( - request=mock_request, mcp_server_name=oauth2_server.server_name - ) - finally: - global_mcp_server_manager.registry.clear() - - import json - - assert response.status_code == 200 - payload = json.loads(response.body.decode("utf-8")) - assert payload == mock_response.json.return_value - - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - assert call_args.args[0] == oauth2_server.registration_url - assert call_args.kwargs["headers"] == { - "Content-Type": "application/json", - "Accept": "application/json", - } - assert call_args.kwargs["json"]["redirect_uris"] == [ - "https://proxy.litellm.example/callback" - ] - assert call_args.kwargs["json"]["grant_types"] == request_payload["grant_types"] - assert ( - call_args.kwargs["json"]["token_endpoint_auth_method"] - == request_payload["token_endpoint_auth_method"] - ) - - -@pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify redirect URL uses HTTPS in the redirect_uri parameter - location = response.headers["location"] - - # The redirect_uri parameter sent to the OAuth provider should use HTTPS - assert ( - "redirect_uri=https%3A%2F%2Flitellm.example.com%2Fcallback" in location - or "redirect_uri=https://litellm.example.com/callback" in location - ) - - -@pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3599, - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_get_client.return_value = mock_async_client - - # Call token endpoint - await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="test_code", - redirect_uri="http://localhost:60108/callback", - client_id="test_client_id", - mcp_server_name="google_mcp", - client_secret="test_secret", - ) - - # Verify that the redirect_uri sent to the provider uses HTTPS - call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://litellm-proxy.example.com/callback" - ) - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_standard_pattern(): - """Test that oauth_protected_resource_mcp_standard returns standard MCP URL pattern (/mcp/{server_name})""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp_standard, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_server", - name="test_server", - server_name="test_server", - alias="test_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Call the standard pattern endpoint - response = await oauth_protected_resource_mcp_standard( - request=mock_request, - mcp_server_name="test_server", - ) - - # Verify response uses standard MCP pattern: /mcp/{server_name} - assert response["resource"] == "https://litellm.example.com/mcp/test_server" - assert ( - response["authorization_servers"][0] - == "https://litellm.example.com/test_server" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_legacy_pattern(): - """Test that oauth_protected_resource_mcp returns legacy URL pattern (/{server_name}/mcp)""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_server", - name="test_server", - server_name="test_server", - alias="test_server", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - # Call the legacy pattern endpoint - response = await oauth_protected_resource_mcp( - request=mock_request, - mcp_server_name="test_server", - ) - - # Verify response uses legacy pattern: /{server_name}/mcp - assert response["resource"] == "https://litellm.example.com/test_server/mcp" - assert ( - response["authorization_servers"][0] - == "https://litellm.example.com/test_server" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_protected_resource_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_protected_resource_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Call the endpoint - response = await oauth_protected_resource_mcp( - request=mock_request, - mcp_server_name="test_oauth", - ) - - # Verify response uses HTTPS URLs - assert response["authorization_servers"][0].startswith( - "https://litellm.example.com/" - ) - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_oauth_authorization_server_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - oauth_authorization_server_mcp, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://litellm.example.com/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - # Call the endpoint - response = await oauth_authorization_server_mcp( - request=mock_request, - mcp_server_name="test_oauth", - ) - - # Verify response uses HTTPS URLs - assert response["authorization_endpoint"].startswith("https://litellm.example.com/") - assert response["token_endpoint"].startswith("https://litellm.example.com/") - assert response["registration_endpoint"].startswith("https://litellm.example.com/") - assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] - assert response["scopes_supported"] == oauth2_server.scopes - - -@pytest.mark.asyncio -async def test_register_client_respects_x_forwarded_proto( - trusted_proxy_origin_headers, -): - """Test that register_client uses X-Forwarded-Proto for redirect_uris""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - register_client, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - global_mcp_server_manager.registry.clear() - - # Mock request with http base_url but X-Forwarded-Proto: https - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://proxy.litellm.example/" # HTTP - mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", - new=AsyncMock(return_value={}), - ): - result = await register_client(request=mock_request) - - # Verify the redirect_uris use HTTPS - assert result == { - "client_id": "dummy_client", - "client_secret": "dummy", - "redirect_uris": ["https://proxy.litellm.example/callback"], - } - - -@pytest.mark.asyncio -async def test_authorize_endpoint_respects_x_forwarded_host( - trusted_proxy_origin_headers, -): - """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - authorize, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="test_oauth_server", - name="test_oauth", - server_name="test_oauth", - alias="test_oauth", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_client_secret", - authorization_url="https://provider.com/oauth/authorize", - token_url="https://provider.com/oauth/token", - scopes=["read", "write"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request simulating nginx proxy: - # Internal: http://localhost:8888/github/mcp - # External: https://proxy.example.com/github/mcp - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:8888/github/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - # Mock the encryption functions - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" - ) as mock_encrypt: - mock_encrypt.return_value = "mocked_encrypted_state" - - # Call authorize endpoint - response = await authorize( - request=mock_request, - client_id="test_client_id", - mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", - state="test_state", - ) - - # Verify redirect URL uses the forwarded host and scheme - location = response.headers["location"] - - # The redirect_uri parameter should use the external URL - assert ( - "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" - in location - or "redirect_uri=https://proxy.example.com/github/mcp/callback" in location - ) - - -@pytest.mark.asyncio -async def test_token_endpoint_respects_x_forwarded_host( - trusted_proxy_origin_headers, -): - """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - token_endpoint, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - from litellm.proxy._types import MCPTransport - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Clear registry - global_mcp_server_manager.registry.clear() - - # Create mock OAuth2 server - oauth2_server = MCPServer( - server_id="google_mcp", - name="google_mcp", - server_name="google_mcp", - alias="google_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - client_id="test_client_id", - client_secret="test_secret", - authorization_url="https://accounts.google.com/o/oauth2/v2/auth", - token_url="https://oauth2.googleapis.com/token", - scopes=["openid", "email"], - ) - global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server - - # Mock request simulating nginx proxy without port in host - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:8888/github/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - # Mock httpx client response - mock_response = MagicMock() - mock_response.json.return_value = { - "access_token": "test_token", - "token_type": "Bearer", - "expires_in": 3599, - } - mock_response.raise_for_status = MagicMock() - - # Mock the async httpx client - mock_async_client = MagicMock() - mock_async_client.post = AsyncMock(return_value=mock_response) - - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" - ) as mock_get_client: - mock_get_client.return_value = mock_async_client - - # Call token endpoint - await token_endpoint( - request=mock_request, - grant_type="authorization_code", - code="test_code", - redirect_uri="http://localhost:60108/callback", - client_id="test_client_id", - mcp_server_name="google_mcp", - client_secret="test_secret", - ) - - # Verify that the redirect_uri sent to the provider uses the external URL - call_args = mock_async_client.post.call_args - assert ( - call_args[1]["data"]["redirect_uri"] - == "https://proxy.example.com/github/mcp/callback" - ) - - -@pytest.mark.parametrize( - "base_url,x_forwarded_proto,x_forwarded_host,x_forwarded_port,expected_url", - [ - # Case 1: No forwarded headers - use original URL as-is (no trailing slash) - ( - "http://localhost:4000/", - None, - None, - None, - "http://localhost:4000", - ), - # Case 2: Only X-Forwarded-Proto - change scheme only - ( - "http://localhost:4000/", - "https", - None, - None, - "https://localhost:4000", - ), - # Case 3: X-Forwarded-Proto + X-Forwarded-Host - change scheme and host - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - None, - "https://proxy.example.com", - ), - # Case 4: X-Forwarded-Host with port included in host header - ( - "http://localhost:4000/", - "https", - "proxy.example.com:8080", - None, - "https://proxy.example.com:8080", - ), - # Case 5: X-Forwarded-Host + X-Forwarded-Port as separate headers - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - "8443", - "https://proxy.example.com:8443", - ), - # Case 6: Only X-Forwarded-Host without proto - use original scheme - ( - "http://localhost:4000/", - None, - "proxy.example.com", - None, - "http://proxy.example.com", - ), - # Case 7: Only X-Forwarded-Port without host - preserves original port if present - # (This is safer behavior - X-Forwarded-Port alone is unusual) - ( - "http://localhost:4000/", - None, - None, - "8443", - "http://localhost:4000", # Original port preserved when already present - ), - # Case 8: Complex internal URL with path (path is preserved) - ( - "http://localhost:8888/github/mcp", - "https", - "proxy.example.com", - None, - "https://proxy.example.com/github/mcp", - ), - # Case 9: IPv6 address in X-Forwarded-Host (should not treat :: as port separator) - ( - "http://localhost:4000/", - "https", - "[2001:db8::1]", - None, - "https://[2001:db8::1]", - ), - # Case 10: IPv6 address with port - ( - "http://localhost:4000/", - "https", - "[2001:db8::1]:8080", - None, - "https://[2001:db8::1]:8080", - ), - # Case 11: X-Forwarded-Host already has port, X-Forwarded-Port also provided (host wins) - ( - "http://localhost:4000/", - "https", - "proxy.example.com:9000", - "8443", - "https://proxy.example.com:9000", - ), - # Case 12: Standard proxy setup (most common case) - ( - "http://127.0.0.1:8888/", - "https", - "chatproxy.company.com", - None, - "https://chatproxy.company.com", - ), - # Case 13: Internal URL already has port, X-Forwarded-Port does NOT override - # (safer behavior - preserves original port when X-Forwarded-Host not provided) - ( - "http://localhost:4000/", - None, - None, - "443", - "http://localhost:4000", # Original port preserved - ), - # Case 14: Original URL with existing port in netloc, X-Forwarded-Host replaces it - ( - "http://internal.local:8888/", - "https", - "external.com", - None, - "https://external.com", - ), - ], -) -def test_get_request_base_url_comprehensive( - base_url, - x_forwarded_proto, - x_forwarded_host, - x_forwarded_port, - expected_url, - trusted_proxy_origin_headers, -): - """Comprehensive test for get_request_base_url with various header combinations""" - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - # Create mock request - mock_request = MagicMock(spec=Request) - mock_request.base_url = base_url - set_request_from_trusted_proxy(mock_request) - - # Build headers dict - headers = {} - if x_forwarded_proto: - headers["X-Forwarded-Proto"] = x_forwarded_proto - if x_forwarded_host: - headers["X-Forwarded-Host"] = x_forwarded_host - if x_forwarded_port: - headers["X-Forwarded-Port"] = x_forwarded_port - - # Mock headers.get() to return our test values - def mock_get(header_name, default=None): - return headers.get(header_name, default) - - mock_request.headers.get = mock_get - - # Test the function - result = get_request_base_url(mock_request) - - # Verify result - assert result == expected_url, ( - f"Expected '{expected_url}' but got '{result}'\n" - f"Input: base_url={base_url}, " - f"X-Forwarded-Proto={x_forwarded_proto}, " - f"X-Forwarded-Host={x_forwarded_host}, " - f"X-Forwarded-Port={x_forwarded_port}" - ) - - -def test_get_request_base_url_ignores_forwarded_headers_from_untrusted_client(): - try: - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://gateway.example.com/mcp" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "attacker.example.com", - "X-Forwarded-Port": "443", - } - mock_request.client = MagicMock() - mock_request.client.host = "203.0.113.10" - - with patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ): - assert get_request_base_url(mock_request) == "https://gateway.example.com/mcp" - - -def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): - try: - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - validate_trusted_redirect_uri, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP OAuth utilities not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://gateway.example.com/" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "attacker.example.com", - } - mock_request.client = MagicMock() - mock_request.client.host = "203.0.113.10" - - with ( - patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ), - pytest.raises(HTTPException), - ): - validate_trusted_redirect_uri( - mock_request, - "https://attacker.example.com/callback", - ) - - -def test_validate_trusted_redirect_uri_allows_forwarded_origin_from_trusted_proxy( - trusted_proxy_origin_headers, -): - try: - from litellm.proxy._experimental.mcp_server.oauth_utils import ( - validate_trusted_redirect_uri, - ) - from fastapi import Request - except ImportError: - pytest.skip("MCP OAuth utilities not available") - - mock_request = MagicMock(spec=Request) - mock_request.base_url = "http://localhost:4000/" - mock_request.headers = { - "X-Forwarded-Proto": "https", - "X-Forwarded-Host": "proxy.example.com", - } - set_request_from_trusted_proxy(mock_request) - - validate_trusted_redirect_uri( - mock_request, - "https://proxy.example.com/callback", - ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index b4d3782ba43..442bfe8a090 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2957,6 +2957,55 @@ def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection(caplog, monk assert "X-Forwarded-Host" in msg +@pytest.mark.parametrize( + "direct_ip,expect_accepted", + [ + ("10.0.0.7", True), + ("203.0.113.5", False), + ], +) +def test_validate_trusted_redirect_uri_follows_the_xff_trust_gate(direct_ip, expect_accepted, monkeypatch): + try: + from fastapi import HTTPException, Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP oauth_utils not available") + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("MCP_TRUSTED_REDIRECT_ORIGINS", raising=False) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = direct_ip + + headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "proxy.example.com", + } + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + mock_request.headers.__contains__ = lambda self_, name: name in headers + + redirect_uri = "https://proxy.example.com/callback" + general_settings = { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + } + + with patch("litellm.proxy.proxy_server.general_settings", general_settings, create=True): + if expect_accepted: + validate_trusted_redirect_uri(mock_request, redirect_uri) + return + with pytest.raises(HTTPException) as exc_info: + validate_trusted_redirect_uri(mock_request, redirect_uri) + + assert exc_info.value.status_code == 400 + assert "proxy.example.com" in str(exc_info.value.detail) + + @pytest.mark.parametrize( "bad_value", [ From e6a6016e3e9c0b35a39f2caf638ce3cd44a8603c Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Fri, 21 Aug 2026 19:48:27 +0000 Subject: [PATCH 05/43] fix(model-costs): apply GPT-5.6 Sol promotional pricing cut Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 60 +++++++++---------- model_prices_and_context_window.json | 60 +++++++++---------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 18 +++--- tests/test_litellm/test_cost_calculator.py | 2 +- 4 files changed, 70 insertions(+), 70 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a1961136d11..4f245fcdfbf 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26104,33 +26104,33 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-sol": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, - "cache_creation_input_token_cost_flex": 3.125e-06, - "cache_creation_input_token_cost_priority": 1.25e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, - "cache_read_input_token_cost_flex": 2.5e-07, - "cache_read_input_token_cost_priority": 1e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_above_272k_tokens_flex": 5e-06, - "input_cost_per_token_batches": 2.5e-06, - "input_cost_per_token_flex": 2.5e-06, - "input_cost_per_token_priority": 1e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_flex": 2e-07, + "cache_read_input_token_cost_priority": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_batches": 2e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 8e-06, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "output_cost_per_token_batches": 1.5e-05, - "output_cost_per_token_flex": 1.5e-05, - "output_cost_per_token_priority": 6e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_batches": 1e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -26372,19 +26372,19 @@ "supports_parallel_function_calling": true }, "daybreak-blue-latest": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a1961136d11..4f245fcdfbf 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26104,33 +26104,33 @@ "supports_xhigh_reasoning_effort": true }, "gpt-5.6-sol": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, - "cache_creation_input_token_cost_flex": 3.125e-06, - "cache_creation_input_token_cost_priority": 1.25e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, - "cache_read_input_token_cost_flex": 2.5e-07, - "cache_read_input_token_cost_priority": 1e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_above_272k_tokens_flex": 5e-06, - "input_cost_per_token_batches": 2.5e-06, - "input_cost_per_token_flex": 2.5e-06, - "input_cost_per_token_priority": 1e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_flex": 2e-07, + "cache_read_input_token_cost_priority": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_batches": 2e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 8e-06, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "output_cost_per_token_batches": 1.5e-05, - "output_cost_per_token_flex": 1.5e-05, - "output_cost_per_token_priority": 6e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_batches": 1e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -26372,19 +26372,19 @@ "supports_parallel_function_calling": true }, "daybreak-blue-latest": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" 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 f66056a54e2..0e7db195865 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 @@ -914,7 +914,7 @@ def test_generic_cost_per_token_gpt55_pro(): "model,input_cost,output_cost,cache_read_cost,cache_write_cost", [ ("gpt-5.6", 5e-6, 3e-5, 5e-7, 6.25e-6), - ("gpt-5.6-sol", 5e-6, 3e-5, 5e-7, 6.25e-6), + ("gpt-5.6-sol", 4e-6, 2e-5, 4e-7, 5e-6), ("gpt-5.6-terra", 2e-6, 1.2e-5, 2e-7, 2.5e-6), ("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7), ], @@ -969,7 +969,7 @@ def test_generic_cost_per_token_gpt56( "model,flex_long_input_cost,flex_long_output_cost", [ ("gpt-5.6", 5e-6, 2.25e-5), - ("gpt-5.6-sol", 5e-6, 2.25e-5), + ("gpt-5.6-sol", 4e-6, 1.5e-5), ("gpt-5.6-terra", 2e-6, 9e-6), ("gpt-5.6-luna", 2e-7, 9e-7), ], @@ -3300,8 +3300,8 @@ def test_generic_cost_per_token_gemini_35_flash_lite(): @pytest.mark.parametrize( "service_tier,input_rate,cache_read_rate,cache_write_rate,output_rate", [ - ("flex", 2.5e-6, 2.5e-7, 3.125e-6, 1.5e-5), - ("priority", 1e-5, 1e-6, 1.25e-5, 6e-5), + ("flex", 2e-6, 2e-7, 2.5e-6, 1e-5), + ("priority", 8e-6, 8e-7, 1e-5, 4e-5), ], ) def test_service_tier_cache_creation_rates_for_gpt_5_6( @@ -3314,7 +3314,7 @@ def test_service_tier_cache_creation_rates_for_gpt_5_6( ): """Regression: gpt-5.6 publishes cache_creation_input_token_cost_flex/_priority, so a flex or priority request must bill cache writes at that tier's rate instead of falling - back to the standard 6.25e-6 rate.""" + back to the standard cache-write rate.""" usage = Usage( prompt_tokens=10_000, completion_tokens=500, @@ -3361,8 +3361,8 @@ def test_fast_service_tier_bills_at_the_priority_rate(_local_model_cost_map): model="gpt-5.6-sol", usage=usage, custom_llm_provider="openai", service_tier="fast" ) - expected_prompt = 800 * 1e-05 + 200 * 1e-06 - expected_completion = 500 * 6e-05 + expected_prompt = 800 * 8e-06 + 200 * 8e-07 + expected_completion = 500 * 4e-05 assert fast == priority assert fast[0] == pytest.approx(expected_prompt, rel=1e-9) @@ -3397,8 +3397,8 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m ) assert fast == priority - assert fast[0] == pytest.approx(300_000 * 1e-05, rel=1e-9) - assert fast[1] == pytest.approx(1_000 * 4.5e-05, rel=1e-9) + assert fast[0] == pytest.approx(300_000 * 8e-06, rel=1e-9) + assert fast[1] == pytest.approx(1_000 * 3e-05, rel=1e-9) def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map): diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 98938dee62e..2b30138faa2 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3774,4 +3774,4 @@ def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(): custom_llm_provider="openai", ) - assert cost == pytest.approx(3 * 5e-6 + 4014 * 5e-7 + 5 * 3e-5, rel=1e-9) + assert cost == pytest.approx(3 * 4e-6 + 4014 * 4e-7 + 5 * 2e-5, rel=1e-9) From 5ed230701a66177e732480a8c5dc901fb0454671 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 21 Aug 2026 16:22:07 -0700 Subject: [PATCH 06/43] test: escape the literal match= patterns PT017 minted --- tests/llm_translation/test_containers_api.py | 2 +- tests/llm_translation/test_litellm_proxy_provider.py | 3 ++- tests/local_testing/test_rules.py | 5 +++-- tests/proxy_admin_ui_tests/test_role_based_access.py | 3 ++- tests/proxy_unit_tests/test_key_generate_prisma.py | 9 +++++++-- .../llms/bedrock/chat/test_converse_transformation.py | 6 +++--- 6 files changed, 18 insertions(+), 10 deletions(-) diff --git a/tests/llm_translation/test_containers_api.py b/tests/llm_translation/test_containers_api.py index 6c7303e7b4d..7fb0c6d21d6 100644 --- a/tests/llm_translation/test_containers_api.py +++ b/tests/llm_translation/test_containers_api.py @@ -63,7 +63,7 @@ def test_container_files_api(): # 3. Try retrieve non-existent file metadata (should raise error) print("3. Testing retrieve_container_file (expect error)...") - with pytest.raises(Exception, match="(?i)not found|invalid"): + with pytest.raises(Exception, match=r"(?i)not found|invalid"): retrieve_container_file( container_id=container.id, file_id="cfile_nonexistent", diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index cea0167472e..7a917c226df 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -1,5 +1,6 @@ import json import os +import re import sys from datetime import datetime from io import BytesIO @@ -578,7 +579,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): def test_litellm_gateway_from_sdk_with_thinking_param(): - with pytest.raises(Exception, match="Connection error.") as exc_info: + with pytest.raises(Exception, match=re.escape("Connection error.")) as exc_info: response = litellm.completion( model="litellm_proxy/anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[{"role": "user", "content": "Hello world"}], diff --git a/tests/local_testing/test_rules.py b/tests/local_testing/test_rules.py index b075821e205..7ffab789d64 100644 --- a/tests/local_testing/test_rules.py +++ b/tests/local_testing/test_rules.py @@ -2,6 +2,7 @@ # This tests setting rules before / after making llm api calls import asyncio import os +import re import sys import time import traceback @@ -82,7 +83,7 @@ def test_post_call_rule(): litellm.post_call_rules = [my_post_call_rule] ### completion - with pytest.raises(Exception, match="This violates LiteLLM Proxy Rules. Response too short") as exc_info: + with pytest.raises(Exception, match=re.escape("This violates LiteLLM Proxy Rules. Response too short")) as exc_info: completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "say sorry"}], @@ -118,7 +119,7 @@ def test_post_call_rule_streaming(): stream=True, ) - with pytest.raises(Exception, match="This violates LiteLLM Proxy Rules. Response too short") as exc_info: + with pytest.raises(Exception, match=re.escape("This violates LiteLLM Proxy Rules. Response too short")) as exc_info: list(response) assert "This violates LiteLLM Proxy Rules. Response too short" in exc_info.value.message diff --git a/tests/proxy_admin_ui_tests/test_role_based_access.py b/tests/proxy_admin_ui_tests/test_role_based_access.py index b5a076d0185..1c4ee2caa04 100644 --- a/tests/proxy_admin_ui_tests/test_role_based_access.py +++ b/tests/proxy_admin_ui_tests/test_role_based_access.py @@ -3,6 +3,7 @@ RBAC tests """ import os +import re import sys import traceback from litellm._uuid import uuid @@ -411,7 +412,7 @@ async def test_org_admin_create_user_team_wrong_org_permissions(prisma_client): request.body = return_body with pytest.raises( - Exception, match="You do not have a role within the selected organization. Passed organization_id" + Exception, match=re.escape("You do not have a role within the selected organization. Passed organization_id") ) as exc_info: response = await user_api_key_auth(request=request, api_key="Bearer " + new_key) e = exc_info.value diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 16507aaaf55..efedc156429 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -20,6 +20,7 @@ # function to validate a request - async def user_auth(request: Request): import os +import re import sys import traceback from litellm._uuid import uuid @@ -1498,7 +1499,9 @@ def test_key_generate_with_custom_auth(prisma_client): await litellm.proxy.proxy_server.prisma_client.connect() request = GenerateKeyRequest() - with pytest.raises(Exception, match="This violates LiteLLM Proxy Rules. No team id provided.") as exc_info: + with pytest.raises( + Exception, match=re.escape("This violates LiteLLM Proxy Rules. No team id provided.") + ) as exc_info: key = await generate_key_fn( request, user_api_key_dict=UserAPIKeyAuth( @@ -3045,7 +3048,9 @@ async def test_custom_api_key_header_name(prisma_client): "headers": [], } ) - with pytest.raises(Exception, match="Malformed API Key passed in. Ensure Key has `Bearer ` prefix") as exc_info: + with pytest.raises( + Exception, match=re.escape("Malformed API Key passed in. Ensure Key has `Bearer ` prefix") + ) as exc_info: result = await user_api_key_auth(request=request, api_key="Bearer sk-1234") e = exc_info.value print("failed with error", e) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 30843e8160b..a2f88138eaa 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -3025,7 +3025,7 @@ def test_request_metadata_key_constraints(): long_key = "a" * 257 invalid_metadata = {long_key: "value"} - with pytest.raises(Exception, match="(?i)key length|256 characters"): + with pytest.raises(Exception, match=r"(?i)key length|256 characters"): config.transform_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", messages=messages, @@ -3037,7 +3037,7 @@ def test_request_metadata_key_constraints(): # Test empty key invalid_metadata = {"": "value"} - with pytest.raises(Exception, match="(?i)key length|empty"): + with pytest.raises(Exception, match=r"(?i)key length|empty"): config.transform_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", messages=messages, @@ -3057,7 +3057,7 @@ def test_request_metadata_value_constraints(): long_value = "a" * 257 invalid_metadata = {"key": long_value} - with pytest.raises(Exception, match="(?i)value length|256 characters"): + with pytest.raises(Exception, match=r"(?i)value length|256 characters"): config.transform_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", messages=messages, From 91599aef693600f45161e001e4a1928867fbc502 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 21 Aug 2026 16:25:33 -0700 Subject: [PATCH 07/43] test: say whether a match= pattern is a regex or a literal (ruff RUF043) --- ruff-tests.toml | 5 +++++ tests/local_testing/test_embedding.py | 5 +++-- tests/multi_instance_e2e_tests/test_update_team_e2e.py | 6 +++--- tests/router_unit_tests/test_router_helper_utils.py | 2 +- .../integrations/bitbucket/test_bitbucket_prompt_manager.py | 3 ++- .../integrations/gitlab/test_gitlab_prompt_manager.py | 3 ++- tests/test_litellm/integrations/test_openmeter.py | 2 +- tests/test_litellm/litellm_core_utils/test_url_utils.py | 4 ++-- .../llms/anthropic/test_anthropic_common_utils.py | 2 +- tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py | 6 +++--- .../llms/vertex_ai/files/test_vertex_ai_files_handler.py | 3 ++- .../llms/vertex_ai/test_vertex_ai_common_utils.py | 2 +- tests/test_litellm/proxy/client/cli/test_agents.py | 2 +- tests/test_litellm/proxy/db/test_db_url_settings.py | 4 ++-- tests/test_litellm/proxy/proxy_server/test_proxy_config.py | 3 ++- tests/test_litellm/proxy/test_proxy_server.py | 3 ++- tests/test_litellm/sandbox/test_opensandbox_sandbox.py | 2 +- tests/test_litellm/test_router_model_cost_isolation.py | 5 +++-- 18 files changed, 37 insertions(+), 25 deletions(-) diff --git a/ruff-tests.toml b/ruff-tests.toml index ff29bcff313..de0931f5e69 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -32,6 +32,10 @@ # PT017 an `assert` on the caught error inside `except`. Nothing runs the handler when # the call stops raising, so the test goes green on the exact regression it was # written to catch. `pytest.raises` fails when the call succeeds +# RUF043 a `match=` pattern carrying regex metacharacters in a plain string. `match=` is +# `re.search`, so a `.` copied out of an error message is a wildcard and the block +# accepts messages the author never meant to accept. Mark a real regex raw, wrap a +# literal message in `re.escape`, and the pattern says which one it is # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -53,4 +57,5 @@ lint.select = [ "PT017", "PLR0133", "PLW0127", + "RUF043", ] diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index ee9d4cdd915..677f8b10cca 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -1,5 +1,6 @@ import json import os +import re import sys import traceback @@ -536,13 +537,13 @@ def test_demo_tokens_as_input_to_embeddings_fails_for_titan(): with pytest.raises( litellm.BadRequestError, - match='litellm.BadRequestError: BedrockException - {"message":"Malformed input request: expected type: String, found: JSONArray, please reformat your input and try again."}', + match=re.escape('litellm.BadRequestError: BedrockException - {"message":"Malformed input request: expected type: String, found: JSONArray, please reformat your input and try again."}'), ): litellm.embedding(model="amazon.titan-embed-text-v1", input=[[1]]) with pytest.raises( litellm.BadRequestError, - match='litellm.BadRequestError: BedrockException - {"message":"Malformed input request: expected type: String, found: Integer, please reformat your input and try again."}', + match=re.escape('litellm.BadRequestError: BedrockException - {"message":"Malformed input request: expected type: String, found: Integer, please reformat your input and try again."}'), ): litellm.embedding( model="amazon.titan-embed-text-v1", diff --git a/tests/multi_instance_e2e_tests/test_update_team_e2e.py b/tests/multi_instance_e2e_tests/test_update_team_e2e.py index 13091fd3df6..ce88e976ce0 100644 --- a/tests/multi_instance_e2e_tests/test_update_team_e2e.py +++ b/tests/multi_instance_e2e_tests/test_update_team_e2e.py @@ -143,7 +143,7 @@ async def test_team_blocking_behavior_multi_instance(): assert team_info_4001["blocked"] is True, "Team should be blocked after update" # 8. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked. - with pytest.raises(Exception, match="(?i)blocked") as excinfo: + with pytest.raises(Exception, match=r"(?i)blocked") as excinfo: await chat_completion_on_port( session, key=key, @@ -157,7 +157,7 @@ async def test_team_blocking_behavior_multi_instance(): ), f"Expected error indicating team blocked, got: {error_msg}" # 9. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked. - with pytest.raises(Exception, match="(?i)blocked") as excinfo: + with pytest.raises(Exception, match=r"(?i)blocked") as excinfo: await chat_completion_on_port( session, key=key, @@ -171,7 +171,7 @@ async def test_team_blocking_behavior_multi_instance(): ), f"Expected error indicating team blocked, got: {error_msg}" # 9. Repeat the chat completion request with another new prompt; expect it to be blocked. - with pytest.raises(Exception, match="(?i)blocked") as excinfo_second: + with pytest.raises(Exception, match=r"(?i)blocked") as excinfo_second: await chat_completion_on_port( session, key=key, diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 755405c8b21..82bdbd7bfc9 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1832,7 +1832,7 @@ def test_init_auto_router_deployment_duplicate_model_name(mock_auto_router, mode ) with pytest.raises( - ValueError, match="Auto-router deployment test-auto-router with tags .* already exists" + ValueError, match=r"Auto-router deployment test-auto-router with tags .* already exists" ): router.init_auto_router_deployment(deployment) diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py index dd97de24df3..4a15da87c89 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_prompt_manager.py @@ -1,5 +1,6 @@ import json import os +import re import sys from unittest.mock import MagicMock, patch @@ -158,7 +159,7 @@ def test_bitbucket_client_get_file_content_access_denied(mock_get): client = BitBucketClient(config) - with pytest.raises(Exception, match="Access denied to file 'test.prompt'"): + with pytest.raises(Exception, match=re.escape("Access denied to file 'test.prompt'")): client.get_file_content("test.prompt") diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py b/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py index 1f7706882f6..adccd94141f 100644 --- a/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_prompt_manager.py @@ -1,4 +1,5 @@ import os +import re import sys from unittest.mock import MagicMock, patch @@ -172,7 +173,7 @@ def test_gitlab_client_get_file_content_access_denied(mock_get): mock_get.side_effect = err client = GitLabClient({"project": "g/s/r", "access_token": "tok"}) - with pytest.raises(Exception, match="Access denied to file 'test.prompt'"): + with pytest.raises(Exception, match=re.escape("Access denied to file 'test.prompt'")): client.get_file_content("test.prompt") diff --git a/tests/test_litellm/integrations/test_openmeter.py b/tests/test_litellm/integrations/test_openmeter.py index 539e3f99cdc..b9da99b6fa9 100644 --- a/tests/test_litellm/integrations/test_openmeter.py +++ b/tests/test_litellm/integrations/test_openmeter.py @@ -33,7 +33,7 @@ class TestOpenMeterIntegration: def test_openmeter_logger_missing_api_key(self): """Test that OpenMeterLogger raises exception when API key is missing""" os.environ.pop("OPENMETER_API_KEY", None) - with pytest.raises(Exception, match="Missing keys.*OPENMETER_API_KEY"): + with pytest.raises(Exception, match=r"Missing keys.*OPENMETER_API_KEY"): OpenMeterLogger() def test_common_logic_with_string_user(self): diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index 751b548adcd..aaaa43a0dc4 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -100,12 +100,12 @@ class TestEncodeUrlPathSegment: @pytest.mark.parametrize("value", ["", ".", "..", None]) def test_rejects_empty_and_dot_segments(self, value): - with pytest.raises(ValueError, match="resource_id (is required|cannot be a dot path segment)"): + with pytest.raises(ValueError, match=r"resource_id (is required|cannot be a dot path segment)"): encode_url_path_segment(value, field_name="resource_id") @pytest.mark.parametrize("value", ["../model", "model/../other", "/model"]) def test_rejects_dot_segments_in_multi_segment_paths(self, value): - with pytest.raises(ValueError, match="model (is required|cannot be a dot path segment)"): + with pytest.raises(ValueError, match=r"model (is required|cannot be a dot path segment)"): encode_url_path_segments(value, field_name="model") diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 25739a978d0..e519fab896a 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -929,7 +929,7 @@ class TestValidateEnvironmentAuthToken: config = AnthropicModelInfo() with mock_patch.dict("os.environ", {}, clear=True): with pytest.raises( - Exception, match="ANTHROPIC_API_KEY.*ANTHROPIC_AUTH_TOKEN" + Exception, match=r"ANTHROPIC_API_KEY.*ANTHROPIC_AUTH_TOKEN" ): config.validate_environment( headers={}, diff --git a/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py b/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py index a4a5f111513..0a47852d085 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_generic_chat.py @@ -106,7 +106,7 @@ class TestGenericToolCallErrors: ) def test_non_string_id_raises(self): - with pytest.raises(OCIError, match="id.*must be a string"): + with pytest.raises(OCIError, match=r"id.*must be a string"): adapt_messages_to_generic_oci_standard_tool_call( "assistant", [ @@ -126,7 +126,7 @@ class TestGenericToolCallErrors: ) def test_non_string_function_name_raises(self): - with pytest.raises(OCIError, match="function.name.*must be a string"): + with pytest.raises(OCIError, match=r"function\.name.*must be a string"): adapt_messages_to_generic_oci_standard_tool_call( "assistant", [ @@ -139,7 +139,7 @@ class TestGenericToolCallErrors: ) def test_non_string_arguments_raises(self): - with pytest.raises(OCIError, match="arguments.*must be a JSON string"): + with pytest.raises(OCIError, match=r"arguments.*must be a JSON string"): adapt_messages_to_generic_oci_standard_tool_call( "assistant", [ diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 5e854bbad70..cf52a5a856d 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -3,6 +3,7 @@ Test Vertex AI files handler functionality """ import asyncio +import re from types import MappingProxyType import pytest from unittest.mock import AsyncMock, patch @@ -180,7 +181,7 @@ class TestVertexAIFilesHandler: # Should raise ValueError for failed download with pytest.raises( ValueError, - match="Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt", + match=re.escape("Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt"), ): await self.handler.afile_content( file_content_request=file_content_request, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 813264c1feb..39a06c68913 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -33,7 +33,7 @@ def test_validate_vertex_location_accepts_valid(location): ["attacker.example/", "evil.com#", "us.attacker.example", "us/../..", "US", "us_central1", "-us", "", None], ) def test_validate_vertex_location_rejects_invalid(location): - with pytest.raises(ValueError, match="vertex_location is required|Invalid vertex_location format"): + with pytest.raises(ValueError, match=r"vertex_location is required|Invalid vertex_location format"): validate_vertex_location(location) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index c2858c84c6d..e504dd6e8a0 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -255,7 +255,7 @@ class TestRunAgent: assert calls["args"] == ("claude", "--resume") def test_missing_binary_raises_with_install_hint(self): - with pytest.raises(AgentRunError, match="claude.*Install it first"): + with pytest.raises(AgentRunError, match=r"claude.*Install it first"): run_agent( "http://localhost:4000", "sk-key", diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index ee4cf7fbb05..0ceec49de12 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -504,7 +504,7 @@ def test_apply_to_env_rejects_pinned_sqlite_direct_url(monkeypatch): monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db") monkeypatch.setenv("DIRECT_URL", "sqlite:///data/litellm.db") - with pytest.raises(RuntimeError, match="DIRECT_URL.*sqlite"): + with pytest.raises(RuntimeError, match=r"DIRECT_URL.*sqlite"): _apply() @@ -514,7 +514,7 @@ def test_apply_to_env_rejects_pinned_non_postgres_reader(monkeypatch): "DATABASE_URL_READ_REPLICA", "mysql://u:p@reader.example.com:3306/db" ) - with pytest.raises(RuntimeError, match="DATABASE_URL_READ_REPLICA.*mysql"): + with pytest.raises(RuntimeError, match=r"DATABASE_URL_READ_REPLICA.*mysql"): _apply() diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 47f01fe096d..b0b2c68e30d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -10,6 +10,7 @@ from __future__ import annotations import json import os +import re from types import SimpleNamespace from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock @@ -303,7 +304,7 @@ def test_resolve_routing_plugins_rejects_non_routing_plugin(tmp_path): plugin_file = tmp_path / "bad_rs_plugin.py" plugin_file.write_text("not_a_plugin = object()\n") - with pytest.raises(ValueError, match="router_settings.plugins"): + with pytest.raises(ValueError, match=re.escape("router_settings.plugins")): resolve_routing_plugins( plugin_paths=["bad_rs_plugin.not_a_plugin"], config_file_path=str(tmp_path / "config.yaml"), diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 76aeea894e1..a0ca33da737 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2,6 +2,7 @@ import asyncio import importlib import json import os +import re import socket import subprocess import sys @@ -2683,7 +2684,7 @@ async def test_get_config_from_file(tmp_path, monkeypatch): with open(empty_file, "w") as f: f.write("") # Write empty content which will result in None when loaded - with pytest.raises(Exception, match="Config cannot be None or Empty."): + with pytest.raises(Exception, match=re.escape("Config cannot be None or Empty.")): await proxy_config._get_config_from_file(str(empty_file)) # Test Case 5: Using global user_config_file_path when no config_file_path provided diff --git a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py index 0d7bcbe1e53..2928dea100e 100644 --- a/tests/test_litellm/sandbox/test_opensandbox_sandbox.py +++ b/tests/test_litellm/sandbox/test_opensandbox_sandbox.py @@ -490,7 +490,7 @@ async def test_create_waits_for_endpoint_resolution(monkeypatch): async def test_create_raises_when_endpoint_is_missing(): client = FakeHTTPClient(endpoint_json={"headers": {"X": "y"}}) - with pytest.raises(TimeoutError, match="execd endpoint.*not ready"): + with pytest.raises(TimeoutError, match=r"execd endpoint.*not ready"): await OpenSandboxSandboxConfig().acreate_sandbox( api_key="", api_base=TEST_API_BASE, ready_timeout=0, client=client ) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 4fdb5faf305..1580ec7f437 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -9,6 +9,7 @@ should still use the built-in pricing. import copy import os +import re import sys from unittest.mock import patch @@ -1892,7 +1893,7 @@ def test_a_reservation_without_a_declared_id_is_refused(): duplicate is permanent.""" anonymous = {k: v for k, v in _PTU_MODEL_INFO.items() if k != "id"} - with pytest.raises(ValueError, match="model_info.id is required"): + with pytest.raises(ValueError, match=re.escape("model_info.id is required")): _ptu_router(model_info=anonymous) @@ -1976,7 +1977,7 @@ def test_a_bare_yaml_date_bound_does_not_escape_the_id_rule(): windowed = {k: v for k, v in _PTU_MODEL_INFO.items() if k != "id"} - with pytest.raises(ValueError, match="model_info.id is required"): + with pytest.raises(ValueError, match=re.escape("model_info.id is required")): _ptu_router(model_info={**windowed, "ptu_effective_to": _dt.date(2027, 1, 1)}) From e917e4b307dfe40a21bce83552a8699d21904df0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:23:47 -0700 Subject: [PATCH 08/43] fix(model_cost): dedupe gemini-3.1-flash-lite-image and correct its capabilities The three lite-image keys landed on the deploy branch separately while this branch was open, so merging left every key defined twice in both price maps. The merge is clean as text and the file still parses, but JSON keeps the last occurrence of a repeated key, so the first copy's supported_endpoints, supported_modalities and supports_system_messages were being dropped without any error. Each key is now one entry, placed next to its gemini-3.1-flash-image sibling rather than at the end of the file. supports_reasoning goes to false on all three, matching every other Gemini image model. Leaving it off is not neutral: _supports_factory falls through to the vertex_ai provider config, which answers true, and reasoning_effort then gets forwarded to an image endpoint that rejects it. That was fixed for the rest of the family in 75dd70a678 and these entries had drifted back. Also fills in what the entries were missing against Google's published pricing: the Vertex implicit cache read rate, batch rates on the Vertex routes, and the pdf/video input flags. The two overlapping test files are folded into one, and the price map suite grows a duplicate-key guard so the next clean-but-lossy merge fails loudly. --- ...odel_prices_and_context_window_backup.json | 125 +----- model_prices_and_context_window.json | 125 +----- ...ini_3_1_flash_lite_image_model_metadata.py | 242 ----------- ...est_gemini_3_1_flash_lite_image_pricing.py | 377 ++++++++++++------ .../test_litellm/test_model_prices_schema.py | 26 ++ 5 files changed, 325 insertions(+), 570 deletions(-) delete mode 100644 tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a5f436ed35c..25c3af07ad4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19498,106 +19498,6 @@ }, "web_search_billing_unit": "per_query" }, - "gemini-3.1-flash-lite-image": { - "input_cost_per_image": 0.00028, - "input_cost_per_token": 2.5e-07, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 65536, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "image_generation", - "output_cost_per_image": 0.0336, - "output_cost_per_image_token": 3e-05, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": false, - "supports_prompt_caching": true, - "supports_response_schema": false, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_vision": true - }, - "gemini/gemini-3.1-flash-lite-image": { - "rpm": 1000, - "tpm": 4000000, - "input_cost_per_image": 0.00028, - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_batches": 1.25e-07, - "litellm_provider": "gemini", - "max_input_tokens": 65536, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "image_generation", - "output_cost_per_image": 0.0336, - "output_cost_per_image_token": 3e-05, - "output_cost_per_token": 1.5e-06, - "output_cost_per_token_batches": 7.5e-07, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": true, - "supports_prompt_caching": false, - "supports_response_schema": false, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_vision": true - }, - "vertex_ai/gemini-3.1-flash-lite-image": { - "input_cost_per_image": 0.00028, - "input_cost_per_token": 2.5e-07, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 65536, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "image_generation", - "output_cost_per_image": 0.0336, - "output_cost_per_image_token": 3e-05, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": false, - "supports_prompt_caching": true, - "supports_response_schema": false, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_vision": true - }, "gemini-3.1-flash-image": { "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, @@ -21557,7 +21457,6 @@ "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "rpm": 1000, - "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", "supported_endpoints": [ "/v1/chat/completions", @@ -21577,7 +21476,8 @@ "supports_reasoning": false, "supports_response_schema": false, "supports_system_messages": true, - "supports_vision": true + "supports_vision": true, + "tpm": 4000000 }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, @@ -41120,12 +41020,29 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], "supports_function_calling": false, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": false, "supports_response_schema": false, - "supports_vision": true, - "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a5f436ed35c..25c3af07ad4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19498,106 +19498,6 @@ }, "web_search_billing_unit": "per_query" }, - "gemini-3.1-flash-lite-image": { - "input_cost_per_image": 0.00028, - "input_cost_per_token": 2.5e-07, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 65536, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "image_generation", - "output_cost_per_image": 0.0336, - "output_cost_per_image_token": 3e-05, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": false, - "supports_prompt_caching": true, - "supports_response_schema": false, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_vision": true - }, - "gemini/gemini-3.1-flash-lite-image": { - "rpm": 1000, - "tpm": 4000000, - "input_cost_per_image": 0.00028, - "input_cost_per_token": 2.5e-07, - "input_cost_per_token_batches": 1.25e-07, - "litellm_provider": "gemini", - "max_input_tokens": 65536, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "image_generation", - "output_cost_per_image": 0.0336, - "output_cost_per_image_token": 3e-05, - "output_cost_per_token": 1.5e-06, - "output_cost_per_token_batches": 7.5e-07, - "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": true, - "supports_prompt_caching": false, - "supports_response_schema": false, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_vision": true - }, - "vertex_ai/gemini-3.1-flash-lite-image": { - "input_cost_per_image": 0.00028, - "input_cost_per_token": 2.5e-07, - "litellm_provider": "vertex_ai-language-models", - "max_input_tokens": 65536, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "image_generation", - "output_cost_per_image": 0.0336, - "output_cost_per_image_token": 3e-05, - "output_cost_per_token": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image" - ], - "supported_output_modalities": [ - "text", - "image" - ], - "supports_function_calling": false, - "supports_prompt_caching": true, - "supports_response_schema": false, - "supports_reasoning": true, - "supports_system_messages": true, - "supports_vision": true - }, "gemini-3.1-flash-image": { "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, @@ -21557,7 +21457,6 @@ "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "rpm": 1000, - "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", "supported_endpoints": [ "/v1/chat/completions", @@ -21577,7 +21476,8 @@ "supports_reasoning": false, "supports_response_schema": false, "supports_system_messages": true, - "supports_vision": true + "supports_vision": true, + "tpm": 4000000 }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, @@ -41120,12 +41020,29 @@ "output_cost_per_image_token": 3e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], "supports_function_calling": false, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": false, "supports_response_schema": false, - "supports_vision": true, - "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py deleted file mode 100644 index aa6f03a47ff..00000000000 --- a/tests/test_litellm/test_gemini_3_1_flash_lite_image_model_metadata.py +++ /dev/null @@ -1,242 +0,0 @@ -import json -from pathlib import Path - -import pytest - -import litellm -from litellm import completion_cost -from litellm.cost_calculator import cost_per_token -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.llms.gemini.image_generation.cost_calculator import ( - cost_calculator as gemini_image_generation_cost_calculator, -) -from litellm.llms.vertex_ai.image_generation.cost_calculator import ( - cost_calculator as vertex_image_generation_cost_calculator, -) -from litellm.types.utils import ( - CompletionTokensDetailsWrapper, - ImageObject, - ImageResponse, - ImageUsage, - ImageUsageInputTokensDetails, - ModelResponse, - PromptTokensDetailsWrapper, - Usage, -) - -REPO_ROOT = Path(__file__).parents[2] -MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" -BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" - -UNPREFIXED = "gemini-3.1-flash-lite-image" -GEMINI = "gemini/gemini-3.1-flash-lite-image" -VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" -ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) - -INPUT_COST = 2.5e-07 -INPUT_COST_BATCHES = 1.25e-07 -OUTPUT_TEXT_COST = 1.5e-06 -OUTPUT_TEXT_COST_BATCHES = 7.5e-07 -OUTPUT_IMAGE_TOKEN_COST = 3e-05 -OUTPUT_COST_PER_1K_IMAGE = 0.0336 -INPUT_COST_PER_IMAGE = 0.00028 -CACHE_READ_COST = 2.5e-08 -MAX_INPUT_TOKENS = 65536 -MAX_OUTPUT_TOKENS = 4096 -TOKENS_PER_1K_IMAGE = 1120 - - -def _load(path: Path) -> dict: - with open(path) as f: - return json.load(f) - - -@pytest.fixture -def local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.get_model_info.cache_clear() - try: - yield - finally: - litellm.model_cost = original_model_cost - litellm.get_model_info.cache_clear() - - -@pytest.mark.parametrize("model", ALL_KEYS) -def test_gemini_3_1_flash_lite_image_is_registered(model: str): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["mode"] == "image_generation" - assert info["input_cost_per_token"] == INPUT_COST - assert info["input_cost_per_token_batches"] == INPUT_COST_BATCHES - assert info["output_cost_per_token"] == OUTPUT_TEXT_COST - assert info["output_cost_per_token_batches"] == OUTPUT_TEXT_COST_BATCHES - assert info["output_cost_per_image"] == OUTPUT_COST_PER_1K_IMAGE - assert info["output_cost_per_image_token"] == OUTPUT_IMAGE_TOKEN_COST - assert info["max_input_tokens"] == MAX_INPUT_TOKENS - assert info["max_output_tokens"] == MAX_OUTPUT_TOKENS - assert info["max_tokens"] == MAX_OUTPUT_TOKENS - assert info["supports_reasoning"] is False - assert info["supports_response_schema"] is False - assert info["supports_vision"] is True - for field in ("supports_web_search", "search_context_cost_per_query", "web_search_billing_unit"): - assert field not in info - - -def test_gemini_3_1_flash_lite_image_provider_specific_fields(): - cost_map = _load(MAIN_PATH) - - unprefixed = cost_map[UNPREFIXED] - assert unprefixed["litellm_provider"] == "vertex_ai-language-models" - assert unprefixed["cache_read_input_token_cost"] == CACHE_READ_COST - assert unprefixed["input_cost_per_image"] == INPUT_COST_PER_IMAGE - assert unprefixed["supports_function_calling"] is False - assert unprefixed["supports_prompt_caching"] is True - assert unprefixed["supports_pdf_input"] is True - assert unprefixed["supports_video_input"] is True - assert unprefixed["supported_modalities"] == ["text", "image", "video"] - - gemini = cost_map[GEMINI] - assert gemini["litellm_provider"] == "gemini" - assert gemini["supports_function_calling"] is True - assert gemini["supports_prompt_caching"] is False - assert "cache_read_input_token_cost" not in gemini - assert gemini["supported_modalities"] == ["text", "image"] - assert gemini["supported_output_modalities"] == ["text", "image"] - assert gemini["rpm"] == 1000 - assert gemini["tpm"] == 4000000 - assert gemini["input_cost_per_image"] == INPUT_COST_PER_IMAGE - - vertex = cost_map[VERTEX] - assert vertex["litellm_provider"] == "vertex_ai-language-models" - assert vertex["cache_read_input_token_cost"] == CACHE_READ_COST - assert vertex["input_cost_per_image"] == INPUT_COST_PER_IMAGE - assert vertex["supports_function_calling"] is False - assert vertex["supports_prompt_caching"] is True - - -def test_one_k_image_price_matches_official_token_math(): - assert TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST == OUTPUT_COST_PER_1K_IMAGE - assert TOKENS_PER_1K_IMAGE * INPUT_COST == INPUT_COST_PER_IMAGE - - -@pytest.mark.parametrize("model", ALL_KEYS) -def test_backup_matches_main(model: str): - main_cost = _load(MAIN_PATH) - backup_cost = _load(BACKUP_PATH) - assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" - - -def test_gemini_prefix_routes_to_gemini(): - routed_model, provider, _, _ = get_llm_provider(model=GEMINI) - assert routed_model == UNPREFIXED - assert provider == "gemini" - - -def test_vertex_prefix_routes_to_vertex(): - routed_model, provider, _, _ = get_llm_provider(model=VERTEX) - assert routed_model == UNPREFIXED - assert provider == "vertex_ai" - - -def test_text_token_cost(local_model_cost_map): - prompt_cost, text_completion_cost = cost_per_token(model=GEMINI, prompt_tokens=1000, completion_tokens=500) - assert prompt_cost == pytest.approx(1000 * INPUT_COST) - assert text_completion_cost == pytest.approx(500 * OUTPUT_TEXT_COST) - - -def test_completion_cost_bills_one_k_image(local_model_cost_map): - response = ModelResponse() - response.model = UNPREFIXED - response.usage = Usage( - prompt_tokens=7, - completion_tokens=TOKENS_PER_1K_IMAGE, - total_tokens=7 + TOKENS_PER_1K_IMAGE, - completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=TOKENS_PER_1K_IMAGE, text_tokens=0), - ) - billed = completion_cost( - completion_response=response, - model=UNPREFIXED, - custom_llm_provider="vertex_ai", - ) - expected = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 7 * INPUT_COST - assert billed == pytest.approx(expected) - - -def test_image_tokens_are_not_billed_as_text(local_model_cost_map): - usage = Usage( - completion_tokens=1345, - prompt_tokens=10, - total_tokens=1355, - completion_tokens_details=CompletionTokensDetailsWrapper( - accepted_prediction_tokens=None, - audio_tokens=None, - reasoning_tokens=225, - rejected_prediction_tokens=None, - text_tokens=0, - image_tokens=TOKENS_PER_1K_IMAGE, - ), - prompt_tokens_details=PromptTokensDetailsWrapper( - audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None - ), - ) - - _, image_completion_cost = generic_cost_per_token( - model=UNPREFIXED, - usage=usage, - custom_llm_provider="vertex_ai", - ) - - expected_completion_cost = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 225 * OUTPUT_TEXT_COST - bugged_text_only_cost = 1345 * OUTPUT_TEXT_COST - assert image_completion_cost > bugged_text_only_cost * 2 - assert image_completion_cost == pytest.approx(expected_completion_cost) - - -def test_gemini_image_generation_uses_token_pricing(local_model_cost_map): - image_response = ImageResponse( - data=[ImageObject(b64_json="img1")], - usage=ImageUsage( - input_tokens=50 + TOKENS_PER_1K_IMAGE, - input_tokens_details=ImageUsageInputTokensDetails( - text_tokens=50, - image_tokens=TOKENS_PER_1K_IMAGE, - ), - output_tokens=TOKENS_PER_1K_IMAGE, - total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, - ), - ) - - cost = gemini_image_generation_cost_calculator(model=GEMINI, image_response=image_response) - expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST - assert cost == pytest.approx(expected) - assert cost != OUTPUT_COST_PER_1K_IMAGE - - -def test_vertex_image_generation_uses_token_pricing(local_model_cost_map): - image_response = ImageResponse( - data=[ImageObject(b64_json="img1")], - usage=ImageUsage( - input_tokens=50 + TOKENS_PER_1K_IMAGE, - input_tokens_details=ImageUsageInputTokensDetails( - text_tokens=50, - image_tokens=TOKENS_PER_1K_IMAGE, - ), - output_tokens=TOKENS_PER_1K_IMAGE, - total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, - ), - ) - - cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) - expected = (50 + TOKENS_PER_1K_IMAGE) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST - assert cost == pytest.approx(expected) - - -def test_vertex_image_generation_falls_back_to_flat_image_price(local_model_cost_map): - image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) - cost = vertex_image_generation_cost_calculator(model=UNPREFIXED, image_response=image_response) - assert cost == pytest.approx(2 * OUTPUT_COST_PER_1K_IMAGE) diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py index 67d6b9e76cf..276f54c116a 100644 --- a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py @@ -1,147 +1,284 @@ -"""Pricing entry for ``gemini-3.1-flash-lite-image`` (Google's Nano Banana 2 Lite). - -Google publishes: $0.25/1M input, $1.50/1M text output, and $30/1M image-output -tokens for the Lite image model (https://cloud.google.com/vertex-ai/generative-ai/pricing). -A 1K image is ~1120 output image tokens => ~$0.0336 / image. - -Without this entry, ``completion_cost`` raises "model isn't mapped yet" and Vertex -generateContent pass-through cost tracking silently logs $0. These tests pin the -values in both the primary price map and the ``litellm/`` backup, and verify -``get_model_info`` / ``completion_cost`` surface them. -""" - import json -import os +from pathlib import Path + +import pytest import litellm from litellm import completion_cost -from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) +from litellm.llms.vertex_ai.image_generation.cost_calculator import ( + cost_calculator as vertex_image_generation_cost_calculator, +) +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) -VARIANTS = [ - "gemini-3.1-flash-lite-image", - "gemini/gemini-3.1-flash-lite-image", - "vertex_ai/gemini-3.1-flash-lite-image", -] +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -EXPECTED = { - "input_cost_per_token": 2.5e-07, - "output_cost_per_token": 1.5e-06, - "output_cost_per_image_token": 3e-05, +UNPREFIXED = "gemini-3.1-flash-lite-image" +GEMINI = "gemini/gemini-3.1-flash-lite-image" +VERTEX = "vertex_ai/gemini-3.1-flash-lite-image" +ALL_KEYS = (UNPREFIXED, GEMINI, VERTEX) + +INPUT_COST = 2.5e-07 +INPUT_COST_BATCHES = 1.25e-07 +OUTPUT_TEXT_COST = 1.5e-06 +OUTPUT_TEXT_COST_BATCHES = 7.5e-07 +OUTPUT_IMAGE_TOKEN_COST = 3e-05 +OUTPUT_COST_PER_1K_IMAGE = 0.0336 +INPUT_COST_PER_IMAGE = 0.00028 +CACHE_READ_COST = 2.5e-08 +MAX_INPUT_TOKENS = 65536 +MAX_OUTPUT_TOKENS = 4096 +TOKENS_PER_1K_IMAGE = 1120 + +SHARED_FIELDS = { "mode": "image_generation", -} - -EXPECTED_CAPABILITIES = { - "max_output_tokens": 4096, - "max_tokens": 4096, + "input_cost_per_token": INPUT_COST, + "input_cost_per_token_batches": INPUT_COST_BATCHES, + "input_cost_per_image": INPUT_COST_PER_IMAGE, + "output_cost_per_token": OUTPUT_TEXT_COST, + "output_cost_per_token_batches": OUTPUT_TEXT_COST_BATCHES, + "output_cost_per_image": OUTPUT_COST_PER_1K_IMAGE, + "output_cost_per_image_token": OUTPUT_IMAGE_TOKEN_COST, + "max_input_tokens": MAX_INPUT_TOKENS, + "max_output_tokens": MAX_OUTPUT_TOKENS, + "max_tokens": MAX_OUTPUT_TOKENS, + "supported_endpoints": ["/v1/chat/completions", "/v1/completions", "/v1/batch"], + "supported_output_modalities": ["text", "image"], + "supports_reasoning": False, "supports_response_schema": False, - "supports_reasoning": True, + "supports_system_messages": True, + "supports_vision": True, } -EXPECTED_PER_ROUTE = { - "gemini-3.1-flash-lite-image": { - "supports_prompt_caching": True, - "supports_function_calling": False, - }, - "vertex_ai/gemini-3.1-flash-lite-image": { - "supports_prompt_caching": True, - "supports_function_calling": False, - }, - "gemini/gemini-3.1-flash-lite-image": { - "supports_prompt_caching": False, +VERTEX_ROUTE_FIELDS = { + "litellm_provider": "vertex_ai-language-models", + "cache_read_input_token_cost": CACHE_READ_COST, + "supported_modalities": ["text", "image", "video"], + "supports_function_calling": False, + "supports_pdf_input": True, + "supports_prompt_caching": True, + "supports_video_input": True, +} + +PER_ROUTE_FIELDS = { + UNPREFIXED: VERTEX_ROUTE_FIELDS, + VERTEX: VERTEX_ROUTE_FIELDS, + GEMINI: { + "litellm_provider": "gemini", + "supported_modalities": ["text", "image"], "supports_function_calling": True, - "input_cost_per_token_batches": 1.25e-07, - "output_cost_per_token_batches": 7.5e-07, + "supports_prompt_caching": False, + "rpm": 1000, + "tpm": 4000000, }, } +GROUNDING_FIELDS = ( + "supports_web_search", + "search_context_cost_per_query", + "web_search_billing_unit", +) -def _load_json(path: str) -> dict: + +def _load(path: Path) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) -def _backup_path() -> str: - return os.path.join( - os.path.dirname(litellm.__file__), - "model_prices_and_context_window_backup.json", +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("model", ALL_KEYS) +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_published_prices_are_registered(model: str, path: Path): + info = _load(path).get(model) + assert info is not None, f"{model} missing from {path.name}" + for field, value in SHARED_FIELDS.items(): + assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" + + +@pytest.mark.parametrize("model", ALL_KEYS) +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_per_route_capabilities_match_model_cards(model: str, path: Path): + info = _load(path)[model] + for field, value in PER_ROUTE_FIELDS[model].items(): + assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" + + +@pytest.mark.parametrize("model", ALL_KEYS) +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_grounding_fields_absent(model: str, path: Path): + info = _load(path)[model] + for field in GROUNDING_FIELDS: + assert field not in info, f"{model} should not define {field}" + + +@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) +def test_ai_studio_route_has_no_implicit_cache_price(path: Path): + assert "cache_read_input_token_cost" not in _load(path)[GEMINI] + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_backup_matches_main(model: str): + assert _load(BACKUP_PATH).get(model) == _load(MAIN_PATH).get(model) + + +def test_one_k_image_price_matches_official_token_math(): + assert TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST == pytest.approx(OUTPUT_COST_PER_1K_IMAGE) + assert TOKENS_PER_1K_IMAGE * INPUT_COST == pytest.approx(INPUT_COST_PER_IMAGE) + + +def test_gemini_prefix_routes_to_gemini(): + routed_model, provider, _, _ = get_llm_provider(model=GEMINI) + assert routed_model == UNPREFIXED + assert provider == "gemini" + + +def test_vertex_prefix_routes_to_vertex(): + routed_model, provider, _, _ = get_llm_provider(model=VERTEX) + assert routed_model == UNPREFIXED + assert provider == "vertex_ai" + + +def test_get_model_info_reports_published_costs(local_model_cost_map): + info = litellm.get_model_info(UNPREFIXED) + assert info["input_cost_per_token"] == INPUT_COST + assert info["output_cost_per_token"] == OUTPUT_TEXT_COST + assert info["cache_read_input_token_cost"] == CACHE_READ_COST + + +@pytest.mark.parametrize("model", ALL_KEYS) +def test_reasoning_params_are_not_offered_on_an_image_endpoint(model: str, local_model_cost_map): + assert litellm.supports_reasoning(model) is False + + +def test_text_token_cost(local_model_cost_map): + prompt_cost, text_completion_cost = cost_per_token( + model=GEMINI, prompt_tokens=1000, completion_tokens=500 + ) + assert prompt_cost == pytest.approx(1000 * INPUT_COST) + assert text_completion_cost == pytest.approx(500 * OUTPUT_TEXT_COST) + + +def test_completion_cost_bills_one_k_image(local_model_cost_map): + response = ModelResponse() + response.model = UNPREFIXED + response.usage = Usage( + prompt_tokens=7, + completion_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=7 + TOKENS_PER_1K_IMAGE, + completion_tokens_details=CompletionTokensDetailsWrapper( + image_tokens=TOKENS_PER_1K_IMAGE, text_tokens=0 + ), + ) + billed = completion_cost( + completion_response=response, + model=UNPREFIXED, + custom_llm_provider="vertex_ai", + ) + expected = TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 7 * INPUT_COST + assert billed == pytest.approx(expected) + + +def test_image_tokens_are_not_billed_as_text(local_model_cost_map): + usage = Usage( + completion_tokens=1345, + prompt_tokens=10, + total_tokens=1355, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=225, + rejected_prediction_tokens=None, + text_tokens=0, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None + ), + ) + + _, image_completion_cost = generic_cost_per_token( + model=UNPREFIXED, + usage=usage, + custom_llm_provider="vertex_ai", + ) + + expected_completion_cost = ( + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + 225 * OUTPUT_TEXT_COST + ) + bugged_text_only_cost = 1345 * OUTPUT_TEXT_COST + assert image_completion_cost > bugged_text_only_cost * 2 + assert image_completion_cost == pytest.approx(expected_completion_cost) + + +def _one_k_image_response() -> ImageResponse: + return ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=50 + TOKENS_PER_1K_IMAGE, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=50, + image_tokens=TOKENS_PER_1K_IMAGE, + ), + output_tokens=TOKENS_PER_1K_IMAGE, + total_tokens=50 + TOKENS_PER_1K_IMAGE + TOKENS_PER_1K_IMAGE, + ), ) -def _main_path() -> str: - return os.path.join( - os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" +def test_gemini_image_generation_uses_token_pricing(local_model_cost_map): + cost = gemini_image_generation_cost_calculator( + model=GEMINI, image_response=_one_k_image_response() ) + expected = ( + 50 + TOKENS_PER_1K_IMAGE + ) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) + assert cost != OUTPUT_COST_PER_1K_IMAGE -class TestGeminiFlashLiteImagePricingData: - """Both price maps must carry Google's published Nano Banana 2 Lite costs.""" - - def test_present_in_both_maps(self): - main = _load_json(_main_path()) - backup = _load_json(_backup_path()) - for key in VARIANTS: - for label, data in (("main", main), ("backup", backup)): - assert key in data, f"{key} missing from {label} JSON" - entry = data[key] - for field, value in EXPECTED.items(): - assert entry[field] == value, f"{key} {field} in {label}: {entry.get(field)} != {value}" - - def test_capabilities_match_model_cards(self): - main = _load_json(_main_path()) - backup = _load_json(_backup_path()) - for key in VARIANTS: - expected = {**EXPECTED_CAPABILITIES, **EXPECTED_PER_ROUTE[key]} - for label, data in (("main", main), ("backup", backup)): - entry = data[key] - for field, value in expected.items(): - assert entry[field] == value, f"{key} {field} in {label}: {entry.get(field)} != {value}" - - def test_grounding_fields_absent(self): - """Grounding with Google Search is unsupported on Lite, so no search pricing.""" - for path in (_main_path(), _backup_path()): - data = _load_json(path) - for key in VARIANTS: - for field in ( - "supports_web_search", - "search_context_cost_per_query", - "web_search_billing_unit", - ): - assert field not in data[key], f"{key} should not define {field}" - - def test_image_output_pricing_consistent(self): - """1120 image-output tokens * output_cost_per_image_token == output_cost_per_image.""" - backup = _load_json(_backup_path()) - entry = backup["gemini-3.1-flash-lite-image"] - assert round(1120 * entry["output_cost_per_image_token"], 6) == entry["output_cost_per_image"] +def test_vertex_image_generation_uses_token_pricing(local_model_cost_map): + cost = vertex_image_generation_cost_calculator( + model=UNPREFIXED, image_response=_one_k_image_response() + ) + expected = ( + 50 + TOKENS_PER_1K_IMAGE + ) * INPUT_COST + TOKENS_PER_1K_IMAGE * OUTPUT_IMAGE_TOKEN_COST + assert cost == pytest.approx(expected) -class TestGeminiFlashLiteImageModelInfo: - """``get_model_info`` and ``completion_cost`` must report the new costs.""" - - def test_get_model_info_and_cost(self): - original = litellm.model_cost - try: - litellm.model_cost = _load_json(_backup_path()) - info = litellm.get_model_info("gemini-3.1-flash-lite-image") - assert info["input_cost_per_token"] == EXPECTED["input_cost_per_token"] - assert info["output_cost_per_token"] == EXPECTED["output_cost_per_token"] - - resp = ModelResponse() - resp.model = "gemini-3.1-flash-lite-image" - resp.usage = Usage( - prompt_tokens=7, - completion_tokens=1120, - total_tokens=1127, - completion_tokens_details=CompletionTokensDetailsWrapper( - image_tokens=1120, text_tokens=0 - ), - ) - cost = completion_cost( - completion_response=resp, - model="gemini-3.1-flash-lite-image", - custom_llm_provider="vertex_ai", - ) - expected_cost = 1120 * 3e-05 + 7 * 2.5e-07 - assert abs(cost - expected_cost) < 1e-6, f"unexpected cost {cost}" - finally: - litellm.model_cost = original +def test_vertex_image_generation_falls_back_to_flat_image_price(local_model_cost_map): + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] + ) + cost = vertex_image_generation_cost_calculator( + model=UNPREFIXED, image_response=image_response + ) + assert cost == pytest.approx(2 * OUTPUT_COST_PER_1K_IMAGE) diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index cb7023e6c12..6114d1d8aba 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -11,6 +11,7 @@ import pytest REPO_ROOT = Path(__file__).parents[2] GENERATOR_PATH = REPO_ROOT / "ci_cd" / "generate_model_prices_schema.py" PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PRICES_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json" @@ -118,6 +119,31 @@ def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_sche assert validator.is_valid({"some-model": entry}) +def find_duplicate_keys(path: Path) -> list[str]: + duplicates: list[str] = [] + + def record_duplicates(pairs): + seen: set[str] = set() + for key, _ in pairs: + if key in seen: + duplicates.append(key) + seen.add(key) + return dict(pairs) + + json.loads(path.read_text(), object_pairs_hook=record_duplicates) + return duplicates + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_price_map_has_no_duplicate_keys(path: Path): + assert find_duplicate_keys(path) == [], ( + f"{path.name} defines the same key twice; JSON parsers keep only the last " + "occurrence, so the earlier entry's fields are silently dropped. This is what " + "a clean text merge of two branches that both added a model looks like: " + "deduplicate the keys into one entry" + ) + + DATED_VARIANT = re.compile(r"^(.*?)-(\d{4}-\d{2}-\d{2})$") SERVICE_TIER_SUFFIXES = ("_flex", "_priority") From 601d6ff2c8ad63c28a68bb8047fe44586adfb9b8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:26:18 -0700 Subject: [PATCH 09/43] [e2e] Pin the OpenAI websocket passthrough prefixes The websocket routes under /openai_passthrough and /openai had no e2e coverage, so nothing catches the regression from issue #36088, where both prefixes carried HTTP routes only and refused every upgrade with a 403 before a socket ever existed. Two tests cover it. The realtime one opens /openai_passthrough/v1/realtime and asserts OpenAI's own session.created frame comes back, which proves the route is registered and relayed upstream. The responses one asserts /openai/v1/responses accepts the upgrade, since a responses.connect socket waits for the client to speak first and has no opening frame to check. A refused upgrade is an HTTP response rather than a close frame, so both assert on the handshake. ws_base_url moves into e2e_config now that a second suite needs it --- .../coverage_registry/llm_conversational.yaml | 1 + .../llm_nonconversational.yaml | 1 + tests/e2e/e2e_config.py | 9 +++ .../e2e/llm_translation/passthrough_client.py | 58 +++++++++++++++++++ .../realtime/realtime_client.py | 9 +-- .../test_realtime_pipecat_audio_e2e.py | 2 +- .../realtime/test_realtime_pipecat_e2e.py | 2 +- .../llm_translation/test_passthrough_e2e.py | 52 +++++++++++++++++ 8 files changed, 124 insertions(+), 10 deletions(-) diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 44ed5765e38..1d4e1e028ca 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -64,6 +64,7 @@ - {id: llm.responses.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Streaming via /v1/responses"} - {id: llm.responses.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "response_api_endpoints/endpoints.py:26", rationale: "Cost logged on responses"} - {id: llm.responses.openai.passthrough.stream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "A streamed POST /openai_passthrough/v1/responses is costed and keyed by the provider response id; it used to log a zero-cost row under a random id (GitHub issue #36523)"} +- {id: llm.responses.openai.passthrough_websocket.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [works], fail_before_fix: proven, source: "test_passthrough_e2e.py", rationale: "A websocket upgrade on /openai/v1/responses is accepted, so a responses.connect client reaches OpenAI through the same prefix its HTTP traffic uses; the prefix carried no websocket route and refused the upgrade with a 403 (GitHub issue #36088)"} - {id: llm.responses.openai.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Tool calls via Responses API"} - {id: llm.responses.openai.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vision via Responses API"} - {id: llm.responses.anthropic.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Anthropic translation (smoke)"} diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 8ae7dd01b5a..f0a990a7f1f 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -44,6 +44,7 @@ - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} - {id: llm.realtime.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: realtime, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.19 / LIT-4778", rationale: "HTTP /v1/realtime/client_secrets returns an ephemeral credential"} +- {id: llm.realtime.openai.passthrough.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: openai, capability: basic, streaming: stream, assertions: [works], fail_before_fix: proven, source: "test_passthrough_e2e.py", rationale: "A websocket upgrade on /openai_passthrough/v1/realtime is accepted and relayed to OpenAI; only HTTP routes were registered under the prefix, so realtime clients were refused with a 403 before a socket existed (GitHub issue #36088)"} - {id: llm.vector_stores.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: vector_stores, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.17 / LIT-4778", rationale: "Vector store create/list/retrieve/delete lifecycle"} - {id: llm.vector_stores.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: vector_stores, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.17 / LIT-4778", rationale: "Vector store search and invalid id errors"} - {id: llm.bedrock_native.bedrock_converse.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: bedrock_native, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.12 / LIT-4778", rationale: "Bedrock native converse happy path"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 8bf39f6021f..0266c75e1a7 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -150,6 +150,15 @@ ANOMALY_SPEND_SETTLE_SECONDS = float( ) +def ws_base_url() -> str: + """PROXY_BASE_URL with its scheme swapped for the websocket one, so a suite + opening a socket points at the same proxy every HTTP suite uses.""" + for scheme, ws_scheme in (("https://", "wss://"), ("http://", "ws://")): + if PROXY_BASE_URL.startswith(scheme): + return ws_scheme + PROXY_BASE_URL[len(scheme) :] + return PROXY_BASE_URL + + def datadog_mcp_url(*, toolsets: str = "core") -> str: """Regional Datadog remote MCP endpoint for this process's DD_SITE. diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index e0dfae679a9..20a8592db20 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -11,9 +11,13 @@ native request models are co-located here because only this suite uses them. from __future__ import annotations from dataclasses import dataclass +from urllib.parse import urlencode from pydantic import BaseModel, Field +from websockets.exceptions import InvalidStatus +from websockets.sync.client import connect +from e2e_config import ws_base_url from proxy_client import ProxyClient from e2e_http import FileUploadForm, Headers, NoBody, Result, StreamingResponse from models import ChatMessage @@ -175,6 +179,26 @@ class OpenAIEmbeddingBody(BaseModel): input: str +class WebsocketEnvelope(BaseModel): + """The one field every provider event carries, so the first frame off a + passthrough socket identifies itself without the suite parsing raw dicts.""" + + type: str + + +class WebsocketHandshake(BaseModel): + """What the proxy did with a websocket upgrade on a passthrough prefix. + + `rejected_status` is the HTTP status of a refused upgrade: a prefix carrying no + websocket route answers 403, before any socket exists. `first_event_type` is the + type of the first frame an accepted socket delivered, which is None when the + provider waits for the client to speak first. + """ + + rejected_status: int | None = None + first_event_type: str | None = None + + class PassthroughBatchList(BaseModel): """OpenAI's own batch page, relayed verbatim. `object` is required so a body that is not an OpenAI list fails validation instead of passing vacuously.""" @@ -339,5 +363,39 @@ class PassthroughClient: ), ) + # ---- OpenAI websocket passthrough ---------------------------------- + # + # The same prefixes over an upgrade instead of a POST, for the provider APIs + # that only speak websocket (realtime, responses.connect). + + def openai_passthrough_websocket( + self, + key: str, + path: str, + *, + model: str | None = None, + open_timeout: float = 30.0, + first_event_timeout: float = 30.0, + ) -> WebsocketHandshake: + query = f"?{urlencode({'model': model})}" if model is not None else "" + try: + connection = connect( + f"{ws_base_url()}{path}{query}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=open_timeout, + ) + except InvalidStatus as rejected: + return WebsocketHandshake(rejected_status=rejected.response.status_code) + with connection: + try: + frame = connection.recv(timeout=first_event_timeout) + except TimeoutError: + return WebsocketHandshake() + text = frame.decode("utf-8") if isinstance(frame, bytes) else frame + return WebsocketHandshake( + first_event_type=WebsocketEnvelope.model_validate_json(text).type + ) + + def build_client(proxy: ProxyClient) -> PassthroughClient: return PassthroughClient(proxy=proxy) diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index e6c5c19cbd1..632a9cf7e57 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -21,20 +21,13 @@ from pydantic import BaseModel, ConfigDict from websockets.sync.client import connect from websockets.sync.connection import Connection -from e2e_config import PROXY_BASE_URL, unique_marker +from e2e_config import unique_marker, ws_base_url from proxy_client import ProxyClient from models import LiteLLMParamsBody _M = TypeVar("_M", bound=BaseModel) -def ws_base_url() -> str: - for scheme, ws_scheme in (("https://", "wss://"), ("http://", "ws://")): - if PROXY_BASE_URL.startswith(scheme): - return ws_scheme + PROXY_BASE_URL[len(scheme) :] - return PROXY_BASE_URL - - def realtime_ws_url(model: str) -> str: return f"{ws_base_url()}/v1/realtime?{urlencode({'model': model})}" diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py index 78955974cd5..2e9cfcfe648 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py @@ -27,10 +27,10 @@ from pathlib import Path import pytest +from e2e_config import ws_base_url from realtime_client import ( PROVIDERS, RealtimeProvider, - ws_base_url, realtime_model, ) diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py index 16628fd257a..f84ce197f88 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py @@ -25,10 +25,10 @@ import asyncio import pytest +from e2e_config import ws_base_url from realtime_client import ( PROVIDERS, RealtimeProvider, - ws_base_url, realtime_model, ) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index 17b0dbe1ae5..78b68de8389 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -28,6 +28,9 @@ from passthrough_client import ( ) EMBEDDING_MODEL = "text-embedding-3-small" +# Relayed to OpenAI untranslated, so this is OpenAI's own realtime model name +# rather than a gateway deployment alias. +REALTIME_MODEL = "gpt-realtime-2" pytestmark = pytest.mark.e2e @@ -339,3 +342,52 @@ class TestOpenAIPassthroughSpend: f"the embeddings row logged no prompt tokens, so whatever cost it carries " f"was not computed from the real usage: {row}" ) + + +class TestOpenAIPassthroughWebsocket: + """The OpenAI passthrough prefixes must answer a websocket upgrade, not only a POST. + + The customer points realtime and responses.connect clients at the same prefixes + their HTTP traffic already uses. Only HTTP routes were registered under those + prefixes, so every upgrade was refused before a socket existed and those clients + could not reach the gateway at all. A refused upgrade is an HTTP response, not a + close frame, which is why these assert on the handshake rather than a close code. + """ + + @pytest.mark.covers("llm.realtime.openai.passthrough.stream.works") + def test_realtime_upgrade_reaches_openai_through_the_passthrough_prefix( + self, client: PassthroughClient, scoped_key: str + ) -> None: + """Pins GitHub issue #36088: /openai_passthrough/v1/realtime accepts the + upgrade and relays OpenAI's own session, instead of rejecting it with a 403.""" + handshake = client.openai_passthrough_websocket( + scoped_key, "/openai_passthrough/v1/realtime", model=REALTIME_MODEL + ) + + assert handshake.rejected_status is None, ( + f"/openai_passthrough/v1/realtime refused the websocket upgrade with HTTP " + f"{handshake.rejected_status}, so a realtime client cannot connect through " + "the gateway at all" + ) + assert handshake.first_event_type == "session.created", ( + "the accepted socket never carried OpenAI's opening session event, so the " + f"upgrade was not relayed upstream; the first frame was " + f"{handshake.first_event_type}" + ) + + @pytest.mark.covers("llm.responses.openai.passthrough_websocket.stream.works") + def test_responses_upgrade_is_accepted_on_the_openai_prefix( + self, client: PassthroughClient, scoped_key: str + ) -> None: + """Pins GitHub issue #36088 on the second prefix: /openai/v1/responses upgrades + as well. A responses.connect socket waits for the client to speak first, so the + accepted handshake is the whole signal here.""" + handshake = client.openai_passthrough_websocket( + scoped_key, "/openai/v1/responses", first_event_timeout=2.0 + ) + + assert handshake.rejected_status is None, ( + f"/openai/v1/responses refused the websocket upgrade with HTTP " + f"{handshake.rejected_status}; the prefix relays this route over HTTP but " + "drops a responses.connect client before the socket opens" + ) From 2ad2bec0f078d3794dc897619802f9d6344b2673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:37:44 -0700 Subject: [PATCH 10/43] fix(model-costs): apply the Sol promo cut to the gpt-5.6 alias OpenAI's model page for gpt-5.6 serves the GPT-5.6 Sol page and states that the gpt-5.6 alias routes requests to GPT-5.6 Sol, so the alias bills at Sol's rates. The registry entry was left on the pre-cut rates while gpt-5.6-sol took the cut, overbilling gpt-5.6 callers by 25 percent on input and 50 percent on output. All 23 cost fields on gpt-5.6 now match gpt-5.6-sol, and a regression test pins the two entries together so they cannot drift again. --- ...odel_prices_and_context_window_backup.json | 44 +++++++++---------- model_prices_and_context_window.json | 44 +++++++++---------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 22 +++++++++- 3 files changed, 64 insertions(+), 46 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4f245fcdfbf..5542265075f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26041,33 +26041,33 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.6": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, - "cache_creation_input_token_cost_flex": 3.125e-06, - "cache_creation_input_token_cost_priority": 1.25e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, - "cache_read_input_token_cost_flex": 2.5e-07, - "cache_read_input_token_cost_priority": 1e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_above_272k_tokens_flex": 5e-06, - "input_cost_per_token_batches": 2.5e-06, - "input_cost_per_token_flex": 2.5e-06, - "input_cost_per_token_priority": 1e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_flex": 2e-07, + "cache_read_input_token_cost_priority": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_batches": 2e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 8e-06, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "output_cost_per_token_batches": 1.5e-05, - "output_cost_per_token_flex": 1.5e-05, - "output_cost_per_token_priority": 6e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_batches": 1e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4f245fcdfbf..5542265075f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26041,33 +26041,33 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.6": { - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, - "cache_creation_input_token_cost_above_272k_tokens_flex": 6.25e-06, - "cache_creation_input_token_cost_flex": 3.125e-06, - "cache_creation_input_token_cost_priority": 1.25e-05, - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, - "cache_read_input_token_cost_flex": 2.5e-07, - "cache_read_input_token_cost_priority": 1e-06, - "input_cost_per_token": 5e-06, - "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_above_272k_tokens_flex": 5e-06, - "input_cost_per_token_batches": 2.5e-06, - "input_cost_per_token_flex": 2.5e-06, - "input_cost_per_token_priority": 1e-05, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_flex": 2e-07, + "cache_read_input_token_cost_priority": 8e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_batches": 2e-06, + "input_cost_per_token_flex": 2e-06, + "input_cost_per_token_priority": 8e-06, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-05, - "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, - "output_cost_per_token_batches": 1.5e-05, - "output_cost_per_token_flex": 1.5e-05, - "output_cost_per_token_priority": 6e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_batches": 1e-05, + "output_cost_per_token_flex": 1e-05, + "output_cost_per_token_priority": 4e-05, "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { 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 0e7db195865..ef088fec5e2 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 @@ -913,7 +913,7 @@ def test_generic_cost_per_token_gpt55_pro(): @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost,cache_write_cost", [ - ("gpt-5.6", 5e-6, 3e-5, 5e-7, 6.25e-6), + ("gpt-5.6", 4e-6, 2e-5, 4e-7, 5e-6), ("gpt-5.6-sol", 4e-6, 2e-5, 4e-7, 5e-6), ("gpt-5.6-terra", 2e-6, 1.2e-5, 2e-7, 2.5e-6), ("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7), @@ -965,10 +965,28 @@ def test_generic_cost_per_token_gpt56( assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) +def test_gpt_5_6_alias_prices_match_sol(): + """Regression: the bare gpt-5.6 alias routes to GPT-5.6 Sol, so every cost field on + the two entries has to hold the same value. They drifted once before, when Sol took + its promotional cut and gpt-5.6 was left on the pre-cut rates, overbilling callers + who used the alias.""" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + alias = litellm.model_cost["gpt-5.6"] + sol = litellm.model_cost["gpt-5.6-sol"] + + cost_fields = sorted(field for field in sol if "cost" in field) + assert len(cost_fields) == 23 + + for field in cost_fields: + assert alias.get(field) == sol.get(field), field + + @pytest.mark.parametrize( "model,flex_long_input_cost,flex_long_output_cost", [ - ("gpt-5.6", 5e-6, 2.25e-5), + ("gpt-5.6", 4e-6, 1.5e-5), ("gpt-5.6-sol", 4e-6, 1.5e-5), ("gpt-5.6-terra", 2e-6, 9e-6), ("gpt-5.6-luna", 2e-7, 9e-7), From 9b27d7a977de64c8df4c5f1ec31791605fdd9e90 Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 21 Aug 2026 17:38:23 -0700 Subject: [PATCH 11/43] test(e2e): drop the REALTIME_MODEL comment --- tests/e2e/llm_translation/test_passthrough_e2e.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index 78b68de8389..7e6a8b25155 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -28,8 +28,6 @@ from passthrough_client import ( ) EMBEDDING_MODEL = "text-embedding-3-small" -# Relayed to OpenAI untranslated, so this is OpenAI's own realtime model name -# rather than a gateway deployment alias. REALTIME_MODEL = "gpt-realtime-2" pytestmark = pytest.mark.e2e From 71400e1029f52a1366ee68f775fcd166701572f7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:44:18 -0700 Subject: [PATCH 12/43] test(model-costs): record that azure gpt-5.6 keeps its own pricing The docstring claimed azure pricing mirrors the openai family, which stopped being true when gpt-5.6-sol took its promotional cut and azure did not. Azure publishes no sol rate of its own today, so the entries stay where they are. --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 ef088fec5e2..0e5a2c05fd1 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 @@ -1136,8 +1136,10 @@ def test_generic_cost_per_token_gpt56_cyber( def test_generic_cost_per_token_azure_gpt56( model, input_cost, output_cost, cache_read_cost ): - """Azure gpt-5.6 (global + us/eu regional): pricing mirrors the openai - family for global deployments and carries the standard 10% regional uplift. + """Azure gpt-5.6 (global + us/eu regional): Azure prices this family on its own + schedule and carries the standard 10% regional uplift on top. It did not take the + promotional cut OpenAI applied to gpt-5.6-sol, so these rates deliberately sit + above the openai ones and must not be lowered to match them. """ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") From 6b088f4bb1ebdec8e3604a9d509426ff424dd903 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 21 Aug 2026 17:48:57 -0700 Subject: [PATCH 13/43] style: wrap the escaped messages under 120 columns --- tests/local_testing/test_embedding.py | 10 ++++++++-- .../vertex_ai/files/test_vertex_ai_files_handler.py | 5 ++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 677f8b10cca..fbbe83ada30 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -537,13 +537,19 @@ def test_demo_tokens_as_input_to_embeddings_fails_for_titan(): with pytest.raises( litellm.BadRequestError, - match=re.escape('litellm.BadRequestError: BedrockException - {"message":"Malformed input request: expected type: String, found: JSONArray, please reformat your input and try again."}'), + match=re.escape( + 'litellm.BadRequestError: BedrockException - {"message":"Malformed input request: ' + 'expected type: String, found: JSONArray, please reformat your input and try again."}' + ), ): litellm.embedding(model="amazon.titan-embed-text-v1", input=[[1]]) with pytest.raises( litellm.BadRequestError, - match=re.escape('litellm.BadRequestError: BedrockException - {"message":"Malformed input request: expected type: String, found: Integer, please reformat your input and try again."}'), + match=re.escape( + 'litellm.BadRequestError: BedrockException - {"message":"Malformed input request: ' + 'expected type: String, found: Integer, please reformat your input and try again."}' + ), ): litellm.embedding( model="amazon.titan-embed-text-v1", diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index cf52a5a856d..0a44f0a9a74 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -181,7 +181,10 @@ class TestVertexAIFilesHandler: # Should raise ValueError for failed download with pytest.raises( ValueError, - match=re.escape("Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt"), + match=re.escape( + "Failed to download file from GCS: " + "gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt" + ), ): await self.handler.afile_content( file_content_request=file_content_request, From 1a55418ea22d5414b60430d8947ba3b1ec089b34 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:56:22 -0700 Subject: [PATCH 14/43] test(model-costs): use the local_model_cost_map fixture in the alias test The new test set LITELLM_LOCAL_MODEL_COST_MAP and reassigned litellm.model_cost by hand, leaking both into every test that ran after it and skipping the get_model_info cache clear. The conftest fixture already does this properly and restores the original map on the way out. --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 0e5a2c05fd1..940c6259070 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 @@ -965,14 +965,11 @@ def test_generic_cost_per_token_gpt56( assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) -def test_gpt_5_6_alias_prices_match_sol(): +def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map): """Regression: the bare gpt-5.6 alias routes to GPT-5.6 Sol, so every cost field on the two entries has to hold the same value. They drifted once before, when Sol took its promotional cut and gpt-5.6 was left on the pre-cut rates, overbilling callers who used the alias.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - alias = litellm.model_cost["gpt-5.6"] sol = litellm.model_cost["gpt-5.6-sol"] From 7cb100af63023cd054730caa179ba3fb3118a325 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 21 Aug 2026 18:16:49 -0700 Subject: [PATCH 15/43] fix(mcp): resolve admin OAuth sessions to the same server set the connect page shows (#37900) * fix(mcp): resolve admin OAuth sessions to the same server set the connect page shows * fix(mcp): bind admitted admin rows through the entitlement ceiling, not the credential scope clause --- .../mcp_server/auth/user_api_key_auth_mcp.py | 40 +++++- .../mcp_server/mcp_server_manager.py | 41 +++---- .../auth/test_user_api_key_auth_mcp.py | 116 +++++++++++++++--- .../mcp_server/test_mcp_server_manager.py | 69 +++++++++-- 4 files changed, 208 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index c1248cafac5..7d85f3c4908 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -39,6 +39,7 @@ from litellm.proxy._types import ( SpecialMCPServerName, SpecialMCPServerNames, UserAPIKeyAuth, + user_api_key_has_admin_view, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import ( @@ -1785,11 +1786,14 @@ class MCPRequestHandler: global_mcp_server_manager, ) - # An OPEN channel (allow_all_keys, the user's own BYOM) makes the server REACHABLE through the - # user, though no grant source names it — without this the union returns [], listable but - # uninvokable. Reachability is ALL it confers, NOT a ceiling waiver: the user's own - # mcp_tool_permissions and org tool ceiling still bind, exactly as a key's do on an allow_all server. - reachable_via_open_channel: Final = server_id in await global_mcp_server_manager.operator_open_server_ids(auth) + # An OPEN channel (allow_all_keys, the user's own BYOM, an unscoped admin-view role) makes the + # server REACHABLE through the user, though no grant source names it — without this the union + # returns [], listable but uninvokable. Reachability is ALL it confers, NOT a ceiling waiver: + # the user's own mcp_tool_permissions and org tool ceiling still bind, exactly as a key's do + # on an allow_all server or an admin key's do on any server. + reachable_via_open_channel: Final = server_id in await global_mcp_server_manager.operator_open_server_ids( + auth + ) or await MCPRequestHandler.admin_view_unscoped(auth) allowed: Final[set[str]] = set() for source, granted in await MCPRequestHandler.admitted_source_grants(auth): @@ -2723,6 +2727,32 @@ class MCPRequestHandler: entitled_servers: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth) return entitled_servers is None or len(entitled_servers) > 0 + @staticmethod + async def admin_view_unscoped(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool: + """Whether this principal's admin-view role grants the unscoped MCP resolution, whatever + credential carries it (admin key, dashboard session, or OAuth-admitted session subject). + + Two bounds disqualify, one per ownership of the row. A CREDENTIAL's explicit + ``object_permission.mcp_servers`` scope wins even for admins, including the empty list. An + admitted subject's object_permission is the user's own row, whose ``mcp_servers`` column is + [] by DB default, so for that shape the row binds through the entitlement ceiling instead + (any non-empty entitlement, or an unresolved one, disqualifies), exactly as + ``operator_open_server_ids`` reads the same row. The one owner of this predicate: the + server-axis registry resolution in ``get_allowed_mcp_servers`` and the tools-axis open + channel in ``_resolve_admitted_subject_tools`` both consult it, so the two axes cannot + disagree.""" + if user_api_key_auth is None or not user_api_key_has_admin_view(user_api_key_auth): + return False + object_permission: Final = user_api_key_auth.object_permission + credential_scoped: Final = ( + not _is_mcp_admitted_user_subject(user_api_key_auth) + and object_permission is not None + and object_permission.mcp_servers is not None + ) + if credential_scoped: + return False + return not await MCPRequestHandler._user_places_mcp_ceiling(user_api_key_auth) + @staticmethod async def _apply_user_tool_ceiling( allowed_tools: Sequence[str] | None, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index dbe97dd5bce..7ab26db0f3e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2943,17 +2943,14 @@ class MCPServerManager: 2. If admin and no object_permission, return all servers 3. Otherwise, use standard permission checks """ - from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view - allow_all_server_ids: Final = self.get_allow_all_keys_server_ids() # A keyless admitted subject is resolved per grant source, and channel decisions that are # absolute for a scoped KEY credential are not absolute for it: its own opt-out silences its - # own source (handled per source in the resolver), never its teams' grants, and its admin - # role does not swallow the grant model — a session bearer is a third-party client - # credential, not the dashboard, so an admin signing in through the connect flow gets their - # grants like anyone else rather than handing the client the full registry ahead of every - # per-team org ceiling. + # own source (handled per source in the resolver), never its teams' grants. Its admin role + # rides the HUMAN, not the credential: an admin's session resolves the same registry their + # dashboard shows (connect-page parity), bounded like an admin key by explicit + # object_permission scope, the entitlement ceiling, and the session resource scope below. is_admitted_subject: Final = _is_mcp_admitted_user_subject(user_api_key_auth) # The key explicitly opted out of every MCP server. Return zero before @@ -2982,26 +2979,16 @@ class MCPServerManager: ) try: - # If admin but NO explicit object permission, get all servers (never for an admitted - # subject — see is_admitted_subject above) - if ( - user_api_key_auth - and not is_admitted_subject - and _user_has_admin_view(user_api_key_auth) - and not has_explicit_object_permission - # An entitlement attached to the HUMAN binds them whatever their role: it is the - # person's scope, not the credential's, so an admin role is not a waiver of it. An - # UNRESOLVED entitlement also skips the shortcut, so the resolver denies rather than - # handing over the whole registry on a transient fault. - and not await MCPRequestHandler._user_places_mcp_ceiling(user_api_key_auth) - ): - verbose_logger.debug("Admin user without explicit object_permission - returning all servers") - return list(self.get_registry().keys()) - - # Get allowed servers from object permissions (respects object_permission even for admins) - allowed_mcp_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) - verbose_logger.debug("Allowed MCP Servers for user api key auth: %s", allowed_mcp_servers) - combined_servers: Final = set(allowed_mcp_servers) + # Admin view with no explicit object permission and no entitlement ceiling resolves the + # whole registry, for keys AND admitted session subjects alike (one predicate owns the + # question). Seeded into the union rather than returned early so the session resource + # scope below still bounds a per-server envelope held by an admin. + combined_servers: Final = ( + set(self.get_registry().keys()) + if await MCPRequestHandler.admin_view_unscoped(user_api_key_auth) + else set(await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)) + ) + verbose_logger.debug("Allowed MCP Servers for user api key auth: %s", combined_servers) combined_servers.update( await self.operator_open_server_ids( user_api_key_auth, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 3bd615a6a33..17c4d773981 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -6430,9 +6430,7 @@ class TestAggregateGatewayDcrChallenge: assert _gateway_dcr_challenge_target("/mcp/srv", None, None) == expected, resolved assert _gateway_dcr_challenge_target("/mcp/a,b", None, None) is None assert _gateway_dcr_challenge_target("/mcp", None, None) is None - with patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_mgr: + with patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr: mock_mgr.get_mcp_server_by_name.return_value = _server(MCPAuth.oauth2) assert _gateway_dcr_challenge_target("/mcp/srv", ["other"], None) is None @@ -7073,25 +7071,109 @@ class TestUserSubjectTeamUnion: assert await manager.operator_open_server_ids(admitted) == {"srv-byom"} assert await manager.operator_open_server_ids(scoped_key) == set(), "explicit key scope still suppresses BYOM" - async def test_admitted_admin_is_scoped_to_grants_not_full_registry(self): - """The wrapper's admin short-circuit hands the FULL registry to any admin-role auth before - the grant union or the per-team org ceilings run. A session bearer is a third-party client - credential, not the dashboard: an admin signing in through the connect flow gets their - grants like anyone else. A real admin key keeps the dashboard behavior unchanged.""" + @pytest.mark.parametrize( + "role", ["PROXY_ADMIN", "PROXY_ADMIN_VIEW_ONLY"], ids=["proxy_admin", "proxy_admin_view_only"] + ) + async def test_admitted_admin_gets_registry_like_an_admin_key(self, role): + """Connect-page parity: admin view rides the HUMAN, not the credential. An admitted session + subject with an admin-view role resolves the same full registry an admin KEY does, so the + servers the dashboard shows an admin are the servers their OAuth session serves. Regression + pin for the customer report where an admin's Claude Code session showed zero tools.""" + from litellm.proxy._types import LitellmUserRoles + + manager = self._manager_with(["srv-granted", "srv-secret"]) + admitted = _make_admitted_subject("admin-user") + admitted.user_role = LitellmUserRoles[role] + key_admin = UserAPIKeyAuth(user_id="admin-user", api_key="sk-hash", user_role=LitellmUserRoles[role]) + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])): + admitted_view = set(await manager.get_allowed_mcp_servers(admitted)) + key_admin_view = set(await manager.get_allowed_mcp_servers(key_admin)) + assert admitted_view == {"srv-granted", "srv-secret"}, "an admitted admin resolves the registry" + assert key_admin_view == admitted_view, "session and key admin views must be identical" + + async def test_admitted_admin_explicit_scope_still_wins(self): + """An admin whose own user row names servers is entitlement-bound whatever their role: the + row binds through the ceiling for an admitted subject (a user row's mcp_servers is the + human's grant list, not a credential scope), so the registry seed must not fire. A KEY + carrying an explicit scope disqualifies directly, empty list included.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles + + manager = self._manager_with(["srv-granted", "srv-secret"]) + admitted = _make_admitted_subject("admin-user", own_servers=["srv-granted"]) + admitted.user_role = LitellmUserRoles.PROXY_ADMIN + with ( + patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", AsyncMock(return_value=["srv-granted"]) + ), + patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])), + ): + assert set(await manager.get_allowed_mcp_servers(admitted)) == {"srv-granted"} + + scoped_key = UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-hash", + user_role=LitellmUserRoles.PROXY_ADMIN, + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-k", mcp_servers=[]), + ) + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=[])): + assert await manager.get_allowed_mcp_servers(scoped_key) == [] + + async def test_admitted_admin_db_default_empty_scope_still_gets_registry(self): + """The admitted subject's object_permission is the user's own row, whose mcp_servers column + is [] by DB default whenever the row exists for any other field: default noise, never an + explicit scope. The registry seed must fire through it, or every admin with a shared + permission row keeps resolving zero servers while their dashboard shows all of them.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles + + manager = self._manager_with(["srv-granted", "srv-secret"]) + admitted = _make_admitted_subject("admin-user") + admitted.user_role = LitellmUserRoles.PROXY_ADMIN + admitted.object_permission = LiteLLM_ObjectPermissionTable(object_permission_id="op-u", mcp_servers=[]) + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=[])): + assert set(await manager.get_allowed_mcp_servers(admitted)) == {"srv-granted", "srv-secret"} + + async def test_non_admin_admitted_subject_never_gets_registry(self): + """The negative control for the registry seed: a plain admitted subject with no admin-view + role resolves only their grant union, however many servers the registry holds.""" + manager = self._manager_with(["srv-granted", "srv-secret"]) + plain = _make_admitted_subject("plain-user") + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])): + assert set(await manager.get_allowed_mcp_servers(plain)) == {"srv-granted"} + + async def test_admitted_admin_entitlement_ceiling_disables_registry(self): + """An entitlement ceiling, including an UNRESOLVED one, binds the human whatever their role: + the registry seed must not fire on a transient fault, and the grant union answers instead.""" from litellm.proxy._types import LitellmUserRoles manager = self._manager_with(["srv-granted", "srv-secret"]) admitted = _make_admitted_subject("admin-user") admitted.user_role = LitellmUserRoles.PROXY_ADMIN - with patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])): - admitted_view = set(await manager.get_allowed_mcp_servers(admitted)) - key_admin_view = set( - await manager.get_allowed_mcp_servers( - UserAPIKeyAuth(user_id="admin-user", api_key="sk-hash", user_role=LitellmUserRoles.PROXY_ADMIN) - ) - ) - assert admitted_view == {"srv-granted"}, "an admitted admin gets their grants, not the registry" - assert key_admin_view == {"srv-granted", "srv-secret"}, "admin KEY behavior must be unchanged" + with ( + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_user", AsyncMock(return_value=None)), + patch.object(MCPRequestHandler, "get_allowed_mcp_servers", AsyncMock(return_value=["srv-granted"])), + ): + assert set(await manager.get_allowed_mcp_servers(admitted)) == {"srv-granted"} + + async def test_admitted_admin_tools_ride_own_source_on_ungranted_server(self): + """Admin view is an open channel on the tools axis too: the user's OWN source resolves the + tools for a server no grant names, so an admin session's registry-wide servers are invokable + rather than listable-but-uninvokable. A non-admin subject on the same server stays denied. + An admin whose row carries any entitlement never reaches this channel: the ceiling clause + disqualifies the predicate first, so their own tool permissions keep binding on the grants path.""" + from litellm.proxy._types import LitellmUserRoles + + admin = _make_admitted_subject("admin-user") + admin.user_role = LitellmUserRoles.PROXY_ADMIN + plain = _make_admitted_subject("plain-user") + with self._patch(teams_by_id={}, user_teams=[]): + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager.operator_open_server_ids", + AsyncMock(return_value=set()), + ): + admin_tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-any", admin) + plain_tools = await MCPRequestHandler.get_allowed_tools_for_server("srv-any", plain) + assert admin_tools is None, "admin channel resolves allow-all through the user's own source" + assert plain_tools == [], "a non-admin subject with no granting source stays denied" async def test_admitted_opt_out_via_wrapper_keeps_team_servers(self): """The wrapper's no_mcp_servers early-return is a KEY rule (a scoped credential's opt-out is diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5ee8143fb8e..09e8c78a3f8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -4846,7 +4846,9 @@ class TestMCPServerManager: @staticmethod def _manager_with_deepwiki_and_huggingface() -> MCPServerManager: manager = MCPServerManager() - deepwiki = MCPServer(server_id="deepwiki-id", name="deepwiki", server_name="deepwiki", transport=MCPTransport.http) + deepwiki = MCPServer( + server_id="deepwiki-id", name="deepwiki", server_name="deepwiki", transport=MCPTransport.http + ) huggingface = MCPServer( server_id="huggingface-id", name="huggingface", server_name="huggingface", transport=MCPTransport.http ) @@ -4867,8 +4869,14 @@ class TestMCPServerManager: with pytest.raises(ValueError, match="Tool hub_repo_search not found"): manager._resolve_mcp_server_for_tool_call("deepwiki", "hub_repo_search") - assert manager._resolve_mcp_server_for_tool_call("deepwiki", "read_wiki_structure") is manager.registry["deepwiki-id"] - assert manager._resolve_mcp_server_for_tool_call("huggingface", "hub_repo_search") is manager.registry["huggingface-id"] + assert ( + manager._resolve_mcp_server_for_tool_call("deepwiki", "read_wiki_structure") + is manager.registry["deepwiki-id"] + ) + assert ( + manager._resolve_mcp_server_for_tool_call("huggingface", "hub_repo_search") + is manager.registry["huggingface-id"] + ) def test_get_mcp_server_from_tool_name_rejects_other_servers_prefix(self): manager = self._manager_with_deepwiki_and_huggingface() @@ -4876,7 +4884,9 @@ class TestMCPServerManager: assert manager._get_mcp_server_from_tool_name("huggingface-read_wiki_structure") is None assert manager._get_mcp_server_from_tool_name("deepwiki-hub_repo_search") is None assert manager._get_mcp_server_from_tool_name("deepwiki-read_wiki_structure") is manager.registry["deepwiki-id"] - assert manager._get_mcp_server_from_tool_name("huggingface-hub_repo_search") is manager.registry["huggingface-id"] + assert ( + manager._get_mcp_server_from_tool_name("huggingface-hub_repo_search") is manager.registry["huggingface-id"] + ) def test_resolve_mcp_server_for_tool_call_shared_bare_name_resolves_via_own_prefixed_spelling(self): manager = MCPServerManager() @@ -10500,6 +10510,39 @@ class TestSessionResourceScopeIntersect: assert MCPServerManager._admitted_session_resource_scope(self._admitted_auth("b")) == "b" + @pytest.mark.asyncio + async def test_admin_registry_seed_still_bounded_by_session_resource_scope(self): + """The admin-view registry seed flows through the same scoped exit as every union: a + session envelope sealed to one server never widens past it, even held by an admin whose + role resolves the whole registry. Pin for the connect-page-parity change; without the + single-exit shape, the old early return would hand a per-server bearer the registry.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + from litellm.proxy._types import LitellmUserRoles + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + manager = MCPServerManager() + for sid in ("granted-id", "other-id"): + manager.registry[sid] = MCPServer( + server_id=sid, name=sid, server_name=sid, url="https://example.com/mcp", transport=MCPTransport.http + ) + auth = self._admitted_auth("granted-id") + auth.user_role = LitellmUserRoles.PROXY_ADMIN + with ( + patch.object(MCPServerManager, "get_allow_all_keys_server_ids", return_value=[]), + patch.object( + MCPServerManager, + "_get_active_submitted_mcp_server_ids_for_user", + new_callable=AsyncMock, + return_value=[], + ), + ): + assert await manager.get_allowed_mcp_servers(auth) == ["granted-id"] + auth.mcp_session_resource_server_id = None + assert set(await manager.get_allowed_mcp_servers(auth)) == {"granted-id", "other-id"} + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_scopes_past_operator_open_union(self): """The intersect applies AFTER the operator-open (allow_all_keys) union, so a scoped @@ -10519,7 +10562,12 @@ class TestSessionResourceScopeIntersect: new_callable=AsyncMock, return_value=["granted-id", "other-id"], ), - patch.object(MCPServerManager, "_get_active_submitted_mcp_server_ids_for_user", new_callable=AsyncMock, return_value=[]), + patch.object( + MCPServerManager, + "_get_active_submitted_mcp_server_ids_for_user", + new_callable=AsyncMock, + return_value=[], + ), ): allowed = await manager.get_allowed_mcp_servers(auth) assert allowed == ["granted-id"] @@ -10531,7 +10579,12 @@ class TestSessionResourceScopeIntersect: new_callable=AsyncMock, side_effect=RuntimeError("resolver down"), ), - patch.object(MCPServerManager, "_get_active_submitted_mcp_server_ids_for_user", new_callable=AsyncMock, return_value=[]), + patch.object( + MCPServerManager, + "_get_active_submitted_mcp_server_ids_for_user", + new_callable=AsyncMock, + return_value=[], + ), ): fallback = await manager.get_allowed_mcp_servers(auth) assert fallback == ["granted-id"] @@ -10645,9 +10698,7 @@ class TestClientForwardedDiscoveryFailureIsNotFatal: @pytest.mark.parametrize("auth_type", [MCPAuth.true_passthrough, MCPAuth.oauth_delegate]) @pytest.mark.asyncio - async def test_client_forwarded_servers_keep_discovering_their_front_door_endpoints( - self, auth_type: MCPAuthType - ): + async def test_client_forwarded_servers_keep_discovering_their_front_door_endpoints(self, auth_type: MCPAuthType): """Exempting these modes from the FAILURE must not exempt them from discovery itself. ``/authorize``, ``/token`` and ``/register`` read the discovered endpoints for these servers From 770bd40f5de3f38ccc37d1b5bf5b942aecdf77c7 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 21 Aug 2026 18:23:25 -0700 Subject: [PATCH 16/43] fix(anthropic_messages): gate sampling params on /v1/messages like /chat/completions (#37868) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(anthropic_messages): gate sampling params on /v1/messages like /chat/completions /v1/messages forwarded temperature/top_p/top_k raw to models that removed sampling params (supports_sampling_params: false — Claude 4.7+/Fable 5), producing provider 400s that router fallbacks mask as silent model downgrades. The chat path already gates these via AnthropicModelInfo._apply_sampling_param; reuse it in get_requested_anthropic_messages_optional_param so both endpoints agree: drop under drop_params, else raise the clean client-side 400. Fixes #35053 * test(anthropic): drive new sampling-param tests off the kwarg, not the global The five tests added here set `litellm.drop_params = True` under a manual try/finally. That trips TQ005 (module-global mutation, 10 new violations over the ceiling) and it leaks process-wide if the finally is ever skipped, which is what the save/restore conftest exists to paper over. `get_requested_anthropic_messages_optional_param` already takes `drop_params` as a kwarg, and that is the path /v1/messages actually uses, so pass it directly. `monkeypatch.setattr` pins the global to False so each test proves the per-request flag alone is sufficient and cannot pass on a leaked global. Verified: TQ gate clean, all 10 tests pass, and the 3 that assert the new gating still fail with the fix in utils.py reverted. --------- Co-authored-by: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> --- .../messages/utils.py | 11 ++++ .../test_request_optional_param_utils.py | 60 +++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index d0dc5d527fe..02d82887dde 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -44,6 +44,7 @@ class AnthropicMessagesRequestUtils: filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} if model is not None: from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.common_utils import AnthropicModelInfo AnthropicConfig._maybe_drop_speed_param( model=model, @@ -51,6 +52,16 @@ class AnthropicMessagesRequestUtils: drop_params=drop_params, custom_llm_provider=custom_llm_provider, ) + for param in ("temperature", "top_p", "top_k"): + if param in filtered_params: + AnthropicModelInfo._apply_sampling_param( # pyright: ignore[reportPrivateUsage] # same gating the /chat/completions path applies; forking it would drift + optional_params=filtered_params, + model=model, + param=param, + value=filtered_params.pop(param), + drop_params=drop_params, + output_key=param, + ) return cast(AnthropicMessagesRequestOptionalParams, filtered_params) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py index f0252e13336..dc2e107928f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_request_optional_param_utils.py @@ -6,6 +6,8 @@ Regression tests for the /v1/messages request-parse fast paths: while resolving the (static) type hints only once per process. """ +import pytest + import litellm from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( AnthropicMessagesRequestUtils, @@ -88,3 +90,61 @@ def test_drop_params_keeps_speed_for_supporting_model(): litellm.drop_params = original assert result == {"speed": "fast"} + + +def test_drop_params_strips_sampling_params_for_unsupported_model(monkeypatch): + # claude-opus-4-7 has supports_sampling_params: false in the model map; the + # API 400s on these rather than ignoring them. + monkeypatch.setattr(litellm, "drop_params", False) + result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"temperature": 0.3, "top_p": 0.9, "top_k": 40, "stream": True}, + model="claude-opus-4-7", + drop_params=True, + ) + + assert result == {"stream": True} + + +def test_drop_params_strips_sampling_params_for_provider_prefixed_model(monkeypatch): + # Vertex-routed ids must resolve the same capability flag. + monkeypatch.setattr(litellm, "drop_params", False) + result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"temperature": 0.3, "top_p": 0.9, "top_k": 40}, + model="vertex_ai/claude-opus-4-7", + drop_params=True, + ) + + assert result == {} + + +def test_sampling_params_kept_for_supporting_model(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"temperature": 0.3, "top_p": 0.9, "top_k": 40}, + model="claude-sonnet-4-6", + drop_params=True, + ) + + assert result == {"temperature": 0.3, "top_p": 0.9, "top_k": 40} + + +def test_temperature_1_kept_for_unsupported_model(monkeypatch): + # temperature=1 is the one value these models still accept. + monkeypatch.setattr(litellm, "drop_params", False) + result = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"temperature": 1}, + model="claude-opus-4-7", + drop_params=True, + ) + + assert result == {"temperature": 1} + + +def test_sampling_param_raises_clean_400_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError, match="does not support temperature"): + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"temperature": 0.3}, + model="claude-opus-4-7", + drop_params=False, + ) From 4e191273508e79f18396a1f307ef776af57d6817 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 21 Aug 2026 18:44:19 -0700 Subject: [PATCH 17/43] fix(pricing): add undated azure aliases for gpt-audio-mini and gpt-realtime-mini (#37867) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(pricing): add undated azure aliases for gpt-audio-mini and gpt-realtime-mini Azure deployments are commonly created against the undated model name, and the cost-tracking docs say to set base_model to azure/ — but only the dated -2025-10-06 entries existed for these two models (the openai provider has undated aliases for both). base_model: azure/gpt-audio-mini therefore resolved to nothing and, depending on the fallback path, text tokens billed at $0 while audio tokens billed fine. Mirror the -2025-10-06 entries as undated aliases, exactly like the undated openai entries mirror their newest dated variant. Fixes #33170 Co-Authored-By: Claude Fable 5 * test(pricing): assert undated azure audio aliases exactly mirror their dated entries Review follow-up: COST_FIELDS missed realtime-specific cost keys (cache_creation_input_audio_token_cost, cache_read_input_token_cost, input_cost_per_image). Full-entry equality catches drift on every field. Co-Authored-By: Claude Fable 5 * fix(pricing): mirror updated mode=realtime on the undated gpt-realtime-mini alias Upstream changed the dated entry's mode from chat to realtime after this branch was cut; the undated alias must stay a byte-for-byte mirror. Co-Authored-By: Claude Fable 5 * fix(pricing): mirror the new deprecation_date onto the undated gpt-audio-mini alias * test(pricing): use shared local_model_cost_map fixture so get_model_info's lru_cache never crosses maps --------- Co-authored-by: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Co-authored-by: Claude Fable 5 --- ...odel_prices_and_context_window_backup.json | 64 ++++++++++++++++ model_prices_and_context_window.json | 64 ++++++++++++++++ .../test_azure_audio_price_aliases.py | 75 +++++++++++++++++++ tests/test_litellm/test_gpt_realtime_mode.py | 1 + 4 files changed, 204 insertions(+) create mode 100644 tests/test_litellm/test_azure_audio_price_aliases.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4c1f777ce3e..3af7d9e5019 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4881,6 +4881,38 @@ "supports_tool_choice": true, "supports_vision": false }, + "azure/gpt-audio-mini": { + "deprecation_date": "2027-04-06", + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "azure/gpt-audio-mini-2025-10-06": { "deprecation_date": "2027-04-06", "input_cost_per_audio_token": 1e-05, @@ -5094,6 +5126,38 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "azure/gpt-realtime-mini": { + "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, + "input_cost_per_audio_token": 1e-05, + "input_cost_per_image": 8e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "azure/gpt-realtime-mini-2025-10-06": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4c1f777ce3e..3af7d9e5019 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4881,6 +4881,38 @@ "supports_tool_choice": true, "supports_vision": false }, + "azure/gpt-audio-mini": { + "deprecation_date": "2027-04-06", + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "azure/gpt-audio-mini-2025-10-06": { "deprecation_date": "2027-04-06", "input_cost_per_audio_token": 1e-05, @@ -5094,6 +5126,38 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "azure/gpt-realtime-mini": { + "cache_creation_input_audio_token_cost": 3e-07, + "cache_read_input_token_cost": 6e-08, + "input_cost_per_audio_token": 1e-05, + "input_cost_per_image": 8e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "azure/gpt-realtime-mini-2025-10-06": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, diff --git a/tests/test_litellm/test_azure_audio_price_aliases.py b/tests/test_litellm/test_azure_audio_price_aliases.py new file mode 100644 index 00000000000..b87744aeae1 --- /dev/null +++ b/tests/test_litellm/test_azure_audio_price_aliases.py @@ -0,0 +1,75 @@ +"""Undated azure aliases for the audio models must exist and match their dated +variants. Azure deployments are commonly created under an admin-chosen name, so +the served model name means nothing to the cost lookup and `base_model: +azure/gpt-audio-mini` is what prices the call. That key resolved to nothing, the +lookup raised "This model isn't mapped yet", and the proxy logged the request at +$0. Issue #33170.""" + +import json +from pathlib import Path + +import pytest + +import litellm + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +COST_FIELDS = ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token", +) + +ALIAS_PAIRS = ( + ("azure/gpt-audio-mini", "azure/gpt-audio-mini-2025-10-06"), + ("azure/gpt-realtime-mini", "azure/gpt-realtime-mini-2025-10-06"), +) + + +def _load_root_cost_map() -> dict: + root_map_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(root_map_path) as f: + return json.load(f) + + +@pytest.mark.parametrize("undated, dated", ALIAS_PAIRS) +def test_undated_azure_audio_alias_matches_dated_entry(undated, dated): + undated_info = litellm.get_model_info(undated) + dated_info = litellm.get_model_info(dated) + + for field in COST_FIELDS: + assert undated_info.get(field) == dated_info.get(field), field + assert (undated_info.get(field) or 0) > 0, f"{undated}.{field} must be non-zero" + + assert undated_info.get("litellm_provider") == "azure" + assert undated_info.get("mode") == dated_info.get("mode") + + +@pytest.mark.parametrize("undated, dated", ALIAS_PAIRS) +def test_undated_azure_audio_alias_is_exact_mirror(undated, dated): + """The undated alias must be a byte-for-byte mirror of its dated entry, covering + every field (incl. realtime-specific cache/audio cost keys) so any future drift + between the pair is caught, not just the core COST_FIELDS.""" + model_map = litellm.model_cost + assert undated in model_map, f"{undated} missing from model cost map" + assert model_map[undated] == model_map[dated], ( + f"{undated} must exactly mirror {dated}; " + f"diff keys: {[k for k in set(model_map[undated]) | set(model_map[dated]) if model_map[undated].get(k) != model_map[dated].get(k)]}" + ) + + +@pytest.mark.parametrize("undated, dated", ALIAS_PAIRS) +def test_undated_azure_audio_alias_is_in_the_root_cost_map(undated, dated): + """`local_model_cost_map` pins `litellm.model_cost` to the packaged backup, but a + proxy left on its defaults fetches the root map instead, and that is the copy + that ships to the CDN. An alias added to only one of the two files still bills + $0 for every proxy reading the other, which is the very bug this file guards, so + assert the root map directly and assert the two files agree.""" + root_map = _load_root_cost_map() + assert undated in root_map, f"{undated} missing from the root cost map" + assert root_map[undated] == root_map[dated], f"{undated} must exactly mirror {dated} in the root cost map" + assert root_map[undated] == litellm.model_cost[undated], ( + f"{undated} differs between the root cost map and the packaged backup" + ) diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py index 4413cbc12ef..ed593228621 100644 --- a/tests/test_litellm/test_gpt_realtime_mode.py +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -10,6 +10,7 @@ from litellm.types.utils import ModelInfoBase REALTIME_ONLY_GPT_MODELS = ( "azure/gpt-realtime-2025-08-28", "azure/gpt-realtime-1.5-2026-02-23", + "azure/gpt-realtime-mini", "azure/gpt-realtime-mini-2025-10-06", "gpt-realtime", "gpt-realtime-1.5", From d4162bd1ca09bd74879560b3a55c3df7f403fc05 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:50:41 -0700 Subject: [PATCH 18/43] test(e2e): record and replay the non-streaming provider flows Chat completions, embeddings, the non-streaming /v1/messages tests, and the OpenAI batch deployment now register through the provider edge, so E2E_FIXTURE_MODE=record captures their provider calls and replay serves them back offline. None of them was wired before, so record was a silent no-op over these suites and replay quietly went live instead of using the bundle Multipart uploads now key on their parsed parts: every ordinary form field, plus the field name, filename, content digest, and length of each file part. The boundary is envelope rather than content, so it stays out of the digest instead of changing the key on every run. A body that does not parse as its declared envelope still has the boundary normalized away before hashing, so the fallback is at least stable, and it records a name that says why Binary uploads hash byte for byte. Canonicalizing them first meant decoding with errors="replace", which collapsed every invalid byte to one U+FFFD and gave two different PDFs of the same length the same key Bundles stay out of the repo: they hold verbatim provider response bodies and expire seven days after recording. Publishing them for CI is LIT-5748, and streaming fidelity is LIT-5742 --- tests/e2e/CLAUDE.md | 15 +- tests/e2e/CONTRIBUTING.md | 8 +- tests/e2e/batches/capabilities.py | 19 +- tests/e2e/batches/test_batches_e2e.py | 9 +- .../test_chat_completions_contract_e2e.py | 9 +- .../test_embeddings_endpoint_e2e.py | 26 ++- .../e2e/llm_translation/test_messages_e2e.py | 38 +-- tests/e2e/provider_edge.py | 216 ++++++++++++++++-- tests/e2e/test_provider_edge.py | 200 ++++++++++++++++ 9 files changed, 475 insertions(+), 65 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 840a40a54cd..05f20ff8b98 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -77,13 +77,24 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover The seam is `provider_edge.py`: `start_provider_edge` boots an in-process HTTP server (one shared instance per pytest process, `e2e_config.provider_edge_base` is the accessor) that mounts each supported provider under a path prefix (`EDGE_MOUNTS`: `/openai` -> `https://api.openai.com`, `/anthropic` -> `https://api.anthropic.com`). A test participates by registering its deployment with `api_base=provider_edge_base("openai")` plus the provider's path suffix; `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` is the reference. In live mode the accessor returns None and the deployment defaults to the real provider, so an edge-wired test runs in all three modes unchanged. Non-wired tests hit their providers live in every mode. The edge binds `E2E_PROVIDER_EDGE_BIND_HOST` (default 127.0.0.1) and advertises `E2E_PROVIDER_EDGE_ADVERTISE_HOST` in the api_base it hands out, for proxies running in containers -A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. `fixture_bundle.py` owns the format. Record serves the proxy the same filtered stored response replay will serve later, so the two modes are byte-identical from the proxy's side of the socket +A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, `multipart/form-data` bodies store their ordinary fields plus a `field:filename` label, a digest of the file part's content, and that part's length, so the per-request random boundary and the envelope never reach the key, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. `fixture_bundle.py` owns the format. Record serves the proxy the same filtered stored response replay will serve later, so the two modes are byte-identical from the proxy's side of the socket Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, edge path, and a content hash, so identity survives re-records and machine changes while any real content drift comes back as an HTTP 599 naming the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a retry loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live providers A replayed response carries the recorded provider response id, and `LiteLLM_SpendLogs.request_id` (the table's primary key) is that id, so a replay against a database that still holds the record run's rows silently dedupes its spend inserts and any spend assertion goes red with zero matching rows and nothing in the proxy log. Run both modes with `E2E_RESET_SPEND_LOGS=1` (plus `DATABASE_URL` in the runner env) so each session truncates the table after itself, or replay against a fresh database, which is the CI shape -Current limits: streaming chunk fidelity is LIT-5742 (a streamed response records as one buffered body), CI wiring is LIT-5748, Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), multipart uploads have per-run random boundaries (the digest changes every run, so they always miss), and deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base) +The same id reuse reaches the managed-object tables. A replayed `/v1/files` or `/v1/batches` response carries the recorded provider object id, and `LiteLLM_ManagedObjectTable.model_object_id` is unique, so a unified batch create replayed against a database that still holds the record run's row fails on a Prisma unique-constraint violation, which surfaces as a 500, makes the router retry, and exhausts the recording. Replay the batches suite against a fresh database, or truncate `LiteLLM_ManagedObjectTable` and `LiteLLM_ManagedFileTable` before the run + +Edge-wired today: `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` (the reference), `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic deployments in `llm_translation/test_messages_e2e.py` except the streaming test, and the OpenAI batch deployment behind `batches/` (`capabilities.openai_batch_params`). The mount base is not the same for both providers: OpenAI deployments register `f"{base}/v1"`, Anthropic deployments register `base` on its own, because litellm's Anthropic handler appends `/v1/messages` to `api_base` itself where the OpenAI handler appends only `/chat/completions`. Recording one suite locally is two runs against a proxy you already have up: + +```bash +E2E_FIXTURE_MODE=record E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 uv run pytest tests/e2e/llm_translation/test_chat_completions_contract_e2e.py +E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 uv run pytest tests/e2e/llm_translation/test_chat_completions_contract_e2e.py +``` + +Point the proxy at bogus provider credentials for the replay run and it still has to pass: that is the whole proof that nothing left the process. Bundles are never committed. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and hard-fails after seven days, and publishing one for CI is LIT-5748 + +Current limits: streaming chunk fidelity is LIT-5742 (a streamed response records as one buffered body), CI wiring is LIT-5748, Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode ## Typing diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 9096050a45a..29778b06d7a 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -57,13 +57,15 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT Record/replay scopes to the proxy's provider-bound traffic only. In `E2E_FIXTURE_MODE=record` the harness boots a local provider-edge server, edge-wired tests register their deployments with an `api_base` pointing at it, and every provider call the proxy makes is forwarded verbatim and written to a fixture bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`). `E2E_FIXTURE_MODE=replay` runs the same tests against the same live proxy and database, but the edge answers the proxy's provider calls from the bundle instead of the provider, so the run makes zero provider calls and spends nothing while key auth, routing, cost calculation, and spend-log writes all still execute for real. Unset (or `live`) behaves exactly as before the knob existed. Both record and replay need the proxy up; only the provider is taken out of the loop ```bash -E2E_FIXTURE_MODE=record uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v -E2E_FIXTURE_MODE=replay uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v +E2E_FIXTURE_MODE=record E2E_FIXTURE_DIR=/tmp/e2e-fixtures uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v +E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v ``` +Bundles stay local. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and expires seven days after it was recorded, so record the suite you want before you replay it and never commit the result; publishing bundles for CI is LIT-5748 + One sharp edge: a replayed response reuses the recorded provider response id, and that id is the primary key of `LiteLLM_SpendLogs`, so replaying against a database that still holds the record run's rows silently dedupes the spend writes and a spend assertion fails with zero rows. Run both commands above with `E2E_RESET_SPEND_LOGS=1` (and `DATABASE_URL` set in the pytest env) so each session truncates the spend log table after itself, or point replay at a fresh database -Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (streaming, Bedrock, multipart) +Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the non-streaming Anthropic tests in `llm_translation/test_messages_e2e.py`, and the OpenAI batch deployment behind `batches/`. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (streaming, Bedrock) Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the proxy isn't up; they never skip for a missing proxy, so an absent proxy can't be mistaken for a pass diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index ce1f68184a7..67eadedbd46 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -7,7 +7,7 @@ import os from dataclasses import dataclass from typing import Literal -from e2e_config import unique_marker +from e2e_config import provider_edge_base, unique_marker from models import LiteLLMParamsBody _BATCH_RUN = unique_marker() @@ -17,6 +17,18 @@ def batch_model_name(base: str) -> str: return f"{base}-{_BATCH_RUN}" +def openai_batch_params() -> LiteLLMParamsBody: + """The OpenAI batch deployment, wired through the record/replay edge when a fixture + mode is active and straight at OpenAI otherwise (LIT-5974). Azure, Vertex, and + Bedrock stay live: none of them has an edge mount.""" + base = provider_edge_base("openai") + return LiteLLMParamsBody( + model="openai/gpt-4o-mini", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + ) + + def _env_ref(*names: str) -> str: for name in names: value = os.environ.get(name) @@ -47,10 +59,7 @@ class Provider: def litellm_params(self) -> LiteLLMParamsBody: match self.name: case "openai": - return LiteLLMParamsBody( - model="openai/gpt-4o-mini", - api_key="os.environ/OPENAI_API_KEY", - ) + return openai_batch_params() case "azure": return LiteLLMParamsBody( model="azure/gpt-5.4-mini-batch", diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 35ad7830e2b..09ef4cfc3a3 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -51,6 +51,7 @@ from capabilities import ( decoded_model_from_id, is_managed_id, matches_id_shape, + openai_batch_params, raw_id_matches_provider, ) from e2e_http import ( @@ -506,13 +507,7 @@ class TestBatchFileContent: self, client: BatchClient, resources: ResourceManager ) -> None: proxy_name = f"e2e-file-content-{unique_marker()}" - model_id = client.create_model( - proxy_name, - LiteLLMParamsBody( - model=f"openai/{OPENAI_FILE_CONTENT_BACKEND}", - api_key="os.environ/OPENAI_API_KEY", - ), - ) + model_id = client.create_model(proxy_name, openai_batch_params()) resources.defer(lambda: client.delete_model(model_id)) key = resources.key() diff --git a/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py b/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py index 2eb7aeb643d..114beaae2fb 100644 --- a/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py @@ -6,7 +6,7 @@ Exercises the gateway against a live OpenAI deployment using customer request sh from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import provider_edge_base, unique_marker from e2e_http import StreamingResponse, assert_client_error, require_successful_call, unwrap from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody @@ -38,10 +38,15 @@ class ChatErrorEnvelope(BaseModel): def _register_chat_model(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: + base = provider_edge_base("openai") model = f"e2e-chat-sec-{unique_marker()}" model_id = proxy.create_model( model, - LiteLLMParamsBody(model=OPENAI_BACKEND, api_key="os.environ/OPENAI_API_KEY"), + LiteLLMParamsBody( + model=OPENAI_BACKEND, + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + ), ) resources.defer(lambda: proxy.delete_model(model_id)) return model, resources.key() diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 35a53f055d8..265cc202ff4 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -9,7 +9,7 @@ covered by tests/e2e/quota_management/spend_tracking/. from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import provider_edge_base, unique_marker from e2e_http import ( assert_client_error, require_successful_call, @@ -27,6 +27,18 @@ class _OptionalEmbeddingsBody(BaseModel): input: str | list[str] | None = None +def _openai_embeddings_params() -> LiteLLMParamsBody: + """The OpenAI embeddings deployment, wired through the record/replay edge when a + fixture mode is active and straight at OpenAI otherwise (LIT-5974). Bedrock and + Vertex stay live: SigV4 signs the Host header, and neither has an edge mount.""" + base = provider_edge_base("openai") + return LiteLLMParamsBody( + model="openai/text-embedding-3-small", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + ) + + class TestEmbeddingsEndpoint: @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_embeddings_returns_vector( @@ -35,9 +47,7 @@ class TestEmbeddingsEndpoint: model = f"e2e-embeddings-{unique_marker()}" model_id = endpoints_client.create_model( model, - LiteLLMParamsBody( - model="openai/text-embedding-3-small", api_key="os.environ/OPENAI_API_KEY" - ), + _openai_embeddings_params(), ) resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() @@ -106,9 +116,7 @@ class TestEmbeddingsEndpoint: model = f"e2e-embeddings-array-{unique_marker()}" model_id = endpoints_client.create_model( model, - LiteLLMParamsBody( - model="openai/text-embedding-3-small", api_key="os.environ/OPENAI_API_KEY" - ), + _openai_embeddings_params(), ) resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() @@ -140,9 +148,7 @@ class TestEmbeddingsEndpoint: model = f"e2e-embeddings-missin-{unique_marker()}" model_id = endpoints_client.create_model( model, - LiteLLMParamsBody( - model="openai/text-embedding-3-small", api_key="os.environ/OPENAI_API_KEY" - ), + _openai_embeddings_params(), ) resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index e0317e0389d..7f81a5e3946 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -9,7 +9,7 @@ litellm-regression-tests/tests/test_inference_endpoints.py. from __future__ import annotations import pytest -from e2e_config import unique_marker +from e2e_config import provider_edge_base, unique_marker from e2e_http import assert_client_error, require_successful_call, unwrap from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager @@ -50,16 +50,27 @@ def _approx_equal(actual: float, expected: float) -> bool: return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) +def _anthropic_params() -> LiteLLMParamsBody: + """The Anthropic deployment, wired through the record/replay edge when a fixture + mode is active (LIT-5974). The mount base carries no ``/v1``: litellm's Anthropic + handler appends ``/v1/messages`` to ``api_base`` itself, where the OpenAI handler + appends only ``/chat/completions``.""" + base = provider_edge_base("anthropic") + return LiteLLMParamsBody( + model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", api_base=base + ) + + class TestAnthropicMessages: def _register( - self, endpoints_client: EndpointsClient, resources: ResourceManager + self, + endpoints_client: EndpointsClient, + resources: ResourceManager, + params: LiteLLMParamsBody | None = None, ) -> tuple[str, str]: model = f"e2e-messages-{unique_marker()}" model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" - ), + model, _anthropic_params() if params is None else params ) resources.defer(lambda: endpoints_client.delete_model(model_id)) return model, resources.key() @@ -81,12 +92,7 @@ class TestAnthropicMessages: self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: model = f"e2e-messages-cost-{unique_marker()}" - model_id = endpoints_client.create_model( - model, - LiteLLMParamsBody( - model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY" - ), - ) + model_id = endpoints_client.create_model(model, _anthropic_params()) resources.defer(lambda: endpoints_client.delete_model(model_id)) key = resources.key() @@ -131,7 +137,13 @@ class TestAnthropicMessages: def test_messages_streams_completion( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - model, key = self._register(endpoints_client, resources) + """Stays on a live Anthropic deployment in every mode: the edge buffers a + streamed response into one body, so chunk fidelity waits on LIT-5742.""" + model, key = self._register( + endpoints_client, + resources, + LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), + ) result = endpoints_client.proxy.messages_stream( key, diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index ab0791e6b74..92ffe75e800 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -20,10 +20,9 @@ headers must never touch disk. An unmatched replay call returns HTTP proxy relays as a provider error the failing test surfaces. v1 limits: only the mounts in ``EDGE_MOUNTS`` (SigV4 providers like Bedrock -sign the Host header, so a forwarding edge breaks their signatures), JSON and -opaque single-part bodies (multipart boundaries are random per request), -streaming fidelity is LIT-5742, and CI wiring is LIT-5748. Suites that do not -wire the edge keep hitting providers live in every mode. +sign the Host header, so a forwarding edge breaks their signatures), streaming +fidelity is LIT-5742, and CI wiring is LIT-5748. Suites that do not wire the +edge keep hitting providers live in every mode. """ from __future__ import annotations @@ -32,6 +31,7 @@ import base64 import difflib import functools import hashlib +import re import threading from collections import deque from collections.abc import Mapping @@ -103,26 +103,194 @@ _RESPONSE_DROPPED_HEADERS: Final[frozenset[str]] = _HOP_BY_HOP_HEADERS | { _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) -def _edge_request(method: str, path: str, query: str, body: bytes | None) -> RecordedRequest: - """The identity replay matches on: the edge path (mount included), the query - as params, and the body as parsed JSON, or as a canonicalized content digest - when it is not JSON so opaque uploads still match across runs.""" - params: Final = dict(parse_qsl(query, keep_blank_values=True)) - if not body: - return RecordedRequest(method=method.lower(), path=path, headers={}, params=params) - decoded: Final = body.decode("utf-8", errors="replace") +_BOUNDARY_PATTERN: Final = re.compile( + r'boundary=(?:"([^"]*)"|([^;,\s]+))', re.IGNORECASE +) +_DISPOSITION_NAME_PATTERN: Final = re.compile(r'(?:^|;)\s*name="([^"]*)"') +_DISPOSITION_FILENAME_PATTERN: Final = re.compile(r'(?:^|;)\s*filename="([^"]*)"') +_UNPARSED_MULTIPART: Final = "" +_BOUNDARY_PLACEHOLDER: Final = b"--" + + +@dataclass(frozen=True) +class _MultipartPart: + field_name: str + filename: str | None + content: bytes + + +def _header_value(headers: Mapping[str, str], name: str) -> str: + wanted: Final = name.lower() + return next((value for key, value in headers.items() if key.lower() == wanted), "") + + +def _multipart_boundary(content_type: str) -> str | None: + if "multipart/form-data" not in content_type.lower(): + return None + match: Final = _BOUNDARY_PATTERN.search(content_type) + return None if match is None else match.group(1) or match.group(2) + + +def _parse_multipart_part(segment: bytes) -> _MultipartPart | None: + head, separator, content = segment.partition(b"\r\n\r\n") + if not separator: + return None + disposition: Final = "".join( + value + for line in head.decode("utf-8", errors="replace").split("\r\n") + for name, _, value in [line.partition(":")] + if name.strip().lower() == "content-disposition" + ) + name_match: Final = _DISPOSITION_NAME_PATTERN.search(disposition) + if name_match is None: + return None + filename_match: Final = _DISPOSITION_FILENAME_PATTERN.search(disposition) + return _MultipartPart( + field_name=name_match.group(1), + filename=None if filename_match is None else filename_match.group(1), + content=content, + ) + + +def _multipart_parts(body: bytes, boundary: str) -> tuple[_MultipartPart, ...] | None: + """The wire body split back into its parts, or None when it does not parse as the + declared envelope so the caller can fall back to the opaque content digest.""" + segments: Final = body.split(b"--" + boundary.encode()) + if len(segments) < 3 or not segments[-1].startswith(b"--"): + return None + parsed: Final = tuple( + _parse_multipart_part(segment.removeprefix(b"\r\n").removesuffix(b"\r\n")) + for segment in segments[1:-1] + ) + if any(part is None for part in parsed): + return None + return tuple(part for part in parsed if part is not None) + + +def _content_digest(content: bytes) -> str: + """Text is canonicalized before hashing so a per-run marker inside an uploaded JSONL + does not move the key; anything that is not UTF-8 is hashed byte for byte, since a + lossy decode collapses every binary payload of one length onto one digest.""" try: - parsed: Final[JsonValue] = _JSON.validate_json(decoded) - except ValueError: - return RecordedRequest( - method=method.lower(), - path=path, - headers={}, - params=params, - file_sha256=hashlib.sha256(canonical_string(decoded).encode()).hexdigest(), - file_bytes=len(body), + text: Final = content.decode("utf-8") + except UnicodeDecodeError: + return hashlib.sha256(content).hexdigest() + return hashlib.sha256(canonical_string(text).encode()).hexdigest() + + +def _form_fields(fields: tuple[_MultipartPart, ...]) -> dict[str, str]: + """The ordinary field parts, flattened into the mapping the bundle format stores. A + name sent more than once takes an index instead of overwriting the earlier value, so + nothing an upload said is dropped from its key.""" + form: dict[str, str] = {} + for part in fields: + name = part.field_name + occurrence = 1 + while name in form: + name = f"{part.field_name}[{occurrence}]" + occurrence += 1 + form[name] = part.content.decode("utf-8", errors="replace") + return form + + +def _file_identity(files: tuple[_MultipartPart, ...]) -> tuple[str | None, str | None, int | None]: + """Name, content digest, and total length for the uploaded file parts. The name + carries each part's field name as well as its filename, so two uploads sending the + same bytes under different field names stay apart. A lone file keeps its own content + digest; several fold into one digest over the per-part identities, which is ordered, + so parts arriving in a different order key differently.""" + if not files: + return None, None, None + names: Final = ", ".join(f"{part.field_name}:{part.filename}" for part in files) + total: Final = sum(len(part.content) for part in files) + if len(files) == 1: + return names, _content_digest(files[0].content), total + folded: Final = _JSON.dump_json( + [ + [part.field_name, part.filename, _content_digest(part.content), len(part.content)] + for part in files + ] + ) + return names, hashlib.sha256(folded).hexdigest(), total + + +def _multipart_request( + method: str, path: str, params: dict[str, str], parts: tuple[_MultipartPart, ...] +) -> RecordedRequest: + """A multipart upload keyed by what it says rather than by its wire bytes: every + ordinary field, plus the identity of the uploaded file. The random per-request + boundary is envelope, never content, so it never reaches the digest.""" + form: Final = _form_fields(tuple(part for part in parts if part.filename is None)) + file_name, file_sha256, file_bytes = _file_identity( + tuple(part for part in parts if part.filename is not None) + ) + return RecordedRequest( + method=method, + path=path, + headers={}, + params=params, + form=form, + file_name=file_name, + file_sha256=file_sha256, + file_bytes=file_bytes, + ) + + +def _opaque_request( + method: str, + path: str, + params: dict[str, str], + body: bytes, + digested: bytes, + file_name: str | None = None, +) -> RecordedRequest: + """A body kept out of the bundle and matched on its digest alone. ``digested`` is + what the digest runs over, which is the body itself unless something in it has to be + normalized away first.""" + return RecordedRequest( + method=method, + path=path, + headers={}, + params=params, + file_name=file_name, + file_sha256=_content_digest(digested), + file_bytes=len(body), + ) + + +def _edge_request( + method: str, path: str, query: str, body: bytes | None, content_type: str = "" +) -> RecordedRequest: + """The identity replay matches on: the edge path (mount included), the query as + params, and the body as parsed JSON, as parsed multipart fields and file identity + when the content type declares an envelope, or as a content digest otherwise so + opaque uploads still match across runs. A multipart body that does not parse still + has its boundary normalized away, because that boundary is fresh every request and + would otherwise guarantee a miss.""" + params: Final = dict(parse_qsl(query, keep_blank_values=True)) + lowered_method: Final = method.lower() + if not body: + return RecordedRequest(method=lowered_method, path=path, headers={}, params=params) + boundary: Final = _multipart_boundary(content_type) + if boundary is not None: + parts = _multipart_parts(body, boundary) + if parts is not None: + return _multipart_request(lowered_method, path, params, parts) + return _opaque_request( + lowered_method, + path, + params, + body, + body.replace(b"--" + boundary.encode(), _BOUNDARY_PLACEHOLDER), + _UNPARSED_MULTIPART, ) - return RecordedRequest(method=method.lower(), path=path, headers={}, params=params, body=parsed) + try: + parsed: Final[JsonValue] = _JSON.validate_json(body) + except ValueError: + return _opaque_request(lowered_method, path, params, body, body) + return RecordedRequest( + method=lowered_method, path=path, headers={}, params=params, body=parsed + ) def _build_pool(recorded: tuple[Interaction, ...]) -> dict[str, deque[Interaction]]: @@ -351,7 +519,9 @@ def handle_edge_request( return _text_reply( 404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}" ) - request: Final = _edge_request(method, split.path, split.query, body) + request: Final = _edge_request( + method, split.path, split.query, body, _header_value(headers, "content-type") + ) match backend: case RecordEdge(): return _handle_record( diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 492eee57aaf..f237ea70ff8 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -53,8 +53,10 @@ from provider_edge import ( ) CHAT_PATH = "/openai/v1/chat/completions" +UPLOAD_PATH = "/openai/v1/files" REPLAY_MOUNTS = {"openai": "https://replay.invalid"} JSON_OBJECT = TypeAdapter(dict[str, object]) +BATCH_JSONL = b'{"custom_id":"one"}\n{"custom_id":"two"}\n' def json_object(body: bytes) -> dict[str, object]: @@ -164,6 +166,46 @@ def chat_body(prompt: str) -> bytes: return json.dumps({"model": "gpt", "messages": [{"role": "user", "content": prompt}]}).encode() +def multipart_body( + boundary: str, + fields: tuple[tuple[str, str], ...] = (), + files: tuple[tuple[str, str, bytes], ...] = (), +) -> bytes: + """One multipart/form-data body on the wire, exactly as ``requests`` writes it, with + the boundary under the caller's control instead of randomly generated.""" + parts = [ + f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n'.encode() + + value.encode() + for name, value in fields + ] + [ + ( + f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"; ' + f'filename="{filename}"\r\nContent-Type: application/octet-stream\r\n\r\n' + ).encode() + + content + for name, filename, content in files + ] + return b"\r\n".join(parts) + f"\r\n--{boundary}--\r\n".encode() + + +def upload_headers(boundary: str) -> dict[str, str]: + return { + "content-type": f"multipart/form-data; boundary={boundary}", + "authorization": "Bearer sk-upload-secret", + } + + +def record_upload(root: Path, body: bytes, boundary: str) -> None: + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + call_edge(edge, "POST", UPLOAD_PATH, body=body, headers=upload_headers(boundary)) + + +def replay_upload(root: Path, body: bytes, boundary: str) -> RawResponse: + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + return call_edge(edge, "POST", UPLOAD_PATH, body=body, headers=upload_headers(boundary)) + + class TestRecordMode: def test_forwards_to_the_provider_and_writes_one_interaction_file(self, tmp_path: Path) -> None: root = tmp_path / "bundle" @@ -328,6 +370,164 @@ class TestReplayMode: assert replayed.status_code == 200 +class TestMultipartIdentity: + """LIT-5974: a multipart upload is keyed by its parsed fields and file identity. + ``requests`` picks a fresh random boundary per request, so hashing the wire body + made every upload miss on replay; parsing the envelope keys the upload on what it + actually says, which is stable across runs and still separates real drift.""" + + def test_a_fresh_boundary_replays_the_same_upload(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + recorded = multipart_body( + "d0a1b2c3d4e5f60718293a4b5c6d7e8f", + fields=(("purpose", "batch"),), + files=(("file", "batch.jsonl", BATCH_JSONL),), + ) + record_upload(root, recorded, "d0a1b2c3d4e5f60718293a4b5c6d7e8f") + + rerun = multipart_body( + "ffffeeeeddddccccbbbbaaaa99998888", + fields=(("purpose", "batch"),), + files=(("file", "batch.jsonl", BATCH_JSONL),), + ) + assert rerun != recorded + replayed = replay_upload(root, rerun, "ffffeeeeddddccccbbbbaaaa99998888") + assert replayed.status_code == 200, replayed.body[:400] + + def test_the_stored_request_carries_fields_and_file_identity_but_no_secrets( + self, tmp_path: Path + ) -> None: + root = tmp_path / "bundle" + boundary = "0123456789abcdef0123456789abcdef" + record_upload( + root, + multipart_body( + boundary, + fields=(("purpose", "batch"),), + files=(("file", "batch.jsonl", BATCH_JSONL),), + ), + boundary, + ) + + raw = this_tests_files(root)[0].read_text(encoding="utf-8") + interaction = Interaction.model_validate_json(raw) + assert interaction.request.form == {"purpose": "batch"} + assert interaction.request.file_name == "file:batch.jsonl" + assert interaction.request.file_bytes == len(BATCH_JSONL) + stored = interaction.request.model_dump_json() + assert boundary not in stored + assert "sk-upload-secret" not in stored + assert "custom_id" not in stored + + @pytest.mark.parametrize( + ("fields", "files"), + [ + pytest.param( + (("purpose", "batch"),), + (("file", "batch.jsonl", b'{"custom_id":"three"}\n'),), + id="file-content", + ), + pytest.param( + (("purpose", "batch"),), + (("file", "other.jsonl", BATCH_JSONL),), + id="file-name", + ), + pytest.param( + (("purpose", "fine-tune"),), + (("file", "batch.jsonl", BATCH_JSONL),), + id="form-field", + ), + pytest.param( + (("purpose", "batch"), ("purpose", "batch")), + (("file", "batch.jsonl", BATCH_JSONL),), + id="repeated-form-field", + ), + pytest.param( + (("purpose", "batch"),), + ( + ("file", "batch.jsonl", BATCH_JSONL), + ("mask", "mask.jsonl", BATCH_JSONL), + ), + id="extra-file-part", + ), + ], + ) + def test_a_structurally_different_upload_misses( + self, + tmp_path: Path, + fields: tuple[tuple[str, str], ...], + files: tuple[tuple[str, str, bytes], ...], + ) -> None: + root = tmp_path / "bundle" + record_upload( + root, + multipart_body( + "aaaaaaaabbbbbbbbccccccccdddddddd", + fields=(("purpose", "batch"),), + files=(("file", "batch.jsonl", BATCH_JSONL),), + ), + "aaaaaaaabbbbbbbbccccccccdddddddd", + ) + + drifted = replay_upload( + root, + multipart_body("11112222333344445555666677778888", fields=fields, files=files), + "11112222333344445555666677778888", + ) + assert drifted.status_code == REPLAY_MISS_STATUS + + def test_several_file_parts_separate_when_their_contents_swap(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + image, mask = b"image-bytes", b"mask-bytes" + record_upload( + root, + multipart_body( + "1a1a1a1a2b2b2b2b3c3c3c3c4d4d4d4d", + fields=(("prompt", "a cat"),), + files=(("image", "a.png", image), ("mask", "b.png", mask)), + ), + "1a1a1a1a2b2b2b2b3c3c3c3c4d4d4d4d", + ) + + swapped = replay_upload( + root, + multipart_body( + "5e5e5e5e6f6f6f6f7070707081818181", + fields=(("prompt", "a cat"),), + files=(("image", "a.png", mask), ("mask", "b.png", image)), + ), + "5e5e5e5e6f6f6f6f7070707081818181", + ) + assert swapped.status_code == REPLAY_MISS_STATUS + + same = replay_upload( + root, + multipart_body( + "9292929203030303a4a4a4a4b5b5b5b5", + fields=(("prompt", "a cat"),), + files=(("image", "a.png", image), ("mask", "b.png", mask)), + ), + "9292929203030303a4a4a4a4b5b5b5b5", + ) + assert same.status_code == 200, same.body[:400] + + def test_a_body_that_does_not_match_its_declared_boundary_stays_opaque( + self, tmp_path: Path + ) -> None: + root = tmp_path / "bundle" + opaque = b"custom_id one\ncustom_id two\n" + absent = "boundary-that-is-absent-from-the-body" + record_upload(root, opaque, absent) + + raw = this_tests_files(root)[0].read_text(encoding="utf-8") + interaction = Interaction.model_validate_json(raw) + assert interaction.request.form is None + assert interaction.request.file_name == "" + assert interaction.request.file_bytes == len(opaque) + assert "custom_id" not in interaction.request.model_dump_json() + assert replay_upload(root, opaque, absent).status_code == 200 + + class TestReplayLeftover: def test_partially_consumed_recording_names_the_leftover(self, tmp_path: Path) -> None: root = tmp_path / "bundle" From f5df60f1062c19758fcea48c8a88354347a413ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:32:02 -0700 Subject: [PATCH 19/43] test(e2e): key multipart uploads by structured part identity Adversarial review of the new multipart keying turned up collisions where two different provider requests computed the same replay key, which is the dangerous failure for a replay harness: the second request silently gets the first one's response instead of missing loudly. - a part counts as an upload when it has a filename or declares its own content type, and the declared content type joins the identity, so two uploads of the same bytes under the same field no longer collapse - the uploaded parts contribute a JSON list of [field, filename, type] triples instead of a "field:filename" string, so a separator inside a filename can no longer impersonate a field boundary - repeated field names get a "name[n]" suffix with a literal "[" doubled first, so a repeated field and a literally indexed one stay distinct - a field value that is not UTF-8 is stored as a base64 sha256 digest; base64 rather than hex because the canonicalizer rewrites 64-character hex runs to and folded every binary value onto one key - a field whose name reads as a credential is stored as . This stays key-preserving because the key is recomputed from the stored request rather than saved beside it, so the live request carrying the real value still matches its redacted fixture - the uploaded byte length leaves the key. The canonicalizer absorbs timestamp and id drift inside a file, and that drift moves the count, so keeping it there made re-records miss Also stops a lookalike parameter such as "xboundary=" from being read as the multipart boundary, and gives the OpenAI batch backend model a single constant instead of three copies of the literal. BUNDLE_FORMAT_VERSION goes to 3 because all of this moves recorded keys. A bundle recorded under the old rules now fails naming both versions instead of missing on every call. --- tests/e2e/CLAUDE.md | 4 +- tests/e2e/batches/capabilities.py | 13 +- tests/e2e/batches/test_batches_e2e.py | 5 +- tests/e2e/fixture_bundle.py | 15 +- tests/e2e/fixture_canonical.py | 1 - tests/e2e/provider_edge.py | 120 +++++++++++----- tests/e2e/test_provider_edge.py | 188 +++++++++++++++++++++++++- 7 files changed, 302 insertions(+), 44 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 05f20ff8b98..15bd2c19ca9 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -77,7 +77,9 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover The seam is `provider_edge.py`: `start_provider_edge` boots an in-process HTTP server (one shared instance per pytest process, `e2e_config.provider_edge_base` is the accessor) that mounts each supported provider under a path prefix (`EDGE_MOUNTS`: `/openai` -> `https://api.openai.com`, `/anthropic` -> `https://api.anthropic.com`). A test participates by registering its deployment with `api_base=provider_edge_base("openai")` plus the provider's path suffix; `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` is the reference. In live mode the accessor returns None and the deployment defaults to the real provider, so an edge-wired test runs in all three modes unchanged. Non-wired tests hit their providers live in every mode. The edge binds `E2E_PROVIDER_EDGE_BIND_HOST` (default 127.0.0.1) and advertises `E2E_PROVIDER_EDGE_ADVERTISE_HOST` in the api_base it hands out, for proxies running in containers -A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, `multipart/form-data` bodies store their ordinary fields plus a `field:filename` label, a digest of the file part's content, and that part's length, so the per-request random boundary and the envelope never reach the key, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. `fixture_bundle.py` owns the format. Record serves the proxy the same filtered stored response replay will serve later, so the two modes are byte-identical from the proxy's side of the socket +A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, `multipart/form-data` bodies store their ordinary fields plus a JSON list of the uploaded parts' `[field, filename, content-type]` triples and a digest of their content, so the per-request random boundary and the envelope never reach the key, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. `fixture_bundle.py` owns the format. Record serves the proxy the same filtered stored response replay will serve later, so the two modes are byte-identical from the proxy's side of the socket + +Multipart identity is the fiddly corner, and the rules exist because each one had a collision behind it. A part counts as an upload when it carries a filename or declares its own content type, and everything else is an ordinary field. Field names get a `name[n]` suffix on repeats, with a literal `[` doubled first, so a form that repeats `purpose` never keys the same as one that literally sends `purpose[1]`. A field whose name reads as a credential is stored as ``, which stays key-preserving because the key is recomputed from the stored request rather than saved alongside it, so the live request carrying the real value still matches its redacted fixture. A field value that is not UTF-8 is stored as a base64 sha256 digest, base64 and not hex because the canonicalizer rewrites any 64-character hex run to `` and would fold every binary value onto one key. The uploaded parts contribute a JSON list rather than a `field:filename` string, so a separator inside a filename cannot impersonate a field boundary, and their byte length is stored for a reader's benefit but deliberately left out of the key, since the canonicalizer absorbs timestamp and id drift inside a file that changes its length Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, edge path, and a content hash, so identity survives re-records and machine changes while any real content drift comes back as an HTTP 599 naming the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a retry loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live providers diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 67eadedbd46..ee44a50d215 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -5,7 +5,7 @@ from __future__ import annotations import base64 import os from dataclasses import dataclass -from typing import Literal +from typing import Final, Literal from e2e_config import provider_edge_base, unique_marker from models import LiteLLMParamsBody @@ -17,13 +17,16 @@ def batch_model_name(base: str) -> str: return f"{base}-{_BATCH_RUN}" +OPENAI_BATCH_BACKEND: Final = "gpt-4o-mini" + + def openai_batch_params() -> LiteLLMParamsBody: """The OpenAI batch deployment, wired through the record/replay edge when a fixture mode is active and straight at OpenAI otherwise (LIT-5974). Azure, Vertex, and Bedrock stay live: none of them has an edge mount.""" base = provider_edge_base("openai") return LiteLLMParamsBody( - model="openai/gpt-4o-mini", + model=f"openai/{OPENAI_BATCH_BACKEND}", api_key="os.environ/OPENAI_API_KEY", api_base=None if base is None else f"{base}/v1", ) @@ -116,7 +119,11 @@ class Capability: PROVIDERS: tuple[Provider, ...] = ( Provider( - "openai", batch_model_name("openai-batch"), "gpt-4o-mini", can_cancel=True, can_list=True + "openai", + batch_model_name("openai-batch"), + OPENAI_BATCH_BACKEND, + can_cancel=True, + can_list=True, ), Provider( "azure", diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 09ef4cfc3a3..7af064b1fdd 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -42,6 +42,7 @@ from capabilities import ( BATCH_ID_SHAPE, CAPABILITIES, FILE_ID_SHAPE, + OPENAI_BATCH_BACKEND, OPENAI_BATCH_MODEL, PROVIDERS, Capability, @@ -480,8 +481,6 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( ) -OPENAI_FILE_CONTENT_BACKEND = "gpt-4o-mini" - FILE_CONTENT_CELLS = { "azure": "llm.files.azure_openai.content.nonstream.works", "vertex_ai": "llm.files.vertex.content.nonstream.works", @@ -511,7 +510,7 @@ class TestBatchFileContent: resources.defer(lambda: client.delete_model(model_id)) key = resources.key() - payload = render_jsonl(OPENAI_FILE_CONTENT_BACKEND) + payload = render_jsonl(OPENAI_BATCH_BACKEND) file = unwrap( client.upload_file( content=payload, diff --git a/tests/e2e/fixture_bundle.py b/tests/e2e/fixture_bundle.py index 6feb40fc8bc..aa0ba100b6c 100644 --- a/tests/e2e/fixture_bundle.py +++ b/tests/e2e/fixture_bundle.py @@ -5,7 +5,9 @@ version + format version) plus one subdirectory per test, holding one JSON file per provider-bound interaction in call order. Bundles older than ``MAX_BUNDLE_AGE`` hard-fail replay at collection time (see conftest), so a green replay run can never certify against fixtures that have drifted more than -a week from the live providers. +a week from the live providers. Bump ``BUNDLE_FORMAT_VERSION`` whenever a change +moves recorded keys: a bundle recorded under the old rules then fails naming +both versions instead of quietly missing on every call. This module owns the format only. The provider-edge server that produces and consumes it lives in provider_edge.py (LIT-5745) and the canonical match keys @@ -28,7 +30,7 @@ from typing import Final from pydantic import BaseModel, JsonValue -BUNDLE_FORMAT_VERSION: Final = 2 +BUNDLE_FORMAT_VERSION: Final = 3 MAX_BUNDLE_AGE: Final = timedelta(days=7) MANIFEST_FILENAME: Final = "manifest.json" @@ -47,7 +49,14 @@ class RecordedRequest(BaseModel): over ``method``, ``path`` (the edge path including the provider mount, query string excluded), and the canonicalized headers, params, body, form, and file identity. Non-JSON bodies store a canonicalized content digest - instead of the bytes.""" + instead of the bytes. + + ``file_name`` is a JSON list of the uploaded parts' ``[field, filename, + content-type]`` triples rather than a flat label, so a separator inside a + filename cannot impersonate a field boundary. ``file_bytes`` is recorded for + a reader's benefit and stays out of the key: the canonicalizer absorbs + timestamp and id drift inside an uploaded file, and that drift moves the + byte count.""" method: str path: str diff --git a/tests/e2e/fixture_canonical.py b/tests/e2e/fixture_canonical.py index 427f06bf8fb..c043951a108 100644 --- a/tests/e2e/fixture_canonical.py +++ b/tests/e2e/fixture_canonical.py @@ -129,7 +129,6 @@ def canonicalize(request: RecordedRequest) -> CanonicalRequest: else { "name": None if request.file_name is None else canonical_string(request.file_name), "sha256": request.file_sha256, - "bytes": request.file_bytes, } ) content: Final[dict[str, JsonValue]] = { diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 92ffe75e800..25a1e8043ed 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -59,7 +59,13 @@ from fixture_bundle import ( prepare_bundle, slug_for_test, ) -from fixture_canonical import CanonicalRequest, canonical_string, canonicalize +from fixture_canonical import ( + SECRET_PLACEHOLDER, + CanonicalRequest, + canonical_string, + canonicalize, + is_secret_field, +) from fixture_mode import ( FIXTURE_MODES, InvalidFixtureMode, @@ -104,12 +110,15 @@ _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _BOUNDARY_PATTERN: Final = re.compile( - r'boundary=(?:"([^"]*)"|([^;,\s]+))', re.IGNORECASE + r'(?:^|;)\s*boundary\s*=\s*(?:"([^"]*)"|([^;,\s]+))', re.IGNORECASE +) +_DISPOSITION_NAME_PATTERN: Final = re.compile(r'(?:^|;)\s*name="([^"]*)"', re.IGNORECASE) +_DISPOSITION_FILENAME_PATTERN: Final = re.compile( + r'(?:^|;)\s*filename="([^"]*)"', re.IGNORECASE ) -_DISPOSITION_NAME_PATTERN: Final = re.compile(r'(?:^|;)\s*name="([^"]*)"') -_DISPOSITION_FILENAME_PATTERN: Final = re.compile(r'(?:^|;)\s*filename="([^"]*)"') _UNPARSED_MULTIPART: Final = "" _BOUNDARY_PLACEHOLDER: Final = b"--" +_BINARY_FIELD_PREFIX: Final = " str: @@ -125,22 +135,34 @@ def _header_value(headers: Mapping[str, str], name: str) -> str: def _multipart_boundary(content_type: str) -> str | None: + """The declared boundary, or None when the envelope is not multipart or names no + usable boundary. ``boundary`` is matched only as a parameter in its own right, so a + longer name ending in it (``myboundary=``) is not mistaken for one, and an empty + boundary is refused rather than splitting the body on a bare ``--``.""" if "multipart/form-data" not in content_type.lower(): return None match: Final = _BOUNDARY_PATTERN.search(content_type) - return None if match is None else match.group(1) or match.group(2) + if match is None: + return None + quoted, bare = match.group(1), match.group(2) + return (quoted if quoted is not None else bare) or None + + +def _part_headers(head: bytes) -> dict[str, str]: + return { + name.strip().lower(): value.strip() + for line in head.decode("utf-8", errors="replace").split("\r\n") + for name, separator, value in [line.partition(":")] + if separator + } def _parse_multipart_part(segment: bytes) -> _MultipartPart | None: head, separator, content = segment.partition(b"\r\n\r\n") if not separator: return None - disposition: Final = "".join( - value - for line in head.decode("utf-8", errors="replace").split("\r\n") - for name, _, value in [line.partition(":")] - if name.strip().lower() == "content-disposition" - ) + headers: Final = _part_headers(head) + disposition: Final = headers.get("content-disposition", "") name_match: Final = _DISPOSITION_NAME_PATTERN.search(disposition) if name_match is None: return None @@ -149,6 +171,7 @@ def _parse_multipart_part(segment: bytes) -> _MultipartPart | None: field_name=name_match.group(1), filename=None if filename_match is None else filename_match.group(1), content=content, + content_type=headers.get("content-type", ""), ) @@ -178,39 +201,72 @@ def _content_digest(content: bytes) -> str: return hashlib.sha256(canonical_string(text).encode()).hexdigest() +def _is_file_part(part: _MultipartPart) -> bool: + """Whether a part is an upload rather than an ordinary field. A filename says so + outright, and so does a declared content type: clients attach one per part only for + a file, and a client that omits the filename (httpx drops the parameter when it is + empty) would otherwise have the file's bytes stored inline as a field value and key + identically to a plain field of the same name.""" + return part.filename is not None or bool(part.content_type) + + +def _field_value(part: _MultipartPart) -> str: + """What a field part contributes to the stored form. A secret-named field never has + its value written out, since the bundle is a file on disk and the key redacts that + field to the same placeholder either way, so replay still matches. A value that is + not UTF-8 is carried as a digest rather than decoded lossily, because a replacing + decode collapses every binary value of one length onto one string. That digest is + base64 rather than hex, since the canonicalizer rewrites any long hex run to a + ```` placeholder and would collapse the values right back together.""" + if is_secret_field(part.field_name): + return SECRET_PLACEHOLDER + try: + return part.content.decode("utf-8") + except UnicodeDecodeError: + digest: Final = base64.b64encode(hashlib.sha256(part.content).digest()).decode() + return f"{_BINARY_FIELD_PREFIX}{digest}>" + + def _form_fields(fields: tuple[_MultipartPart, ...]) -> dict[str, str]: """The ordinary field parts, flattened into the mapping the bundle format stores. A - name sent more than once takes an index instead of overwriting the earlier value, so - nothing an upload said is dropped from its key.""" + name sent more than once takes an occurrence suffix instead of overwriting the + earlier value, so nothing an upload said is dropped from its key. The suffix is + escaped so a field literally named ``x[1]`` cannot collide with a second ``x``.""" form: dict[str, str] = {} for part in fields: - name = part.field_name + name = part.field_name.replace("[", "[[") occurrence = 1 while name in form: - name = f"{part.field_name}[{occurrence}]" + name = f"{part.field_name.replace('[', '[[')}[{occurrence}]" occurrence += 1 - form[name] = part.content.decode("utf-8", errors="replace") + form[name] = _field_value(part) return form def _file_identity(files: tuple[_MultipartPart, ...]) -> tuple[str | None, str | None, int | None]: - """Name, content digest, and total length for the uploaded file parts. The name - carries each part's field name as well as its filename, so two uploads sending the - same bytes under different field names stay apart. A lone file keeps its own content - digest; several fold into one digest over the per-part identities, which is ordered, - so parts arriving in a different order key differently.""" + """Name, content digest, and total length for the uploaded file parts. + + The name is a structured list of every part's field name, filename, and declared + content type rather than a joined string, so a filename containing the separator + cannot be confused for a different split, and two parts that differ only in the type + they declare stay apart. It goes through the canonicalizer as one string, which is + why per-run markers inside a filename do not move the key in the multi-file case any + more than they do in the single-file one. + + The digest covers content only. A lone file keeps its own canonicalized digest; + several fold into one ordered digest, so parts arriving in a different order key + differently. Total length is recorded for a reader but deliberately kept out of the + key: it is the raw byte count, and keying on it would undo exactly the drift the + canonicalized digest exists to absorb.""" if not files: return None, None, None - names: Final = ", ".join(f"{part.field_name}:{part.filename}" for part in files) + names: Final = _JSON.dump_json( + [[part.field_name, part.filename, part.content_type] for part in files] + ).decode() total: Final = sum(len(part.content) for part in files) if len(files) == 1: return names, _content_digest(files[0].content), total - folded: Final = _JSON.dump_json( - [ - [part.field_name, part.filename, _content_digest(part.content), len(part.content)] - for part in files - ] - ) + folded: Final = _JSON.dump_json([_content_digest(part.content) for part in files]) return names, hashlib.sha256(folded).hexdigest(), total @@ -220,9 +276,9 @@ def _multipart_request( """A multipart upload keyed by what it says rather than by its wire bytes: every ordinary field, plus the identity of the uploaded file. The random per-request boundary is envelope, never content, so it never reaches the digest.""" - form: Final = _form_fields(tuple(part for part in parts if part.filename is None)) + form: Final = _form_fields(tuple(part for part in parts if not _is_file_part(part))) file_name, file_sha256, file_bytes = _file_identity( - tuple(part for part in parts if part.filename is not None) + tuple(part for part in parts if _is_file_part(part)) ) return RecordedRequest( method=method, @@ -258,7 +314,7 @@ def _opaque_request( ) -def _edge_request( +def edge_request( method: str, path: str, query: str, body: bytes | None, content_type: str = "" ) -> RecordedRequest: """The identity replay matches on: the edge path (mount included), the query as @@ -519,7 +575,7 @@ def handle_edge_request( return _text_reply( 404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}" ) - request: Final = _edge_request( + request: Final = edge_request( method, split.path, split.query, body, _header_value(headers, "content-type") ) match backend: diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index f237ea70ff8..14a9fd53393 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -23,11 +23,13 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +from typing import Final import pytest from pydantic import TypeAdapter from e2e_http import RawResponse, forward +from fixture_canonical import canonicalize from fixture_bundle import ( BundleRecorder, Interaction, @@ -46,6 +48,7 @@ from provider_edge import ( RecordEdge, ReplayEdge, ReplaySource, + edge_request, handle_edge_request, provider_edge_api_base, replay_leftover_error, @@ -412,7 +415,9 @@ class TestMultipartIdentity: raw = this_tests_files(root)[0].read_text(encoding="utf-8") interaction = Interaction.model_validate_json(raw) assert interaction.request.form == {"purpose": "batch"} - assert interaction.request.file_name == "file:batch.jsonl" + assert interaction.request.file_name == json.dumps( + [["file", "batch.jsonl", "application/octet-stream"]], separators=(",", ":") + ) assert interaction.request.file_bytes == len(BATCH_JSONL) stored = interaction.request.model_dump_json() assert boundary not in stored @@ -528,6 +533,187 @@ class TestMultipartIdentity: assert replay_upload(root, opaque, absent).status_code == 200 +def raw_multipart(boundary: str, *parts: tuple[str, bytes]) -> bytes: + """A body assembled from literal part headers, so a test can send the shapes a + well-formed helper cannot: a file part with no filename, a declared per-part content + type, a repeated or bracketed field name, or a non-UTF-8 value.""" + return ( + b"".join( + f"--{boundary}\r\n{head}\r\n\r\n".encode() + content + b"\r\n" + for head, content in parts + ) + + f"--{boundary}--\r\n".encode() + ) + + +def upload_key(body: bytes, boundary: str) -> str: + content_type: Final = f"multipart/form-data; boundary={boundary}" + return canonicalize(edge_request("POST", UPLOAD_PATH, "", body, content_type)).key + + +DISPOSITION = 'Content-Disposition: form-data; name="{name}"' +FILE_DISPOSITION = DISPOSITION + '; filename="{filename}"' + + +class TestMultipartIdentityEdges: + """The identity a multipart upload keys on, pinned against the ways two materially + different uploads could otherwise collapse onto one key. A collision here is the + dangerous failure: replay would answer one request with another's response.""" + + def test_a_declared_part_content_type_separates_otherwise_identical_uploads(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + as_json = raw_multipart( + boundary, + (FILE_DISPOSITION.format(name="file", filename="a") + "\r\nContent-Type: application/json", b"xy"), + ) + as_csv = raw_multipart( + boundary, + (FILE_DISPOSITION.format(name="file", filename="a") + "\r\nContent-Type: text/csv", b"xy"), + ) + + assert upload_key(as_json, boundary) != upload_key(as_csv, boundary) + + def test_a_file_part_without_a_filename_is_not_mistaken_for_a_plain_field(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + upload = raw_multipart( + boundary, + (DISPOSITION.format(name="file") + "\r\nContent-Type: application/octet-stream", b"CONTENT"), + ) + plain_field = raw_multipart(boundary, (DISPOSITION.format(name="file"), b"CONTENT")) + + request = edge_request( + "POST", UPLOAD_PATH, "", upload, f"multipart/form-data; boundary={boundary}" + ) + + assert upload_key(upload, boundary) != upload_key(plain_field, boundary) + assert request.form == {} + assert b"CONTENT".decode() not in request.model_dump_json() + + def test_a_filename_carrying_a_per_run_marker_keys_the_same_next_run(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + + def upload(marker: str) -> str: + body = raw_multipart( + boundary, + (FILE_DISPOSITION.format(name="one", filename=f"{marker}.jsonl"), b"first"), + (FILE_DISPOSITION.format(name="two", filename="steady.jsonl"), b"second"), + ) + return upload_key(body, boundary) + + assert upload("a1b2c3d4e5f6") == upload("0f9e8d7c6b5a") + + def test_a_separator_inside_a_filename_cannot_forge_a_different_split(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + colon_in_filename = raw_multipart( + boundary, (FILE_DISPOSITION.format(name="file", filename="a:b.jsonl"), b"same") + ) + colon_in_field = raw_multipart( + boundary, (FILE_DISPOSITION.format(name="file:a", filename="b.jsonl"), b"same") + ) + + assert upload_key(colon_in_filename, boundary) != upload_key(colon_in_field, boundary) + + def test_a_repeated_field_cannot_collide_with_a_literal_indexed_name(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + repeated = raw_multipart( + boundary, + (DISPOSITION.format(name="purpose"), b"x"), + (DISPOSITION.format(name="purpose"), b"y"), + ) + literal_index = raw_multipart( + boundary, + (DISPOSITION.format(name="purpose"), b"x"), + (DISPOSITION.format(name="purpose[1]"), b"y"), + ) + + assert upload_key(repeated, boundary) != upload_key(literal_index, boundary) + + def test_two_binary_field_values_of_one_length_stay_apart(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + first = raw_multipart(boundary, (DISPOSITION.format(name="blob"), b"\xff\xfe\xfd")) + second = raw_multipart(boundary, (DISPOSITION.format(name="blob"), b"\xf0\xf1\xf2")) + + assert upload_key(first, boundary) != upload_key(second, boundary) + + def test_a_secret_named_field_never_reaches_the_stored_request(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + body = raw_multipart( + boundary, + (DISPOSITION.format(name="openai_api_key"), b"sk-live-DEADBEEF-0123456789abcd"), + (DISPOSITION.format(name="purpose"), b"batch"), + ) + + request = edge_request( + "POST", UPLOAD_PATH, "", body, f"multipart/form-data; boundary={boundary}" + ) + + assert "sk-live-DEADBEEF-0123456789abcd" not in request.model_dump_json() + assert request.form == {"openai_api_key": "", "purpose": "batch"} + + def test_a_redacted_field_still_matches_the_live_request_that_carried_the_secret( + self, + ) -> None: + boundary = "0123456789abcdef0123456789abcdef" + + def upload(secret: str) -> str: + body = raw_multipart( + boundary, + (DISPOSITION.format(name="openai_api_key"), secret.encode()), + (DISPOSITION.format(name="purpose"), b"batch"), + ) + return upload_key(body, boundary) + + assert upload("sk-live-DEADBEEF-0123456789abcd") == upload("") + + def test_a_length_change_the_canonicalizer_absorbs_does_not_move_the_key(self) -> None: + boundary = "0123456789abcdef0123456789abcdef" + + def upload(created: str) -> str: + body = raw_multipart( + boundary, + ( + FILE_DISPOSITION.format(name="file", filename="batch.jsonl"), + b'{"created_at":"' + created.encode() + b'"}', + ), + ) + return upload_key(body, boundary) + + assert upload("2026-08-21T02:08:19Z") == upload("2026-08-21T02:08:19.123456Z") + + @pytest.mark.parametrize( + "content_type", + [ + pytest.param("multipart/form-data; myboundary=zzz; boundary={boundary}", id="lookalike-parameter"), + pytest.param("multipart/form-data; BOUNDARY={boundary}", id="uppercase-parameter"), + ], + ) + def test_the_boundary_parameter_is_read_the_way_the_client_meant_it( + self, content_type: str + ) -> None: + boundary = "0123456789abcdef0123456789abcdef" + body = raw_multipart( + boundary, (FILE_DISPOSITION.format(name="file", filename="batch.jsonl"), BATCH_JSONL) + ) + + request = edge_request( + "POST", UPLOAD_PATH, "", body, content_type.format(boundary=boundary) + ) + + assert request.form == {} + assert request.file_name is not None + assert "batch.jsonl" in request.file_name + + def test_an_empty_declared_boundary_falls_back_instead_of_splitting_on_dashes(self) -> None: + body = b'--\r\nContent-Disposition: form-data; name="a"\r\n\r\nvalue\r\n----\r\n' + + request = edge_request( + "POST", UPLOAD_PATH, "", body, 'multipart/form-data; boundary=""' + ) + + assert request.form is None + assert request.file_sha256 is not None + + class TestReplayLeftover: def test_partially_consumed_recording_names_the_leftover(self, tmp_path: Path) -> None: root = tmp_path / "bundle" From 060e40021def9354de8732bdfad0ffb578c634c4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 21 Aug 2026 19:47:47 -0700 Subject: [PATCH 20/43] fix(anthropic): resolve the provider exactly once on /v1/messages (#37757) get_llm_provider ran in the messages handler and again inside completion, so a provider/vendor/model id lost its vendor segment and reached upstream bare. Pass the caller's unresolved model down the bridge instead, and move the responses marker into the canonical provider/responses/model slot. Reporting stays provider-local on both bridges: message_start names the id the provider itself knows, through a shared local_model_name helper. Fixes #37716 --- .../adapters/handler.py | 11 +- .../messages/handler.py | 41 ++++++- .../responses_adapters/handler.py | 9 +- .../experimental_pass_through/utils.py | 5 + .../adapters/test_handler_prompt_cache_key.py | 2 +- ...erimental_pass_through_messages_handler.py | 112 +++++++++++++++++- .../test_responses_adapters_handler.py | 39 ++++++ 7 files changed, 204 insertions(+), 15 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 89066e33cbc..9d61701d26d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -21,6 +21,7 @@ from litellm.llms.anthropic.experimental_pass_through.context_management import ) from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, + local_model_name, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -358,9 +359,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: except Exception: pass - if isinstance(model, str) and model and not model.startswith("responses/"): - # Prefix model with "responses/" to route to OpenAI Responses API - completion_kwargs["model"] = f"responses/{model}" + if isinstance(model, str) and model and "responses/" not in model: + local_model: Final = model.removeprefix(f"{custom_llm_provider}/") + completion_kwargs["model"] = f"{custom_llm_provider}/responses/{local_model}" auto_summary: Final = is_reasoning_auto_summary_enabled() @@ -616,7 +617,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( completion_response, - model=model, + model=local_model_name(model, kwargs.get("custom_llm_provider")), tool_name_mapping=tool_name_mapping, polyfill_result=polyfill_result, is_async=True, @@ -750,7 +751,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( completion_response, - model=model, + model=local_model_name(model, kwargs.get("custom_llm_provider")), tool_name_mapping=tool_name_mapping, polyfill_result=polyfill_result, is_async=False, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 26aef666172..f4d24bb933c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -42,15 +42,46 @@ from .utils import AnthropicMessagesRequestUtils, mock_response _RESPONSES_API_PROVIDERS: Final = frozenset({"openai"}) -def _should_route_to_responses_api(custom_llm_provider: str | None) -> bool: - """Return True when the provider should use the Responses API path. +def _bridges_to_responses_api(model: str, custom_llm_provider: str) -> bool: + from litellm.main import responses_api_bridge_check + + model_info, _ = responses_api_bridge_check(model=model, custom_llm_provider=custom_llm_provider) + return model_info.get("mode") == "responses" + + +def _responses_mode_is_lost_by_prefix_strip( + requested_model: str, resolved_model: str, custom_llm_provider: str +) -> bool: + """Whether a Responses-only deployment stops looking like one once its provider prefix is stripped. + + ``litellm.completion`` re-derives the Responses bridge from the stripped id alone, so a + deployment id such as ``perplexity/perplexity/sonar`` (mode ``responses``) is shadowed by the + chat entry ``perplexity/sonar`` and would otherwise be sent to chat/completions. + """ + if requested_model == resolved_model: + return False + return _bridges_to_responses_api(requested_model, custom_llm_provider) and not _bridges_to_responses_api( + resolved_model, custom_llm_provider + ) + + +def _should_route_to_responses_api( + custom_llm_provider: str | None, + requested_model: str | None = None, + resolved_model: str | None = None, +) -> bool: + """Return True when the request should use the Responses API path. Set ``litellm.use_chat_completions_url_for_anthropic_messages = True`` to opt out and route OpenAI/Azure requests through chat/completions instead. """ if litellm.use_chat_completions_url_for_anthropic_messages: return False - return custom_llm_provider in _RESPONSES_API_PROVIDERS + if custom_llm_provider in _RESPONSES_API_PROVIDERS: + return True + if custom_llm_provider is None or requested_model is None or resolved_model is None: + return False + return _responses_mode_is_lost_by_prefix_strip(requested_model, resolved_model, custom_llm_provider) def _deployment_passes_through_anthropic_messages(model_info: object) -> bool: @@ -533,7 +564,7 @@ def anthropic_messages_handler( _shared_kwargs: Final = dict( max_tokens=max_tokens, messages=messages, - model=model, + model=original_model, metadata=metadata, stop_sequences=stop_sequences, stream=stream, @@ -551,7 +582,7 @@ def anthropic_messages_handler( custom_llm_provider=custom_llm_provider, **kwargs, ) - if _should_route_to_responses_api(custom_llm_provider): + if _should_route_to_responses_api(custom_llm_provider, original_model, model): return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(**_shared_kwargs) # The in-gateway context_management polyfill runs inside diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 843cda249c5..c1ea39fd72c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( ) from litellm.types.llms.openai import ResponsesAPIResponse +from ..utils import local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper from .transformation import LiteLLMAnthropicToResponsesAPIAdapter @@ -179,7 +180,9 @@ class LiteLLMMessagesToResponsesAPIHandler: result: Final = await litellm.aresponses(**responses_kwargs) if stream: - wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) + wrapper: Final = AnthropicResponsesStreamWrapper( + responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")) + ) return wrapper.async_anthropic_sse_wrapper() if not isinstance(result, ResponsesAPIResponse): @@ -257,7 +260,9 @@ class LiteLLMMessagesToResponsesAPIHandler: result: Final = litellm.responses(**responses_kwargs) if stream: - wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=result, model=model) + wrapper: Final = AnthropicResponsesStreamWrapper( + responses_stream=result, model=local_model_name(model, kwargs.get("custom_llm_provider")) + ) return wrapper.async_anthropic_sse_wrapper() if not isinstance(result, ResponsesAPIResponse): diff --git a/litellm/llms/anthropic/experimental_pass_through/utils.py b/litellm/llms/anthropic/experimental_pass_through/utils.py index c5abcf8c04c..29661572b73 100644 --- a/litellm/llms/anthropic/experimental_pass_through/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/utils.py @@ -13,6 +13,11 @@ def prompt_cache_key_from_user_id(user_id: object) -> str | None: return str(user_id)[:OPENAI_MAX_PROMPT_CACHE_KEY_LENGTH] or None +def local_model_name(model: str, custom_llm_provider: object) -> str: + """The id the provider itself knows, for reporting back to the caller in ``message_start``.""" + return model.removeprefix(f"{custom_llm_provider}/") if isinstance(custom_llm_provider, str) else model + + def is_reasoning_auto_summary_enabled() -> bool: """Check whether the default 'summary: detailed' injection is enabled (opt-in).""" return litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py index 5b7f2a60f68..f48d51dbe1e 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_prompt_cache_key.py @@ -66,5 +66,5 @@ def test_prepare_completion_kwargs_keeps_prompt_cache_key_through_responses_rero {"custom_llm_provider": "openai"}, thinking={"type": "enabled", "budget_tokens": 1024}, ) - assert completion_kwargs["model"] == "responses/openai/gpt-5.6-luna" + assert completion_kwargs["model"] == "openai/responses/gpt-5.6-luna" assert completion_kwargs["prompt_cache_key"] == "session-abc" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 91f5023496a..15ac73ed352 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -217,7 +217,10 @@ async def _async_return(value): def test_anthropic_experimental_pass_through_messages_handler_custom_llm_provider(): """ - Test that litellm.completion is called when a custom LLM provider is given + Test that litellm.completion is called when a custom LLM provider is given. + + Provider resolution now happens exactly once, inside litellm.completion itself + (BerriAI/litellm#37716), so the handler passes the original unresolved model through. """ from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( anthropic_messages_handler, @@ -241,7 +244,7 @@ def test_anthropic_experimental_pass_through_messages_handler_custom_llm_provide # Verify that the custom provider was passed through call_kwargs = mock_completion.call_args.kwargs assert call_kwargs["custom_llm_provider"] == "my-custom-llm" - assert call_kwargs["model"] == "my-custom-llm/my-custom-model" + assert call_kwargs["model"] == "my-custom-model" assert call_kwargs["api_key"] == "test-api-key" @@ -997,3 +1000,108 @@ def test_first_party_claude_4_8_plus_cost_map_entries_carry_mid_conversation_sys and info.get("supports_mid_conversation_system") is not True ] assert missing == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model, expected_wire_model, expected_url", + [ + ( + "perplexity/perplexity/kimi-k3", + "perplexity/kimi-k3", + "https://api.perplexity.ai/v1/responses", + ), + ( + "perplexity/perplexity/sonar", + "perplexity/sonar", + "https://api.perplexity.ai/v1/responses", + ), + ("perplexity/sonar", "sonar", "https://api.perplexity.ai/chat/completions"), + ], +) +async def test_messages_strips_provider_prefix_exactly_once( + requested_model, expected_wire_model, expected_url +): + """ + BerriAI/litellm#37716: only the leading provider segment may be stripped on the way upstream. + + A multi-segment id such as perplexity/perplexity/kimi-k3 must reach the provider as + perplexity/kimi-k3, matching what /v1/chat/completions and /v1/responses already send. + + The endpoint is asserted alongside the body because perplexity/perplexity/sonar is a + Responses-only deployment whose bare id perplexity/sonar is an ordinary chat model, so + stripping the prefix must not also move the request onto chat/completions. + + The subject is the outbound request, so the transport is cut at the wire rather than + stubbed with a response body: these ids take different bridges (chat completions + versus the Responses API) and would otherwise need different response shapes. + """ + captured = {} + + async def fake_send(self, request, **kwargs): + captured["body"] = json.loads(request.content) + captured["url"] = str(request.url) + raise httpx.ConnectError("cut at the wire", request=request) + + with ( + patch.object(httpx.AsyncClient, "send", fake_send), + pytest.raises(litellm.exceptions.InternalServerError), + ): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model=requested_model, + api_key="test-api-key", + ) + + assert captured["body"]["model"] == expected_wire_model + assert captured["url"] == expected_url + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model, expected_reported_model", + [ + ("perplexity/perplexity/kimi-k3", "perplexity/kimi-k3"), + ("perplexity/sonar", "sonar"), + ], +) +async def test_messages_streaming_reports_provider_local_model(requested_model, expected_reported_model): + """ + BerriAI/litellm#37716: the wire keeps every segment, so ``message_start`` must still + report the id the provider itself knows rather than the caller's prefixed deployment id. + """ + + class _EmptyStream: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + with patch("litellm.acompletion", new=AsyncMock(return_value=_EmptyStream())): + stream = await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model=requested_model, + api_key="test-api-key", + stream=True, + ) + first_event = await stream.__anext__() + + assert json.loads(first_event.decode().split("data: ", 1)[1])["message"]["model"] == expected_reported_model + + +def test_messages_sync_streaming_reports_provider_local_model(): + """Same guarantee as the async bridge, at the sync call site.""" + with patch("litellm.completion", new=MagicMock(return_value=iter(()))): + stream = litellm.anthropic.messages.create( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="perplexity/perplexity/kimi-k3", + api_key="test-api-key", + stream=True, + ) + first_event = next(iter(stream)) + + assert json.loads(first_event.decode().split("data: ", 1)[1])["message"]["model"] == "perplexity/kimi-k3" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py index 7ef3077f9d7..589dc64f9b9 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_handler.py @@ -1,9 +1,15 @@ +import json import os import sys +from unittest.mock import AsyncMock, patch + +import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) +import litellm from litellm.llms.anthropic.experimental_pass_through.responses_adapters.handler import ( + LiteLLMMessagesToResponsesAPIHandler, _build_responses_kwargs, ) @@ -43,3 +49,36 @@ def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key(): ) assert "user" not in responses_kwargs assert "prompt_cache_key" not in responses_kwargs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model, expected_reported_model", + [ + ("openai/gpt-5.6-luna", "gpt-5.6-luna"), + ("perplexity/perplexity/kimi-k3", "perplexity/kimi-k3"), + ], +) +async def test_streaming_message_start_reports_the_provider_local_model(requested_model, expected_reported_model): + """ + BerriAI/litellm#37716 sends the caller's unresolved id down this bridge so the provider + resolves it once. ``message_start`` is a reporting field rather than a wire value, so it + keeps naming the model as the provider knows it, with only the leading provider segment gone. + """ + + async def empty_stream(): + return + yield + + with patch.object(litellm, "aresponses", AsyncMock(return_value=empty_stream())): + sse = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=1024, + messages=MESSAGES, + model=requested_model, + stream=True, + custom_llm_provider=requested_model.split("/")[0], + ) + events = [json.loads(chunk.decode().split("data: ", 1)[1]) async for chunk in sse] + + message_start = next(e for e in events if e["type"] == "message_start") + assert message_start["message"]["model"] == expected_reported_model From afec9b8ab9013fad3156690df1b0b0505052fe64 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:15:08 -0700 Subject: [PATCH 21/43] perf(ci): cache the Rust build the unit shards compile from scratch (#37795) * perf(ci): cache the Rust build the unit shards compile from scratch Every unit shard installs the workspace, and the root package builds through maturin, so each of the eleven jobs compiles litellm-rust/crates/python-bridge in release mode before a single test runs. That step measured 2m40s a shard on 2026-08-21, which is more wall clock than the entire unit tier spends running tests, and none of it was cached: the uv cache covers wheels it downloads, not wheels it builds, and a path dependency whose source moves every commit can never hit that cache anyway. A composite action now exports CARGO_TARGET_DIR to a fixed workspace path and caches it alongside the Cargo registry, keyed on Cargo.lock. Cargo rebuilds only what changed, so a warm job pays for the bridge crate rather than its whole dependency graph. Measured locally, that is 34s cold against 8s warm, including after a Python-only or Rust-only edit. The absolute path matters: uv builds the wheel from its own working directory, so a relative target directory lands the artifacts where nothing can find them again. * perf(ci): cache the Rust build in the other four workflows that sync the workspace code-quality, mcp, documentation and the schema.d.ts check each install the workspace and so each compile the bridge from scratch, measured at 138s, 177s, 163s and 154s on 2026-08-21. The lint job pays the same and is left to #37783, which already owns that file's setup section. * perf(ci): cache the Rust build in the lint job too * fix(ci): cache cargo's own target directory instead of redirecting it uv builds the wheel in place, so cargo already writes to litellm-rust/target, which test-rust.yml has cached all along. Redirecting CARGO_TARGET_DIR bought nothing and cost a GITHUB_ENV write that zizmor rejects as a code-execution path. * chore(ci): raise the job backstop for the added setup step The cargo cache is a fifth bounded setup step, so the base's setup ceiling goes 30m to 35m and every job budget follows: 55 to 60, and proxy-server's 95 to 100. check_workflow_startup_safety enforces exactly this sum, and failed on the first push without it. * perf(ci): cache the Rust build in the four remaining workflows that sync Six workflows were wired; ten install the workspace. The four left out still compile the pyo3 bridge from scratch. test-terraform-provider.yml is the one that matters per PR: its endpoint-drift job triggers on any change under litellm/proxy/**. The other three are a scheduled load check, a manual mutation run, and the staging-push counts publisher, whose gate syncs the project inside scripts/type_check_gate.py rather than in a workflow step, so nothing in the file names the build it pays for. --- .github/actions/cache-cargo-build/action.yml | 31 +++++++++++++++++++ .github/workflows/_test-unit-base.yml | 9 ++++-- .github/workflows/check-ui-api-types.yml | 4 +++ .github/workflows/mutation-test.yml | 3 ++ .../publish-basedpyright-base-counts.yml | 3 ++ .github/workflows/test-code-quality.yml | 3 ++ .github/workflows/test-linting.yml | 4 +++ .github/workflows/test-mcp.yml | 4 +++ .github/workflows/test-terraform-provider.yml | 3 ++ .github/workflows/test-unit-documentation.yml | 4 +++ .github/workflows/test-unit.yml | 22 ++++++------- .github/workflows/weekly_load_anomaly.yml | 3 ++ 12 files changed, 80 insertions(+), 13 deletions(-) create mode 100644 .github/actions/cache-cargo-build/action.yml diff --git a/.github/actions/cache-cargo-build/action.yml b/.github/actions/cache-cargo-build/action.yml new file mode 100644 index 00000000000..36c6c790b84 --- /dev/null +++ b/.github/actions/cache-cargo-build/action.yml @@ -0,0 +1,31 @@ +name: "Cache the Rust build" +description: >- + Cache the Cargo registry and target directory the root package's build needs, + so only the first job on a given Cargo.lock compiles the bridge from scratch. + + litellm builds through maturin, which compiles litellm-rust/crates/python-bridge + in release mode before it can produce a wheel. `uv sync` therefore pays a full + build in every job that installs the workspace: measured at 2m40s per unit shard + on 2026-08-21, more than the whole unit tier spends running tests. Nothing caught + it, because the uv cache holds wheels uv downloads rather than wheels it builds, + and a path dependency whose source moves every commit could never hit that cache + anyway. Cargo rebuilds only what changed when its target directory survives, so a + warm job pays for the bridge crate alone. + + The key namespace is separate from test-rust.yml's. Both cache the same directory, + but that workflow fills it with debug and clippy artifacts, which a release build + cannot reuse, and a shared key would let whichever ran first deny the other a save. + +runs: + using: composite + steps: + - name: Restore the Cargo registry and target directory + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-release- diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 54f50524a39..b7d185bd0b9 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -27,7 +27,7 @@ on: default: 20 job-timeout-minutes: description: >- - Backstop for the whole job. Keep it >= `timeout-minutes` plus 35: 30 for + Backstop for the whole job. Keep it >= `timeout-minutes` plus 40: 35 for the per-step ceilings on the setup steps below, and 5 for the runner overhead the job clock charges but no step owns (job init, step transitions, post-job cleanup). That headroom is what makes the test @@ -36,7 +36,7 @@ on: arithmetic, so the sum is passed in rather than computed. required: false type: number - default: 55 + default: 60 max-failures: description: "Stop after this many failures" required: false @@ -103,6 +103,11 @@ jobs: restore-keys: | ${{ runner.os }}-uv- + - name: Cache the Rust build + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' timeout-minutes: 8 diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index dbd663a2efa..285676a0ddd 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -67,6 +67,10 @@ jobs: restore-keys: | ${{ runner.os }}-uv- + - name: Cache the Rust build + if: steps.changes.outputs.relevant == 'true' + uses: ./.github/actions/cache-cargo-build + - name: Install backend dependencies if: steps.changes.outputs.relevant == 'true' run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 68317d5dd12..602c26a3e98 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -53,6 +53,9 @@ jobs: restore-keys: | ${{ runner.os }}-uv- + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml diff --git a/.github/workflows/publish-basedpyright-base-counts.yml b/.github/workflows/publish-basedpyright-base-counts.yml index 71e196d8361..cd443a8e9db 100644 --- a/.github/workflows/publish-basedpyright-base-counts.yml +++ b/.github/workflows/publish-basedpyright-base-counts.yml @@ -43,6 +43,9 @@ jobs: with: version: "0.10.9" + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + - name: Cache Prisma binaries uses: ./.github/actions/cache-prisma-binaries diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 8f62837d29a..2a832d1956e 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -56,6 +56,9 @@ jobs: restore-keys: | ${{ runner.os }}-uv- + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies run: uv sync --frozen --all-groups --all-extras diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index e031ba46773..ccb58f5cc9c 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -78,6 +78,10 @@ jobs: run: | uv lock --check || (echo "❌ uv.lock is out of sync with pyproject.toml. Run 'uv lock' locally and commit the result." && exit 1) + - name: Cache the Rust build + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 95187ef2835..6ea814dc2de 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -47,6 +47,10 @@ jobs: with: version: "0.10.9" + - name: Cache the Rust build + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | diff --git a/.github/workflows/test-terraform-provider.yml b/.github/workflows/test-terraform-provider.yml index 7ea22825f4f..e46432e0e31 100644 --- a/.github/workflows/test-terraform-provider.yml +++ b/.github/workflows/test-terraform-provider.yml @@ -88,6 +88,9 @@ jobs: restore-keys: | ${{ runner.os }}-uv- + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index cb8035aafa1..90b6b28374e 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -67,6 +67,10 @@ jobs: restore-keys: | ${{ runner.os }}-uv- + - name: Cache the Rust build + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 3d6fffe7304..71eb0958bec 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -55,7 +55,7 @@ jobs: workers: 2 reruns: 1 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: enterprise-routing artifact-name: enterprise-routing @@ -67,7 +67,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: integrations artifact-name: integrations @@ -75,7 +75,7 @@ jobs: workers: 2 reruns: 3 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: Vertex AI artifact-name: llm-vertex-ai @@ -83,7 +83,7 @@ jobs: workers: 1 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: All Other Providers artifact-name: llm-other-providers @@ -91,7 +91,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: misc artifact-name: misc @@ -122,7 +122,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: proxy-auth artifact-name: proxy-auth @@ -134,7 +134,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: proxy-endpoints artifact-name: proxy-endpoints @@ -171,7 +171,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: proxy-server artifact-name: proxy-server @@ -179,7 +179,7 @@ jobs: workers: 4 reruns: 2 timeout-minutes: 60 - job-timeout-minutes: 95 + job-timeout-minutes: 100 - shard: proxy-infra artifact-name: proxy-infra @@ -198,7 +198,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: responses-caching-types artifact-name: responses-caching-types @@ -209,7 +209,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 2dffc889d0e..3e1fca89645 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -47,6 +47,9 @@ jobs: with: version: "0.10.9" + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + - name: Install dependencies run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy From 35fcc9f7b849101b35d954e7a11852b58d30ac1a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:15:13 -0700 Subject: [PATCH 22/43] test(proxy): pin the request-body rules `proxy/_types.py` enforces (#37811) Eight validators in that module decide what a request body may say, and none of them was asserted anywhere. Reversing any one of the eight left the file green. Cover them at the API boundary: a JWT issuer must pick audience validation or opt out, a temp budget needs both halves, an empty max budget reads as no limit, an organization member can only take a role the organization has, an LLM-backed injection check needs the call it would make, and four server-only markers are never taken from the caller. The injection case builds each incomplete body as its own value rather than deleting a key out of the one it is iterating. --- tests/test_litellm/proxy/test_proxy_types.py | 104 +++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py index 4a93e9ac7ba..77083af48c0 100644 --- a/tests/test_litellm/proxy/test_proxy_types.py +++ b/tests/test_litellm/proxy/test_proxy_types.py @@ -177,3 +177,107 @@ def test_project_io_token_limits_are_stored_in_metadata(request_type): assert request.metadata == limits assert request.model_dump(exclude_none=True)["metadata"] == limits + + +def test_a_jwt_issuer_must_pick_audience_validation_or_opt_out(): + from pydantic import ValidationError + + from litellm.proxy._types import JWTIssuerConfig + + with pytest.raises(ValidationError, match="must configure audience or set disable_audience_validation"): + JWTIssuerConfig(issuer="https://issuer.example.com") + + with pytest.raises(ValidationError, match="cannot set audience and disable_audience_validation"): + JWTIssuerConfig( + issuer="https://issuer.example.com", + audience="litellm-proxy", + disable_audience_validation=True, + ) + + assert JWTIssuerConfig(issuer="https://issuer.example.com", audience="litellm-proxy").audience == "litellm-proxy" + assert ( + JWTIssuerConfig(issuer="https://issuer.example.com", disable_audience_validation=True).audience + is None + ) + + +def test_a_jwt_issuer_rejects_a_field_it_does_not_define(): + from pydantic import ValidationError + + from litellm.proxy._types import JWTIssuerConfig + + with pytest.raises(ValidationError, match="Extra inputs are not permitted"): + JWTIssuerConfig(issuer="https://issuer.example.com", audience="a", jwks_uri="https://issuer/jwks") + + +def test_a_temp_budget_needs_both_halves_or_neither(): + from pydantic import ValidationError + + from litellm.proxy._types import UpdateKeyRequest + + with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"): + UpdateKeyRequest(key="sk-1234", temp_budget_increase=10) + + with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"): + UpdateKeyRequest(key="sk-1234", temp_budget_expiry="2026-01-01") + + both = UpdateKeyRequest(key="sk-1234", temp_budget_increase=10, temp_budget_expiry="2026-01-01") + assert both.temp_budget_increase == 10 + + +def test_an_empty_max_budget_is_read_as_no_limit(): + from litellm.proxy._types import GenerateKeyRequest + + assert GenerateKeyRequest(max_budget="").max_budget is None + assert GenerateKeyRequest(max_budget=25).max_budget == 25 + + +def test_an_organization_member_can_only_take_a_role_the_organization_has(): + from pydantic import ValidationError + + from litellm.proxy._types import LitellmUserRoles, OrganizationMemberUpdateRequest + + with pytest.raises(ValidationError, match="Invalid role"): + OrganizationMemberUpdateRequest( + organization_id="org-1", user_id="user-1", role=LitellmUserRoles.PROXY_ADMIN + ) + + allowed = OrganizationMemberUpdateRequest( + organization_id="org-1", user_id="user-1", role=LitellmUserRoles.ORG_ADMIN + ) + assert allowed.role == LitellmUserRoles.ORG_ADMIN + + +def test_an_llm_backed_injection_check_needs_the_call_it_would_make(): + from pydantic import ValidationError + + from litellm.proxy._types import LiteLLMPromptInjectionParams + + for missing in ("llm_api_name", "llm_api_system_prompt", "llm_api_fail_call_string"): + complete = { + "llm_api_name": "gpt-4o", + "llm_api_system_prompt": "is this an injection", + "llm_api_fail_call_string": "yes", + } + del complete[missing] + with pytest.raises(ValidationError, match=f"{missing} must be provided"): + LiteLLMPromptInjectionParams(llm_api_check=True, **complete) + + assert LiteLLMPromptInjectionParams(llm_api_check=False).llm_api_name is None + + +@pytest.mark.parametrize( + "field, forged, default", + [ + ("mcp_admitted_user_subject", "someone-else", False), + ("mcp_source_team_rpm_limits", {"team-1": 10_000}, None), + ("mcp_session_resource_server_id", "server-1", None), + ("via_virtual_key", "sk-someone-elses-key", False), + ], +) +def test_a_server_only_marker_is_not_taken_from_the_caller(field, forged, default): + from litellm.proxy._types import UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-1234", **{field: forged}) + + assert getattr(auth, field) == default From b416bdadd36564e6675864ef88b55740a4745baf Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:15:27 -0700 Subject: [PATCH 23/43] test(main): pin what a streamed response costs, end to end (#37812) Rebuilding a streamed response and pricing it is the path a spend row comes from, and nothing asserted it end to end. Reversing either half of the usage the provider reported left the file green. Three cases: the rebuilt response bills the usage the last chunk carried, streaming and not streaming bill the same usage the same, and a stream that reported no usage is still billed rather than dropped. The cost is asserted against the catalog prices the run itself reads, with a non-zero guard in front of it so an all-zeros lookup cannot satisfy it vacuously. Pinning the dollar figure as a literal would have made a routine gpt-4o price update fail a test about usage reconstruction. --- tests/test_litellm/test_main.py | 106 ++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4b223a3a900..4ab09d9d85b 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2754,3 +2754,109 @@ def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6() {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} ] assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "explicit"} + + +STREAM_COST_MODEL = "gpt-4o" +STREAMED_USAGE = {"prompt_tokens": 137, "completion_tokens": 42, "total_tokens": 179} + + +def _text_chunk(content, finish_reason=None, usage=None): + chunk = { + "id": "chatcmpl-stream-cost", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": STREAM_COST_MODEL, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": content}, + "finish_reason": finish_reason, + } + ], + } + if usage is not None: + chunk["usage"] = usage + return chunk + + +def _priced_at(prompt_tokens, completion_tokens): + prices = litellm.model_cost[STREAM_COST_MODEL] + return ( + prompt_tokens * prices["input_cost_per_token"] + + completion_tokens * prices["output_cost_per_token"] + ) + + +@pytest.fixture +def local_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + +def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), + ], + messages=[{"role": "user", "content": "hi"}], + ) + + assert rebuilt.choices[0].message.content == "Hello there" + assert rebuilt.usage.prompt_tokens == STREAMED_USAGE["prompt_tokens"] + assert rebuilt.usage.completion_tokens == STREAMED_USAGE["completion_tokens"] + + cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) + + assert cost == pytest.approx(_priced_at(137, 42)) + assert cost == pytest.approx(0.0007625) + + +def test_streaming_and_not_streaming_bill_the_same_usage_the_same(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop", usage=STREAMED_USAGE), + ], + messages=[{"role": "user", "content": "hi"}], + ) + whole = litellm.ModelResponse( + id="chatcmpl-stream-cost", + model=STREAM_COST_MODEL, + object="chat.completion", + created=1700000000, + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello there"}, + "finish_reason": "stop", + } + ], + usage=STREAMED_USAGE, + ) + + assert litellm.completion_cost( + completion_response=rebuilt, model=STREAM_COST_MODEL + ) == pytest.approx(litellm.completion_cost(completion_response=whole, model=STREAM_COST_MODEL)) + + +def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): + rebuilt = litellm.stream_chunk_builder( + chunks=[ + _text_chunk("Hello"), + _text_chunk(" there"), + _text_chunk(None, finish_reason="stop"), + ], + messages=[{"role": "user", "content": "hi"}], + ) + + assert rebuilt.usage.prompt_tokens > 0 + assert rebuilt.usage.completion_tokens > 0 + + cost = litellm.completion_cost(completion_response=rebuilt, model=STREAM_COST_MODEL) + + assert cost > 0 + assert cost == pytest.approx( + _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) + ) From 89649e4141c72914a94e2596bd4ca4288112851f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:15:32 -0700 Subject: [PATCH 24/43] test(proxy): pin what a failed request records as usage and spend (#37813) Six helpers in `litellm/proxy/utils.py` decide the usage a failed request records, and none of them is named anywhere in the suite. Two of their decisions could be reversed with the file still green: a request with nothing countable in it lifted as a zero-token usage, and a request that never reached a provider billed for input it never sent. Twelve cases asserting those contracts directly, plus a canary pinning the literal no-upstream-call key the module branches on, so a rename cannot pass silently. --- tests/test_litellm/proxy/test_proxy_utils.py | 190 +++++++++++++++++++ 1 file changed, 190 insertions(+) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index fe79ef25da6..deb49ff9f54 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1644,3 +1644,193 @@ async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks): assert guardrail.call_count == 0 assert [item.text for item in returned.content] == ["jane@example.com"] + + +FAILURE_USAGE_MODEL = "gpt-4o" +ONE_USER_MESSAGE = [{"role": "user", "content": "hi"}] + + +class _LoggingObj: + def __init__(self, model_call_details): + self.model_call_details = model_call_details + + +@pytest.mark.parametrize( + "system_input, expected", + [ + ("be brief", "be brief"), + ([{"type": "text", "text": "a"}, {"type": "text", "text": "b"}], "ab"), + (["a", {"text": "b"}], "ab"), + ([{"type": "image"}], ""), + (None, ""), + (17, ""), + ], +) +def test_a_system_prompt_reads_the_same_whatever_shape_it_arrived_in(system_input, expected): + from litellm.proxy.utils import _system_prompt_text + + assert _system_prompt_text(system_input) == expected + + +def test_a_system_prompt_is_counted_on_top_of_the_request(): + from litellm.proxy.utils import _count_request_input_tokens + + without = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) + with_system = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") + + assert without > 0 + assert with_system > without + + +def test_a_request_with_nothing_in_it_counts_zero(): + from litellm.proxy.utils import _count_request_input_tokens + + assert _count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 + assert _count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 + + +def test_a_failed_dispatch_is_estimated_as_input_only(): + from litellm.proxy.utils import _count_request_input_tokens, _estimate_dispatched_failure_usage + + usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + + assert usage is not None + assert usage.prompt_tokens == _count_request_input_tokens( + FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None + ) + assert usage.completion_tokens == 0 + assert usage.total_tokens == usage.prompt_tokens + + +@pytest.mark.parametrize("request_input", [[], object()]) +def test_nothing_is_estimated_when_there_is_nothing_to_count(request_input): + from litellm.proxy.utils import _estimate_dispatched_failure_usage + + assert _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, request_input, None) is None + + +def test_usage_the_stream_already_recovered_beats_an_estimate(): + from litellm.proxy.utils import _failure_usage_to_lift + from litellm.types.utils import Usage + + recovered = Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + + lifted = _failure_usage_to_lift( + model_call_details={"combined_usage_object": recovered, "response_cost": 0.25}, + request_body={}, + dispatched=True, + ) + + assert lifted == (recovered, 0.25) + + +def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): + from litellm.proxy.utils import _failure_usage_to_lift + + lifted = _failure_usage_to_lift( + model_call_details={ + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + }, + request_body={}, + dispatched=True, + ) + + assert lifted is not None + usage, response_cost = lifted + assert usage.prompt_tokens > 0 + assert usage.completion_tokens == 0 + assert response_cost == 0.0 + + +@pytest.mark.parametrize( + "model_call_details, dispatched", + [ + ({"call_type": "acompletion", "model": FAILURE_USAGE_MODEL, "messages": ONE_USER_MESSAGE}, False), + ( + { + "litellm_no_upstream_llm_call": True, + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + }, + True, + ), + ({"call_type": "afile_content", "model": FAILURE_USAGE_MODEL, "messages": ONE_USER_MESSAGE}, True), + ], + ids=["never dispatched", "no upstream call", "call type has no input to price"], +) +def test_a_failure_that_cost_the_provider_nothing_lifts_nothing(model_call_details, dispatched): + from litellm.proxy.utils import _failure_usage_to_lift + + assert _failure_usage_to_lift( + model_call_details=model_call_details, request_body={}, dispatched=dispatched + ) is None + + +def test_the_no_upstream_call_key_the_module_uses_is_the_one_asserted_above(): + from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL + + assert LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL == "litellm_no_upstream_llm_call" + + +def test_the_dispatched_system_prompt_wins_over_the_one_in_the_request_body(): + from litellm.proxy.utils import _failure_usage_to_lift + + def lift(model_call_details, request_body): + lifted = _failure_usage_to_lift( + model_call_details=model_call_details, request_body=request_body, dispatched=True + ) + assert lifted is not None + return lifted[0].prompt_tokens + + base = { + "call_type": "aanthropic_messages", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + } + long_system = "answer as briefly as you possibly can, in one short sentence" + + from_body = lift(base, {"system": long_system}) + from_params = lift({**base, "optional_params": {"system": "x"}}, {"system": long_system}) + body_only_short = lift(base, {"system": "x"}) + + assert from_body > body_only_short + assert from_params == body_only_short + + +def test_a_failure_with_no_logging_object_lifts_nothing(): + from litellm.proxy.utils import _failure_fields_to_lift + + assert dict(_failure_fields_to_lift({})) == {} + assert dict(_failure_fields_to_lift({"litellm_logging_obj": _LoggingObj({})})) == {} + + +def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): + from litellm.proxy.utils import _failure_fields_to_lift + + lifted = _failure_fields_to_lift( + { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + "standard_logging_object": {"id": "log-1"}, + } + ) + } + ) + + assert set(lifted) == { + "first_api_call_start_time", + "combined_usage_object", + "response_cost", + "standard_logging_object", + } + assert lifted["first_api_call_start_time"] == 1700000000.0 + assert lifted["response_cost"] == 0.0 + assert lifted["combined_usage_object"].prompt_tokens > 0 + assert lifted["standard_logging_object"] == {"id": "log-1"} From f88421bb4346f9b1a079f817421a3201bc872a79 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:15:42 -0700 Subject: [PATCH 25/43] test(llm_http_handler): pin the websocket and callback gates the request path branches on (#37814) The Responses WebSocket path, the pre-call deployment hook and the per-frame project quota hook are all selected by small predicates that nothing asserted directly. Mutating those four decisions left 4 of 6 mutants alive against the mapped test file. Cover them at the boundary: the rust WebSocket path needs both the openai provider and the rust flag, a plain CustomLogger must not advertise a pre-call deployment hook while an overriding or inheriting one must, and only callbacks that actually expose a callable enforce_project_io_token_quota_for_frame reach the WebSocket loop. Kill rate on those four decisions goes 2/6 -> 6/6; the file goes 66 -> 75 passing. --- .../custom_httpx/test_llm_http_handler.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index c87abbd8bc4..3c972ae9c84 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -24,7 +24,10 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, + _collect_ws_project_quota_callbacks, _google_genai_streaming_hidden_params, + _has_pre_call_deployment_hook, + _rust_responses_websocket_enabled, ) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams @@ -2445,3 +2448,82 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h collected = [chunk async for chunk in response] assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi" + + +@pytest.mark.parametrize( + "custom_llm_provider, litellm_params, expected", + [ + ("openai", GenericLiteLLMParams(rust=True), True), + ("openai", GenericLiteLLMParams(), False), + ("openai", GenericLiteLLMParams(rust=False), False), + ("azure", GenericLiteLLMParams(rust=True), False), + ("hosted_vllm", GenericLiteLLMParams(rust=True), False), + (None, GenericLiteLLMParams(rust=True), False), + ], +) +def test_the_rust_responses_websocket_needs_both_openai_and_the_rust_flag( + custom_llm_provider, litellm_params, expected +): + assert _rust_responses_websocket_enabled(custom_llm_provider, litellm_params) is expected + + +def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch): + from litellm.integrations.custom_logger import CustomLogger + + class _PlainLogger(CustomLogger): + pass + + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = [] + + monkeypatch.setattr(litellm, "callbacks", []) + assert _has_pre_call_deployment_hook(logging_obj) is False + + monkeypatch.setattr(litellm, "callbacks", [_PlainLogger()]) + assert _has_pre_call_deployment_hook(logging_obj) is False + + +def test_a_callback_that_overrides_the_deployment_hook_is_detected(monkeypatch): + from litellm.integrations.custom_logger import CustomLogger + + class _DeploymentHookLogger(CustomLogger): + async def async_pre_call_deployment_hook(self, kwargs, call_type): + return None + + class _InheritsTheHook(_DeploymentHookLogger): + pass + + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = [] + + monkeypatch.setattr(litellm, "callbacks", [_DeploymentHookLogger()]) + assert _has_pre_call_deployment_hook(logging_obj) is True + + monkeypatch.setattr(litellm, "callbacks", [_InheritsTheHook()]) + assert _has_pre_call_deployment_hook(logging_obj) is True + + monkeypatch.setattr(litellm, "callbacks", []) + logging_obj.dynamic_success_callbacks = [_DeploymentHookLogger()] + assert _has_pre_call_deployment_hook(logging_obj) is True + + +def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monkeypatch): + from litellm.integrations.custom_logger import CustomLogger + + class _PlainLogger(CustomLogger): + pass + + class _QuotaLogger(CustomLogger): + async def enforce_project_io_token_quota_for_frame(self, *args, **kwargs): + return None + + class _NotCallableAttribute: + enforce_project_io_token_quota_for_frame = "not a method" + + plain, quota, decoy = _PlainLogger(), _QuotaLogger(), _NotCallableAttribute() + + monkeypatch.setattr(litellm, "callbacks", [plain, decoy]) + assert _collect_ws_project_quota_callbacks() == () + + monkeypatch.setattr(litellm, "callbacks", [plain, quota, decoy]) + assert _collect_ws_project_quota_callbacks() == (quota,) From 9146667f801571ef9d11076eec5d3849f85c334f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:15:52 -0700 Subject: [PATCH 26/43] fix(ci): stop the mutation report publishing a score it never measured (#37825) * fix(ci): stop the mutation report publishing a score it never measured Run 32475268575 was the first dispatch of this workflow since May. Every setup step passed and mutmut generated all 48 mutant files, so the suspected zero-mutants bug is not what stops it. It dies in the stats phase, where mutmut times the configured test set once up front. That set included tests/proxy_behavior/management/, a behaviour tier that talks to a real seeded database, so the run ended having mutated nothing. Narrow tests_dir to the unit tier that maps to paths_to_mutate. Run 32476663383 proved a Postgres service is not enough on its own: with a schema but no seed rows the same test fails on a foreign key instead, and a mutation score is only meaningful against the tests that claim to cover the mutated code. The second half is the one that matters. With no results at all, mutation_report.py printed "No surviving mutants, the test suite caught every mutation" and exited 0, so a run that mutated nothing published a perfect score. It now separates no survivors from no results, says which it got, and exits 1. * fix(ci): count mutmut's multi-word verdicts as results The verdict capture was `\w+`, so it matched only single-word statuses. mutmut's status_by_exit_code table has four that are not: `no tests`, `not checked`, `caught by type check` and `check was interrupted by user`. A finished run made entirely of those parsed as zero results, which is exactly the state this script now treats as an unfinished run, so it would have failed a run that had in fact completed. The regression test asserting `reported == 2` on a three-verdict fixture was codifying that, and now asserts 3. A second test walks all four multi-word statuses and checks the report does not call the run unfinished. Caught by Greptile on #37825. * fix(ci): keep the saml tests out of the mutmut stats phase Run 32477695014 got past the database blocker and ran 208 of the configured tests, then ended on one error: test_saml_sso.py builds an x509 certificate in a fixture, and inside mutmut's mutants/ sandbox cryptography's hash classes are imported under a second identity, so .sign() rejects the SHA256 instance with "Algorithm must be a registered hash algorithm". That is a property of the sandbox, not of the tests or the code being mutated, and one erroring test ends the stats phase before a single mutant runs. * fix(ci): only claim a clean sweep when something was shown to be killed `mutmut results` skips killed mutants by design, so its silence means either that everything was killed or that nothing ran. Counting the verdicts it does print cannot tell those apart, which left the report still able to say the suite caught every mutation on a run whose mutants were all `no tests` or `not checked`. The clean-sweep sentence is now gated on mutmut-cicd-stats.json reporting a non-zero killed count, which is the only signal that positively distinguishes the two. Without it the report says so in as many words and main returns 1. A run with zero kills and a stats file says that too. The test asserting a non-killed run was not called unfinished was codifying the same confusion; it is replaced by three that pin each branch. Caught by Greptile on #37825. * fix(ci): treat stats that count survivors the report never listed as untrusted clean_sweep_is_provable passed on any positive kill count, so a stats file reporting 48 killed and 3 survived, next to a `mutmut results` that listed no survivors, still published a clean sweep. The two sources contradict each other there, and neither one is worth believing. It now requires the stats file to agree that nothing survived, and the report says which disagreement it found. * fix(ci): refuse a clean sweep while mutants never reached the tests A run can end with kills, no survivors, and a pile of mutants marked no tests, skipped, suspicious, timeout or segfault. Those never got put in front of the suite, so "caught every mutation" says more than the run measured. The verdict now names which of them it found and withholds the pass, and the status list those five come from is one constant the summary and the verdict share. * fix(ci): read anything that is not a kill or a survivor as unresolved The unresolved statuses were a list of five, so a run ending in a status the reporter had never met, "check was interrupted by user" among them, still counted as a clean sweep. The rule is now the other way round: killed, survived and total are the keys with a meaning here, and every other non-zero count is a mutant that did not reach the tests, whatever mutmut chose to call it. --- pyproject.toml | 12 +- scripts/mutation_report.py | 106 +++++++++++----- tests/test_litellm/test_mutation_report.py | 139 +++++++++++++++++++++ 3 files changed, 228 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/test_mutation_report.py diff --git a/pyproject.toml b/pyproject.toml index 6e3c181ae1d..57df956bdc9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -341,9 +341,13 @@ filterwarnings = [ paths_to_mutate = [ "litellm/proxy/management_endpoints/", ] +# Only the unit tier that maps to paths_to_mutate. mutmut times and +# coverage-maps this whole set once before mutating, so a tier that needs a +# seeded database (tests/proxy_behavior/) kills the run before it starts, and +# a mutation score is only meaningful against the tests that claim to cover +# the mutated code anyway. tests_dir = [ "tests/test_litellm/proxy/management_endpoints/", - "tests/proxy_behavior/management/", ] also_copy = [ "litellm/", @@ -360,10 +364,16 @@ mutate_only_covered_lines = true # - rerunning a "failed" test on a mutant would mask which mutants are killed # vs. survive, so reruns are wrong for mutation testing regardless. # - xdist is unnecessary inside mutmut (mutmut handles its own parallelism). +# test_saml_sso.py cannot run inside mutmut's mutants/ sandbox: the copied tree +# re-imports cryptography's hash classes under a second identity, so x509 .sign() +# rejects the SHA256 instance the fixture builds with "Algorithm must be a +# registered hash algorithm". Nothing to do with mutation coverage, and one +# erroring test is enough to end the stats phase before any mutant runs. pytest_add_cli_args = [ "-p", "no:retry", "-p", "no:rerunfailures", "-p", "no:xdist", + "--ignore=tests/test_litellm/proxy/management_endpoints/test_saml_sso.py", ] [tool.coverage.run] diff --git a/scripts/mutation_report.py b/scripts/mutation_report.py index a606e3f71cf..e0d4d569484 100644 --- a/scripts/mutation_report.py +++ b/scripts/mutation_report.py @@ -22,6 +22,7 @@ import tomllib from collections import defaultdict from difflib import SequenceMatcher from pathlib import Path +from typing import Final, NamedTuple from textwrap import dedent ROOT = Path(__file__).resolve().parent.parent @@ -33,16 +34,24 @@ def load_mutmut_config() -> dict: return tomllib.load(f)["tool"]["mutmut"] -def get_survivors() -> list[str]: +class MutmutResults(NamedTuple): + survivors: tuple[str, ...] + reported: int + + +def get_survivors() -> MutmutResults: proc = subprocess.run( [*MUTMUT_INVOCATION, "results"], capture_output=True, text=True, check=False ) - survivors = [] - for line in proc.stdout.splitlines(): - m = re.match(r"\s*(\S+):\s*survived\s*$", line) - if m: - survivors.append(m.group(1)) - return survivors + verdicts = tuple( + m.groups() + for line in proc.stdout.splitlines() + if (m := re.match(r"\s*(\S+):\s*(\S.*?)\s*$", line)) + ) + return MutmutResults( + survivors=tuple(name for name, verdict in verdicts if verdict == "survived"), + reported=len(verdicts), + ) def get_mutmut_show(mutant_name: str) -> str: @@ -222,7 +231,52 @@ def render_meta_style_mutant( return "\n".join(out) -def render(config: dict, survivors: list[str], stats: dict | None) -> str: +RESOLVED_KEYS: Final = frozenset({"killed", "survived", "total"}) + + +def unresolved_counts(stats: dict) -> dict[str, int]: + """Every non-zero count that is neither a kill nor a survivor means a mutant did not + reach the tests. Reading it as "anything else" rather than as a list of known statuses + keeps a status this reporter has never met from passing as a clean sweep.""" + return {k: v for k, v in sorted(stats.items()) if k not in RESOLVED_KEYS and isinstance(v, int) and v > 0} + + +def clean_sweep_is_provable(stats: dict | None) -> bool: + """`mutmut results` omits killed mutants, so its silence is equally consistent with a + perfect run and with a run that never started. Only the stats file can tell them apart, + and only when it agrees that nothing survived and every mutant reached the tests.""" + if not stats or stats.get("killed", 0) <= 0 or stats.get("survived", 0) != 0: + return False + return not unresolved_counts(stats) + + +def no_survivors_verdict(results: MutmutResults, stats: dict | None) -> str: + if clean_sweep_is_provable(stats): + return "**No surviving mutants, and the run killed some, so the test suite caught every mutation.**" + if stats and stats.get("survived", 0) > 0: + return ( + f"**mutmut-cicd-stats.json counts {stats['survived']} surviving mutant(s) that " + "`mutmut results` did not list, so the two disagree and neither can be trusted. " + "This is not a passing score.**" + ) + if stats and unresolved_counts(stats): + unresolved = ", ".join(f"{v} {k.replace('_', ' ')}" for k, v in unresolved_counts(stats).items()) + return ( + f"**No survivors, but {unresolved}, so those mutants never reached the tests " + "and the suite was not shown to catch them. This is not a passing score.**" + ) + if stats: + return "**Not one mutant was killed. This is not a passing score.**" + return ( + f"**mutmut-cicd-stats.json is missing and `mutmut results` printed {results.reported} " + "verdict(s), none of them a survivor. Since that command never lists killed mutants, a " + "clean sweep and a run that mutated nothing look identical from here. This is not a " + "passing score.**" + ) + + +def render(config: dict, results: MutmutResults, stats: dict | None) -> str: + survivors = list(results.survivors) by_function: dict[tuple[str, str], list[tuple[str, str]]] = defaultdict(list) for survivor in survivors: module_path, function_name, mutant_num = parse_mutant_name(survivor) @@ -235,17 +289,8 @@ def render(config: dict, survivors: list[str], stats: dict | None) -> str: out.append("## Summary") out.append("") if stats: - total = stats.get("total", 0) or sum( - stats.get(k, 0) - for k in ( - "killed", - "survived", - "no_tests", - "skipped", - "suspicious", - "timeout", - "segfault", - ) + total = stats.get("total", 0) or ( + stats.get("killed", 0) + stats.get("survived", 0) + sum(unresolved_counts(stats).values()) ) killed = stats.get("killed", 0) survived = stats.get("survived", 0) @@ -254,17 +299,15 @@ def render(config: dict, survivors: list[str], stats: dict | None) -> str: out.append(f"- Killed: **{killed}**") out.append(f"- Survived: **{survived}**") out.append(f"- Mutation score: **{score:.1f}%**") - for k in ("no_tests", "skipped", "suspicious", "timeout", "segfault"): - v = stats.get(k, 0) - if v: - out.append(f"- {k.replace('_', ' ').title()}: {v}") + for k, v in unresolved_counts(stats).items(): + out.append(f"- {k.replace('_', ' ').title()}: {v}") else: out.append(f"- Survivors found: **{len(survivors)}**") out.append("- (mutmut-cicd-stats.json not available — full counts unavailable)") out.append("") if not survivors: - out.append("**No surviving mutants — the test suite caught every mutation.**") + out.append(no_survivors_verdict(results, stats)) out.append("") return "\n".join(out) @@ -407,15 +450,22 @@ def main() -> int: except json.JSONDecodeError as exc: print(f"warning: could not parse {stats_file}: {exc}", file=sys.stderr) - survivors = get_survivors() - report = render(config, survivors, stats) + results = get_survivors() + report = render(config, results, stats) out_path = ROOT / "mutation-report.md" out_path.write_text(report) print( - f"Wrote {out_path} ({len(survivors)} survivor" - f"{'s' if len(survivors) != 1 else ''}, {len(report)} chars)" + f"Wrote {out_path} ({len(results.survivors)} survivor" + f"{'s' if len(results.survivors) != 1 else ''}, {len(report)} chars)" ) + if not results.survivors and not clean_sweep_is_provable(stats): + print( + "error: nothing was shown to have been killed, so the report cannot say " + "anything about the suite", + file=sys.stderr, + ) + return 1 return 0 diff --git a/tests/test_litellm/test_mutation_report.py b/tests/test_litellm/test_mutation_report.py new file mode 100644 index 00000000000..60b29ef2628 --- /dev/null +++ b/tests/test_litellm/test_mutation_report.py @@ -0,0 +1,139 @@ +"""Tests for scripts/mutation_report.py. + +The report is the only thing anyone reads after a mutation run, so the one thing it +must never do is describe a run that produced nothing as a run that killed everything. +`render` decides that wording and `get_survivors` supplies the evidence for it, so both +are tested directly. +""" + +import importlib.util +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "scripts" / "mutation_report.py" +_spec = importlib.util.spec_from_file_location("mutation_report", _MODULE_PATH) +report = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = report +_spec.loader.exec_module(report) + +_CONFIG = {"paths_to_mutate": ["litellm/proxy/management_endpoints/"], "tests_dir": ["tests/"]} + + +def test_a_run_that_reported_nothing_is_not_a_clean_sweep(): + rendered = report.render(_CONFIG, report.MutmutResults(survivors=(), reported=0), None) + + assert "not a passing score" in rendered + assert "caught every mutation" not in rendered + + +def test_a_run_that_killed_every_mutant_says_so(): + rendered = report.render( + _CONFIG, report.MutmutResults(survivors=(), reported=0), {"killed": 48, "survived": 0} + ) + + assert "caught every mutation" in rendered + assert "not a passing score" not in rendered + + +def test_stats_counting_survivors_results_never_listed_is_not_a_clean_sweep(): + rendered = report.render( + _CONFIG, report.MutmutResults(survivors=(), reported=0), {"killed": 48, "survived": 3} + ) + + assert "not a passing score" in rendered + assert "caught every mutation" not in rendered + assert "3 surviving mutant(s)" in rendered + + +def test_mutants_that_never_reached_the_tests_are_not_a_clean_sweep(): + rendered = report.render( + _CONFIG, + report.MutmutResults(survivors=(), reported=0), + {"killed": 48, "survived": 0, "no_tests": 4, "timeout": 1}, + ) + + assert "not a passing score" in rendered + assert "caught every mutation" not in rendered + assert "4 no tests" in rendered + assert "1 timeout" in rendered + + +def test_a_status_the_reporter_has_never_met_still_blocks_a_clean_sweep(): + rendered = report.render( + _CONFIG, + report.MutmutResults(survivors=(), reported=0), + {"killed": 48, "survived": 0, "check_was_interrupted_by_user": 2}, + ) + + assert "not a passing score" in rendered + assert "caught every mutation" not in rendered + assert "2 check was interrupted by user" in rendered + + +def test_no_survivors_without_a_kill_is_not_a_clean_sweep(): + rendered = report.render( + _CONFIG, report.MutmutResults(survivors=(), reported=48), {"killed": 0, "survived": 0} + ) + + assert "not a passing score" in rendered + assert "caught every mutation" not in rendered + + +def test_no_survivors_and_no_stats_cannot_claim_a_sweep(): + """`mutmut results` never lists killed mutants, so with the stats file missing an + empty survivor list is equally consistent with a perfect run and a dead one.""" + rendered = report.render(_CONFIG, report.MutmutResults(survivors=(), reported=48), None) + + assert "not a passing score" in rendered + assert "caught every mutation" not in rendered + + +def test_survivors_are_read_out_of_the_verdicts_they_came_with(monkeypatch): + class _Proc: + stdout = ( + "litellm.proxy.management_endpoints.key_management_endpoints.x_1: killed\n" + "litellm.proxy.management_endpoints.key_management_endpoints.x_2: survived\n" + "litellm.proxy.management_endpoints.key_management_endpoints.x_3: no tests\n" + "not a verdict line at all\n" + ) + + monkeypatch.setattr(report.subprocess, "run", lambda *a, **k: _Proc()) + + results = report.get_survivors() + + assert results.survivors == ( + "litellm.proxy.management_endpoints.key_management_endpoints.x_2", + ) + assert results.reported == 3 + + +def test_every_multi_word_verdict_mutmut_can_emit_still_counts(monkeypatch): + class _Proc: + stdout = "".join( + f"litellm.proxy.management_endpoints.key_management_endpoints.x_{i}: {verdict}\n" + for i, verdict in enumerate( + ( + "no tests", + "not checked", + "caught by type check", + "check was interrupted by user", + ) + ) + ) + + monkeypatch.setattr(report.subprocess, "run", lambda *a, **k: _Proc()) + + results = report.get_survivors() + + assert results.survivors == () + assert results.reported == 4 + + +def test_an_empty_mutmut_results_reports_nothing_rather_than_zero_survivors(monkeypatch): + class _Proc: + stdout = "" + + monkeypatch.setattr(report.subprocess, "run", lambda *a, **k: _Proc()) + + assert report.get_survivors() == report.MutmutResults(survivors=(), reported=0) From 4a008b67efb1dd59f8a0d0c22fe0cbc4617f0362 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:16:01 -0700 Subject: [PATCH 27/43] test(bedrock): let monkeypatch own bedrock_request_metadata_fields (#37840) Twenty tests in test_request_metadata.py assigned the global directly and leaned on an autouse fixture to put it back afterwards. monkeypatch.setattr does both jobs at the point of use, so each test now says what it sets and the fixture that existed only to undo them goes away. --- .../llms/bedrock/test_request_metadata.py | 105 ++++++++++-------- 1 file changed, 57 insertions(+), 48 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/test_request_metadata.py b/tests/test_litellm/llms/bedrock/test_request_metadata.py index ad14db5c85f..79b8990a3af 100644 --- a/tests/test_litellm/llms/bedrock/test_request_metadata.py +++ b/tests/test_litellm/llms/bedrock/test_request_metadata.py @@ -36,13 +36,6 @@ ALL_FIELDS = [ IDENTITY = {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform"} -@pytest.fixture(autouse=True) -def reset_setting(): - previous = litellm.bedrock_request_metadata_fields - yield - litellm.bedrock_request_metadata_fields = previous - - def litellm_params(metadata_key, **metadata): return {metadata_key: dict(metadata)} @@ -73,8 +66,8 @@ CONVERSE_DRIVERS = [converse_body, converse_body_async] @pytest.mark.parametrize("setting", [None, []]) -def test_feature_off_by_default_leaves_body_and_headers_untouched(setting): - litellm.bedrock_request_metadata_fields = setting +def test_feature_off_by_default_leaves_body_and_headers_untouched(setting, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", setting) params = litellm_params("metadata", spend_logs_metadata={"team": "x"}, **IDENTITY) assert "requestMetadata" not in converse_body(params) @@ -88,18 +81,18 @@ def test_feature_off_by_default_leaves_body_and_headers_untouched(setting): @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_resolver_reads_both_metadata_variable_names(metadata_key): +def test_resolver_reads_both_metadata_variable_names(metadata_key, monkeypatch: pytest.MonkeyPatch): """`/v1/chat/completions` populates `metadata`; the LITELLM_METADATA_ROUTES populate `litellm_metadata`. Reading only one silently forwards nothing on the other route.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) params = litellm_params(metadata_key, spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY) assert converse_body(params)["requestMetadata"] == {**IDENTITY, "cost_center": "cc-1"} @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key): - litellm.bedrock_request_metadata_fields = ALL_FIELDS +def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) params = litellm_params(metadata_key, **IDENTITY) headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( @@ -112,10 +105,15 @@ def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key) @pytest.mark.parametrize("reverse_client_keys", [False, True]) @pytest.mark.parametrize("field_order", [ALL_FIELDS, list(reversed(ALL_FIELDS))]) @pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"]) -def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, field_order, client_source): +def test_identity_survives_a_caller_filling_every_slot( + reverse_client_keys, + field_order, + client_source, + monkeypatch: pytest.MonkeyPatch, +): """A caller sending 16 keys of its own must not evict the identity the feature exists to produce. Driven over every input ordering so the invariant is not an accident of one.""" - litellm.bedrock_request_metadata_fields = field_order + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", field_order) client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS)] client_pairs = {key: "v" for key in (reversed(client_keys) if reverse_client_keys else client_keys)} if client_source == "spend_logs_metadata": @@ -141,11 +139,14 @@ def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, fiel ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata", "user_api_key_team_alias"], ], ) -def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field_order): +def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot( + field_order, + monkeypatch: pytest.MonkeyPatch, +): """An operator repeating a field in YAML must not inflate the reserved count and shrink the client budget. Asserts the client keys that should have fitted actually reach the wire, since asserting only that identity survives passes with or without the deduplication.""" - litellm.bedrock_request_metadata_fields = field_order + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", field_order) client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - 1)] params = litellm_params("metadata", spend_logs_metadata={key: "v" for key in client_keys}, **IDENTITY) @@ -162,11 +163,15 @@ def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field "forged_key", ["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"], ) -def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, client_source): +def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( + forged_key, + client_source, + monkeypatch: pytest.MonkeyPatch, +): """`user_api_key_org_alias` and `user_api_key_hash` are names the proxy does not set here, so an exact-key reservation would let the forged value through under a name that reads as proxy-authoritative in the AWS billing record.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) forged = {forged_key: "attacker-controlled"} if client_source == "spend_logs_metadata": params, optional_params = litellm_params("metadata", spend_logs_metadata=forged, **IDENTITY), {} @@ -179,10 +184,10 @@ def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, clien assert "attacker-controlled" not in resolved.values() -def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(): +def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(monkeypatch: pytest.MonkeyPatch): """A team alias with an apostrophe must not turn a working request into a 400 the moment an operator flips the setting on.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) params = litellm_params( "metadata", user_api_key_alias="prod-key", @@ -196,8 +201,8 @@ def test_identity_violating_the_character_class_is_dropped_and_the_request_succe assert body["messages"] -def test_caller_supplied_violation_still_raises_bad_request(): - litellm.bedrock_request_metadata_fields = ALL_FIELDS +def test_caller_supplied_violation_still_raises_bad_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) with pytest.raises(litellm.exceptions.BadRequestError): converse_body( @@ -206,34 +211,34 @@ def test_caller_supplied_violation_still_raises_bad_request(): ) -def test_non_string_and_absent_identity_values_are_dropped(): - litellm.bedrock_request_metadata_fields = ALL_FIELDS + ["user_api_key_spend"] +def test_non_string_and_absent_identity_values_are_dropped(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS + ["user_api_key_spend"]) params = litellm_params("metadata", user_api_key_alias="prod-key", user_api_key_spend=1.25) assert converse_body(params)["requestMetadata"] == {"user_api_key_alias": "prod-key"} -def test_email_is_separately_opt_in(): +def test_email_is_separately_opt_in(monkeypatch: pytest.MonkeyPatch): """PII crossing into CloudTrail only when the operator names the field.""" identity_with_email = {**IDENTITY, "user_api_key_user_email": "owner@example.com"} - litellm.bedrock_request_metadata_fields = ["user_api_key_alias", "user_api_key_team_alias"] + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_alias", "user_api_key_team_alias"]) assert ( "user_api_key_user_email" not in converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] ) - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) assert converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] == identity_with_email -def test_resolver_returns_none_when_nothing_survives(): - litellm.bedrock_request_metadata_fields = ALL_FIELDS +def test_resolver_returns_none_when_nothing_survives(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) assert resolve_bedrock_request_metadata(litellm_params=None) is None assert resolve_bedrock_request_metadata(litellm_params={"metadata": {"unrelated": "x"}}) is None -def test_invoke_header_is_json_encoded_and_signed(): - litellm.bedrock_request_metadata_fields = ALL_FIELDS +def test_invoke_header_is_json_encoded_and_signed(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) params = litellm_params("metadata", spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY) headers = AmazonInvokeConfig().validate_environment( @@ -250,10 +255,10 @@ def test_invoke_header_is_json_encoded_and_signed(): assert "anthropic-version" not in signed -def test_a_caller_supplied_guardrail_header_still_wins(): +def test_a_caller_supplied_guardrail_header_still_wins(monkeypatch: pytest.MonkeyPatch): """The no-displace rule is deliberate for the guardrail headers and must survive the request-metadata header becoming proxy-owned.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) headers = AmazonInvokeConfig().validate_environment( headers={"X-Amzn-Bedrock-GuardrailIdentifier": "caller-set"}, @@ -318,10 +323,10 @@ def metadata_header_values(headers): return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()] -def test_converse_still_sets_the_bearer_authorization_header(): +def test_converse_still_sets_the_bearer_authorization_header(monkeypatch: pytest.MonkeyPatch): """Converse owns the metadata header now, and that must not disturb the api_key path its validate_environment existed for. Closing the forgery hole cannot break authentication.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) headers = AmazonConverseConfig().validate_environment( headers={}, @@ -341,11 +346,11 @@ def test_converse_still_sets_the_bearer_authorization_header(): "caller_header_name", [BEDROCK_REQUEST_METADATA_HEADER, BEDROCK_REQUEST_METADATA_HEADER.lower(), "x-AMZN-bedrock-Request-METADATA"], ) -def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name): +def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name, monkeypatch: pytest.MonkeyPatch): """`extra_headers` puts caller-supplied names into the same dict the proxy merges into, so a deferring merge would sign the caller's forged identity into the AWS billing record. Every spelling must lose, or a second variant is left for the transport to choose between.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) headers = driver({caller_header_name: FORGED}, litellm_params("metadata", **IDENTITY)) @@ -355,11 +360,11 @@ def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header @pytest.mark.parametrize("driver", HEADER_DRIVERS) -def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver): +def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver, monkeypatch: pytest.MonkeyPatch): """Forwarding enabled but nothing resolvable, which a caller can arrange by supplying values that all fail Bedrock's rules. Owned-but-empty must mean no header on the wire, never a fallback to the caller's.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) unresolvable = litellm_params("metadata", user_api_key_alias="O'Brien's key", user_api_key_team_alias="x" * 300) headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, unresolvable) @@ -373,11 +378,15 @@ def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(drive "forged_key", ["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"], ) -def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing(forged_key, driver): +def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing( + forged_key, + driver, + monkeypatch: pytest.MonkeyPatch, +): """The Converse body has the same fail-open shape as the header: with forwarding on and nothing resolvable, leaving the caller's `requestMetadata` in place would keep their reserved-prefix keys on the wire. Owned-but-empty must remove the field outright.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) body = driver(litellm_params("metadata"), {"requestMetadata": {forged_key: "FORGED"}}) @@ -386,10 +395,10 @@ def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothin @pytest.mark.parametrize("driver", CONVERSE_DRIVERS) -def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver): +def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver, monkeypatch: pytest.MonkeyPatch): """Removing the field must be scoped to the reserved keys being the only thing left, not a blanket drop of the caller's own attribution pairs.""" - litellm.bedrock_request_metadata_fields = ALL_FIELDS + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) body = driver( litellm_params("metadata"), @@ -400,10 +409,10 @@ def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(dr @pytest.mark.parametrize("driver", CONVERSE_DRIVERS) -def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver): +def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver, monkeypatch: pytest.MonkeyPatch): """With the feature off the proxy does not own the field, so the pre-existing pass-through behaviour for a caller-supplied `requestMetadata` must be unchanged.""" - litellm.bedrock_request_metadata_fields = None + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) caller_supplied = {"user_api_key_team_alias": "caller-set", "cost_center": "cc-9"} body = driver(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_supplied}) @@ -412,10 +421,10 @@ def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver): @pytest.mark.parametrize("driver", HEADER_DRIVERS) -def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver): +def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver, monkeypatch: pytest.MonkeyPatch): """The proxy only claims the name when the operator turned forwarding on; with the feature off this is an ordinary passthrough header and stripping it would be a regression.""" - litellm.bedrock_request_metadata_fields = None + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, litellm_params("metadata", **IDENTITY)) From 49da936efb8f096c424d1e046732c4ca936042b0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:16:10 -0700 Subject: [PATCH 28/43] test(audit-logs): let monkeypatch own the audit log and s3 callback globals (#37842) Fifteen tests assigned litellm.audit_log_callbacks, s3_callback_params or s3_audit_callback_params directly and leaned on two autouse fixtures to put them back. monkeypatch.setattr does that at the point of use, so each test now says what it sets, including the one that swaps the value mid-test to prove the cache does not serve the stale params. The fixtures keep only the work monkeypatch cannot do: the per-test empty callback list, and clearing the logger and audit caches around each test. --- .../test_audit_log_callbacks.py | 83 +++++++++---------- 1 file changed, 40 insertions(+), 43 deletions(-) diff --git a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py index 2d54d249713..b1d111bf1f9 100644 --- a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py +++ b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py @@ -25,12 +25,9 @@ from litellm.types.utils import StandardAuditLogPayload @pytest.fixture(autouse=True) -def reset_audit_log_callbacks(): - """Reset audit_log_callbacks before and after each test.""" - original = litellm.audit_log_callbacks - litellm.audit_log_callbacks = [] - yield - litellm.audit_log_callbacks = original +def reset_audit_log_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + """Every test starts with no audit log callbacks registered.""" + monkeypatch.setattr(litellm, "audit_log_callbacks", []) def _make_audit_log( @@ -115,10 +112,10 @@ class TestBuildAuditLogPayload: class TestDispatchAuditLogToCallbacks: @pytest.mark.asyncio - async def test_dispatches_to_custom_logger_instance(self): + async def test_dispatches_to_custom_logger_instance(self, monkeypatch: pytest.MonkeyPatch): mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) audit_log = _make_audit_log() await _dispatch_audit_log_to_callbacks(audit_log) @@ -132,18 +129,18 @@ class TestDispatchAuditLogToCallbacks: assert payload["action"] == "created" @pytest.mark.asyncio - async def test_no_dispatch_when_callbacks_empty(self): - litellm.audit_log_callbacks = [] + async def test_no_dispatch_when_callbacks_empty(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "audit_log_callbacks", []) audit_log = _make_audit_log() # Should return immediately without error await _dispatch_audit_log_to_callbacks(audit_log) @pytest.mark.asyncio - async def test_resolves_string_callback(self): + async def test_resolves_string_callback(self, monkeypatch: pytest.MonkeyPatch): mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = ["s3_v2"] + monkeypatch.setattr(litellm, "audit_log_callbacks", ["s3_v2"]) with patch( "litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback", @@ -156,13 +153,13 @@ class TestDispatchAuditLogToCallbacks: mock_logger.async_log_audit_log_event.assert_called_once() @pytest.mark.asyncio - async def test_nonblocking_on_callback_failure(self): + async def test_nonblocking_on_callback_failure(self, monkeypatch: pytest.MonkeyPatch): """Callback errors should not propagate.""" mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock( side_effect=RuntimeError("boom") ) - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) audit_log = _make_audit_log() # Should not raise @@ -170,8 +167,8 @@ class TestDispatchAuditLogToCallbacks: await asyncio.sleep(0.1) @pytest.mark.asyncio - async def test_skips_unresolvable_string_callback(self): - litellm.audit_log_callbacks = ["nonexistent_callback"] + async def test_skips_unresolvable_string_callback(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "audit_log_callbacks", ["nonexistent_callback"]) with patch( "litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback", @@ -184,10 +181,10 @@ class TestDispatchAuditLogToCallbacks: class TestCreateAuditLogForUpdateWithCallbacks: @pytest.mark.asyncio - async def test_dispatches_to_callbacks_after_db_write(self): + async def test_dispatches_to_callbacks_after_db_write(self, monkeypatch: pytest.MonkeyPatch): mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) with ( patch("litellm.proxy.proxy_server.premium_user", True), @@ -206,10 +203,10 @@ class TestCreateAuditLogForUpdateWithCallbacks: mock_logger.async_log_audit_log_event.assert_called_once() @pytest.mark.asyncio - async def test_no_dispatch_when_not_premium(self): + async def test_no_dispatch_when_not_premium(self, monkeypatch: pytest.MonkeyPatch): mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) with ( patch("litellm.proxy.proxy_server.premium_user", False), @@ -224,10 +221,10 @@ class TestCreateAuditLogForUpdateWithCallbacks: mock_prisma.db.litellm_auditlog.create.assert_not_called() @pytest.mark.asyncio - async def test_no_dispatch_when_store_audit_logs_false(self): + async def test_no_dispatch_when_store_audit_logs_false(self, monkeypatch: pytest.MonkeyPatch): mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) with patch("litellm.store_audit_logs", False): audit_log = _make_audit_log() @@ -237,11 +234,11 @@ class TestCreateAuditLogForUpdateWithCallbacks: mock_logger.async_log_audit_log_event.assert_not_called() @pytest.mark.asyncio - async def test_dispatches_even_when_prisma_client_is_none(self): + async def test_dispatches_even_when_prisma_client_is_none(self, monkeypatch: pytest.MonkeyPatch): """Callbacks should fire even if DB is unavailable.""" mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) with ( patch("litellm.proxy.proxy_server.premium_user", True), @@ -256,11 +253,11 @@ class TestCreateAuditLogForUpdateWithCallbacks: mock_logger.async_log_audit_log_event.assert_called_once() @pytest.mark.asyncio - async def test_dispatches_even_when_db_write_fails(self): + async def test_dispatches_even_when_db_write_fails(self, monkeypatch: pytest.MonkeyPatch): """Callbacks should fire even if the DB write raises.""" mock_logger = MagicMock(spec=CustomLogger) mock_logger.async_log_audit_log_event = AsyncMock() - litellm.audit_log_callbacks = [mock_logger] + monkeypatch.setattr(litellm, "audit_log_callbacks", [mock_logger]) with ( patch("litellm.proxy.proxy_server.premium_user", True), @@ -384,21 +381,21 @@ class TestS3AuditCallbackParamsDecoupling: S3Logger instance, distinct from the singleton serving normal logs.""" @pytest.fixture(autouse=True) - def _isolate_caches_and_globals(self): + def _isolate_caches_and_globals(self, monkeypatch: pytest.MonkeyPatch): from litellm.litellm_core_utils import litellm_logging as ll_logging from litellm.proxy.management_helpers import audit_logs as ll_audit_logs - original_s3 = litellm.s3_callback_params - original_audit = getattr(litellm, "s3_audit_callback_params", None) + monkeypatch.setattr(litellm, "s3_callback_params", litellm.s3_callback_params) + monkeypatch.setattr( + litellm, "s3_audit_callback_params", getattr(litellm, "s3_audit_callback_params", None) + ) ll_audit_logs._audit_log_callback_cache.clear() ll_logging._in_memory_loggers.clear() yield - litellm.s3_callback_params = original_s3 - litellm.s3_audit_callback_params = original_audit ll_audit_logs._audit_log_callback_cache.clear() ll_logging._in_memory_loggers.clear() - def test_opt_in_constructs_separate_instance_with_audit_config(self): + def test_opt_in_constructs_separate_instance_with_audit_config(self, monkeypatch: pytest.MonkeyPatch): """Audit config set → audit resolver returns a fresh S3Logger pointing at the audit bucket, distinct from the normal-log singleton.""" from litellm.integrations.s3_v2 import S3Logger @@ -409,8 +406,8 @@ class TestS3AuditCallbackParamsDecoupling: _resolve_audit_log_callback, ) - litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"} - litellm.s3_audit_callback_params = {"s3_bucket_name": "audit-bucket"} + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"}) + monkeypatch.setattr(litellm, "s3_audit_callback_params", {"s3_bucket_name": "audit-bucket"}) with patch("asyncio.create_task"): audit_instance = _resolve_audit_log_callback("s3_v2") @@ -426,7 +423,7 @@ class TestS3AuditCallbackParamsDecoupling: assert audit_instance.s3_bucket_name == "audit-bucket" assert normal_instance.s3_bucket_name == "normal-bucket" - def test_opt_out_preserves_singleton_behavior(self): + def test_opt_out_preserves_singleton_behavior(self, monkeypatch: pytest.MonkeyPatch): """No `s3_audit_callback_params` → audit and normal share the singleton (existing behavior, regression guard).""" from litellm.integrations.s3_v2 import S3Logger @@ -437,8 +434,8 @@ class TestS3AuditCallbackParamsDecoupling: _resolve_audit_log_callback, ) - litellm.s3_callback_params = {"s3_bucket_name": "shared-bucket"} - litellm.s3_audit_callback_params = None + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "shared-bucket"}) + monkeypatch.setattr(litellm, "s3_audit_callback_params", None) with patch("asyncio.create_task"): normal_instance = _init_custom_logger_compatible_class( @@ -452,7 +449,7 @@ class TestS3AuditCallbackParamsDecoupling: assert id(audit_instance) == id(normal_instance) assert audit_instance.s3_bucket_name == "shared-bucket" - def test_empty_dict_opts_in(self): + def test_empty_dict_opts_in(self, monkeypatch: pytest.MonkeyPatch): """`s3_audit_callback_params = {}` is opt-in (truthy-by-presence) and produces a separate instance with no bucket configured (env/IAM-only).""" from litellm.integrations.s3_v2 import S3Logger @@ -463,8 +460,8 @@ class TestS3AuditCallbackParamsDecoupling: _resolve_audit_log_callback, ) - litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"} - litellm.s3_audit_callback_params = {} + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"}) + monkeypatch.setattr(litellm, "s3_audit_callback_params", {}) with patch("asyncio.create_task"): audit_instance = _resolve_audit_log_callback("s3_v2") @@ -478,7 +475,7 @@ class TestS3AuditCallbackParamsDecoupling: assert audit_instance.s3_bucket_name is None assert normal_instance.s3_bucket_name == "normal-bucket" - def test_reset_audit_log_callback_cache_clears_audit_instance(self): + def test_reset_audit_log_callback_cache_clears_audit_instance(self, monkeypatch: pytest.MonkeyPatch): """`reset_audit_log_callback_cache()` must drop the cached audit instance so a config reload picks up the new params.""" from litellm.proxy.management_helpers.audit_logs import ( @@ -487,7 +484,7 @@ class TestS3AuditCallbackParamsDecoupling: reset_audit_log_callback_cache, ) - litellm.s3_audit_callback_params = {"s3_bucket_name": "first"} + monkeypatch.setattr(litellm, "s3_audit_callback_params", {"s3_bucket_name": "first"}) with patch("asyncio.create_task"): first = _resolve_audit_log_callback("s3_v2") assert first is not None and "s3_v2" in _audit_log_callback_cache @@ -495,7 +492,7 @@ class TestS3AuditCallbackParamsDecoupling: reset_audit_log_callback_cache() assert "s3_v2" not in _audit_log_callback_cache - litellm.s3_audit_callback_params = {"s3_bucket_name": "second"} + monkeypatch.setattr(litellm, "s3_audit_callback_params", {"s3_bucket_name": "second"}) second = _resolve_audit_log_callback("s3_v2") assert second is not None assert id(second) != id(first) From 693797420df83ba58839bead1b44c9f76cbf24fb Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:28:37 -0700 Subject: [PATCH 29/43] test: unwind environment writes in tests/test_litellm with monkeypatch (#37806) * test: use monkeypatch.setenv for env writes in tests/test_litellm `os.environ["X"] = v` inside a test leaks the value into every test that runs after it in the same worker, so ordering decides the result. 262 of those writes across 40 files now go through pytest's `monkeypatch` fixture, which restores the previous value at teardown. The rewrite skips any test that a mock.patch-family decorator wraps, any test with defaulted positional parameters, any test whose own name is called directly elsewhere, and rebinds nothing inside nested defs, because in each of those cases appending a fixture parameter changes what pytest or mock binds. Ratchets the TQ004 ceiling from 768 to 506. * fix(test): delete the key through monkeypatch instead of popping it first Five tests popped a key straight out of `os.environ`, ran, then restored it with `monkeypatch.setenv`. By the time monkeypatch saw the name it was already gone, so it recorded "absent" as the value to go back to and deleted the key at teardown. On a worker that inherited a real `RESEND_API_KEY`, `SENDGRID_API_KEY`, `UI_PASSWORD`, `LITELLM_SALT_KEY` or `OPENAI_API_KEY`, every test after the first one ran without it. `monkeypatch.delenv(..., raising=False)` removes the key and restores whatever was there, so the try/finally the manual restore needed goes with it. * chore(test): leave the two cost-calc files to the PR that rewrites them fully Both files are also in #37815, which converts the module-global writes as well as the env writes and folds them into one fixture. Two PRs rewriting the same lines differently is a conflict nobody benefits from resolving, so this one drops back to staging on those two and keeps the other 39. TQ004 clears 200 here instead of 275; the rest moves with #37815. --- test-quality-budget.json | 2 +- ...responses_transformation_transformation.py | 6 +- .../test_container_transformation.py | 4 +- .../send_emails/test_resend_email.py | 65 +++++----- .../send_emails/test_sendgrid_email.py | 24 ++-- .../gcs_bucket/test_gcs_bucket_base.py | 4 +- .../integrations/test_openmeter.py | 12 +- .../llm_cost_calc/test_guardrail_cost.py | 4 +- .../test_tool_call_cost_tracking.py | 4 +- ...llm_core_utils_prompt_templates_factory.py | 12 +- .../test_litellm_logging.py | 18 +-- ...erimental_pass_through_messages_handler.py | 4 +- .../test_responses_adapters_transformation.py | 4 +- .../llms/apiserpent/test_apiserpent_search.py | 8 +- .../test_mai_image_generation.py | 12 +- .../chat/test_converse_transformation.py | 42 +++---- .../test_agentcore_search_transformation.py | 44 +++---- .../llms/bedrock/test_bedrock_ssl_verify.py | 16 +-- tests/test_litellm/llms/crusoe/test_crusoe.py | 6 +- .../test_datarobot_chat_transformation.py | 8 +- .../test_deepinfra_chat_transformation.py | 4 +- .../llms/gemini/test_cost_calculator.py | 24 ++-- .../test_inception_chat_transformation.py | 8 +- ...est_inception_completion_transformation.py | 4 +- ...tex_ai_image_generation_cost_calculator.py | 8 +- .../llms/zai/test_zai_provider.py | 20 +-- .../proxy/auth/test_login_utils.py | 60 +++++---- .../guardrail_hooks/test_deepkeep.py | 16 +-- .../guardrail_hooks/test_hiddenlayer.py | 80 ++++++------ .../guardrails/guardrail_hooks/test_lasso.py | 4 +- .../guardrails/guardrail_hooks/test_onyx.py | 118 +++++++++--------- .../guardrail_hooks/test_repelloai.py | 22 ++-- .../test_prompt_security_guardrails.py | 68 +++++----- .../hooks/test_dynamic_rate_limiter_v3.py | 60 ++++----- .../proxy/hooks/test_rate_limiter_toctou.py | 12 +- .../test_add_deployment_no_master_key.py | 89 +++++++------ .../test_count_tokens_public_api.py | 24 ++-- .../test_register_model_custom_pricing.py | 4 +- tests/test_litellm/test_utils.py | 22 ++-- 39 files changed, 462 insertions(+), 484 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 55c14ec2680..143378efec7 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -9,7 +9,7 @@ "limit": 1078 }, "TQ004": { - "limit": 757 + "limit": 557 }, "TQ005": { "limit": 2810 diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 382b41807d4..858ca482eb7 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1508,7 +1508,7 @@ def test_multiple_tool_calls_in_single_choice(): print("✓ Multiple tool calls are correctly grouped in a single choice") -def test_map_reasoning_effort_adds_summary_detailed(): +def test_map_reasoning_effort_adds_summary_detailed(monkeypatch): """ Test that _map_reasoning_effort behavior with reasoning_auto_summary flag. @@ -1571,7 +1571,7 @@ def test_map_reasoning_effort_adds_summary_detailed(): # Test 3: With env var enabled (flag disabled) - summary IS added litellm.reasoning_auto_summary = False - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") result = handler._map_reasoning_effort("high") assert ( @@ -1603,7 +1603,7 @@ def test_map_reasoning_effort_adds_summary_detailed(): # Restore original values litellm.reasoning_auto_summary = original_flag if original_env is not None: - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = original_env + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", original_env) elif "LITELLM_REASONING_AUTO_SUMMARY" in os.environ: del os.environ["LITELLM_REASONING_AUTO_SUMMARY"] diff --git a/tests/test_litellm/containers/test_container_transformation.py b/tests/test_litellm/containers/test_container_transformation.py index 555fe7773f0..f0432816fce 100644 --- a/tests/test_litellm/containers/test_container_transformation.py +++ b/tests/test_litellm/containers/test_container_transformation.py @@ -341,10 +341,10 @@ class TestOpenAIContainerTransformation: assert data["expires_after"] is None assert data["file_ids"] is None - def test_container_create_response_includes_cost(self): + def test_container_create_response_includes_cost(self, monkeypatch): """Test that container create response includes code interpreter cost calculation.""" # Force use of local model cost map for CI/CD consistency - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py index 88cc2275ae2..fbfd609cca6 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py @@ -88,49 +88,44 @@ async def test_send_email_success(mock_env_vars): @pytest.mark.asyncio -async def test_send_email_missing_api_key(): +async def test_send_email_missing_api_key(monkeypatch): # Remove the API key from environment before initializing logger - original_key = os.environ.pop("RESEND_API_KEY", None) + monkeypatch.delenv("RESEND_API_KEY", raising=False) - try: - # Initialize the logger after removing the API key - logger = ResendEmailLogger() + # Initialize the logger after removing the API key + logger = ResendEmailLogger() - # Test data - from_email = "test@example.com" - to_email = ["recipient@example.com"] - subject = "Test Subject" - html_body = "

Test email body

" + # Test data + from_email = "test@example.com" + to_email = ["recipient@example.com"] + subject = "Test Subject" + html_body = "

Test email body

" - # Create mock HTTP client and inject it directly into the logger - # This ensures the mock is used regardless of any caching issues - mock_response = mock.Mock(spec=Response) - mock_response.raise_for_status.return_value = None - mock_response.status_code = 200 - mock_response.json.return_value = {"id": "test_email_id"} + # Create mock HTTP client and inject it directly into the logger + # This ensures the mock is used regardless of any caching issues + mock_response = mock.Mock(spec=Response) + mock_response.raise_for_status.return_value = None + mock_response.status_code = 200 + mock_response.json.return_value = {"id": "test_email_id"} - mock_async_client = mock.AsyncMock() - mock_async_client.post.return_value = mock_response + mock_async_client = mock.AsyncMock() + mock_async_client.post.return_value = mock_response - # Directly inject the mock client to bypass any caching - logger.async_httpx_client = mock_async_client + # Directly inject the mock client to bypass any caching + logger.async_httpx_client = mock_async_client - # Send email - await logger.send_email( - from_email=from_email, - to_email=to_email, - subject=subject, - html_body=html_body, - ) + # Send email + await logger.send_email( + from_email=from_email, + to_email=to_email, + subject=subject, + html_body=html_body, + ) - # Verify the HTTP client was called with None as the API key - mock_async_client.post.assert_called_once() - call_args = mock_async_client.post.call_args - assert call_args[1]["headers"] == {"Authorization": "Bearer None"} - finally: - # Restore the original key if it existed - if original_key is not None: - os.environ["RESEND_API_KEY"] = original_key + # Verify the HTTP client was called with None as the API key + mock_async_client.post.assert_called_once() + call_args = mock_async_client.post.call_args + assert call_args[1]["headers"] == {"Authorization": "Bearer None"} @pytest.mark.asyncio diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index 5fe4b217e4f..b7fcce8dbf3 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -98,22 +98,18 @@ async def test_send_email_success(mock_env_vars, mock_async_client): @pytest.mark.asyncio -async def test_send_email_missing_api_key(): - original_key = os.environ.pop("SENDGRID_API_KEY", None) +async def test_send_email_missing_api_key(monkeypatch): + monkeypatch.delenv("SENDGRID_API_KEY", raising=False) - try: - logger = SendGridEmailLogger() + logger = SendGridEmailLogger() - with pytest.raises(ValueError, match='SENDGRID_API_KEY is not set'): - await logger.send_email( - from_email="test@example.com", - to_email=["recipient@example.com"], - subject="Test Subject", - html_body="

Test email body

", - ) - finally: - if original_key is not None: - os.environ["SENDGRID_API_KEY"] = original_key + with pytest.raises(ValueError, match='SENDGRID_API_KEY is not set'): + await logger.send_email( + from_email="test@example.com", + to_email=["recipient@example.com"], + subject="Test Subject", + html_body="

Test email body

", + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index a4e16500aee..8d662311da1 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -8,10 +8,10 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase class TestGCSBucketBase: - def test_construct_request_headers_with_project_id(self): + def test_construct_request_headers_with_project_id(self, monkeypatch): """Test that construct_request_headers correctly uses project_id if passed from env""" test_project_id = "test-project" - os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = test_project_id + monkeypatch.setenv("GOOGLE_SECRET_MANAGER_PROJECT_ID", test_project_id) try: # Create handler diff --git a/tests/test_litellm/integrations/test_openmeter.py b/tests/test_litellm/integrations/test_openmeter.py index b9da99b6fa9..2d09e1572db 100644 --- a/tests/test_litellm/integrations/test_openmeter.py +++ b/tests/test_litellm/integrations/test_openmeter.py @@ -236,9 +236,9 @@ class TestOpenMeterIntegration: assert result["data"]["completion_tokens"] == 8 assert result["data"]["total_tokens"] == 23 - def test_custom_event_type(self): + def test_custom_event_type(self, monkeypatch): """Test that custom event type is used when set""" - os.environ["OPENMETER_EVENT_TYPE"] = "custom_event_type" + monkeypatch.setenv("OPENMETER_EVENT_TYPE", "custom_event_type") logger = OpenMeterLogger() @@ -374,10 +374,10 @@ class TestOpenMeterIntegration: assert isinstance(result["subject"], str) assert result["subject"] == "12345" - def test_common_logic_trust_request_user_false_ignores_request_user(self): + def test_common_logic_trust_request_user_false_ignores_request_user(self, monkeypatch): """OPENMETER_TRUST_REQUEST_USER=false makes the key-bound user_id win over a request-supplied `user` (forge-attribution mitigation).""" - os.environ["OPENMETER_TRUST_REQUEST_USER"] = "false" + monkeypatch.setenv("OPENMETER_TRUST_REQUEST_USER", "false") logger = OpenMeterLogger() kwargs = { @@ -400,11 +400,11 @@ class TestOpenMeterIntegration: assert result["subject"] == "real-tenant-id" assert result["subject"] != "forged-by-client" - def test_common_logic_trust_request_user_false_still_raises_without_key_user(self): + def test_common_logic_trust_request_user_false_still_raises_without_key_user(self, monkeypatch): """OPENMETER_TRUST_REQUEST_USER=false still raises when no user_api_key_user_id is available — the request `user` is not a fallback in this mode.""" - os.environ["OPENMETER_TRUST_REQUEST_USER"] = "false" + monkeypatch.setenv("OPENMETER_TRUST_REQUEST_USER", "false") logger = OpenMeterLogger() kwargs = { diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index cf36a2b9b25..052c08a86b5 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -56,8 +56,8 @@ def test_bedrock_guardrail_cost_no_pricing_entry(monkeypatch): assert bedrock_guardrail_cost(usage_units={"contentPolicyUnits": 1}, aws_region_name="us-east-1") == 0.0 -def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_shipped_bedrock_guardrail_prices_match_aws_pricing_page(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") assert litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] == { "automatedReasoningPolicyUnits": 0.00017, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 2f32145580d..a5128228742 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -377,12 +377,12 @@ def test_get_cost_for_vertex_ai_gemini_web_search(model, custom_llm_provider): assert cost == 0.035, f"Expected $0.035 grounding cost, got ${cost}" -def test_azure_assistant_features_integrated_cost_tracking(): +def test_azure_assistant_features_integrated_cost_tracking(monkeypatch): """ Test integrated cost tracking for Azure assistant features. """ # Force use of local model cost map for CI/CD consistency - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure/gpt-4o" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index a10dc46eb42..3a7e06d085a 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2844,7 +2844,7 @@ def test_anthropic_messages_pt_file_block_preserves_cache_control(): assert text_block["cache_control"]["type"] == "ephemeral" -def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(): +def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(monkeypatch): """ Tools with cache_control ttl should preserve the ttl in the cachePoint block for Claude 4.5+ models on Bedrock, matching the behavior of system @@ -2867,7 +2867,7 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: tool_with_1h = { @@ -2927,10 +2927,10 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) -def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): +def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(monkeypatch): """ End-to-end: _bedrock_tools_pt should produce cachePoint blocks with ttl for Claude 4.5+ models when tools have cache_control with ttl. @@ -2944,7 +2944,7 @@ def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: tools = [ @@ -2980,7 +2980,7 @@ def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 82de634b488..a3dcdaf1737 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -64,7 +64,7 @@ def test_post_call_serializes_dict_with_datetime(logging_obj): assert "2026-05-11" in serialized -def test_sentry_sample_rate(): +def test_sentry_sample_rate(monkeypatch): existing_sample_rate = os.getenv("SENTRY_API_SAMPLE_RATE") try: # test with default value by removing the environment variable @@ -76,7 +76,7 @@ def test_sentry_sample_rate(): assert os.environ.get("SENTRY_API_SAMPLE_RATE") == "1.0" # test with custom value - os.environ["SENTRY_API_SAMPLE_RATE"] = "0.5" + monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", "0.5") set_callbacks(["sentry"]) # Check if the custom sample rate is set correctly @@ -86,13 +86,13 @@ def test_sentry_sample_rate(): finally: # Restore the original environment variable if existing_sample_rate: - os.environ["SENTRY_API_SAMPLE_RATE"] = existing_sample_rate + monkeypatch.setenv("SENTRY_API_SAMPLE_RATE", existing_sample_rate) else: if "SENTRY_API_SAMPLE_RATE" in os.environ: del os.environ["SENTRY_API_SAMPLE_RATE"] -def test_sentry_environment(): +def test_sentry_environment(monkeypatch): """Test that SENTRY_ENVIRONMENT is properly handled during Sentry initialization""" existing_environment = os.getenv("SENTRY_ENVIRONMENT") existing_dsn = os.getenv("SENTRY_DSN") @@ -115,7 +115,7 @@ def test_sentry_environment(): try: # Set a mock DSN to allow Sentry initialization - os.environ["SENTRY_DSN"] = "https://test@sentry.io/123456" + monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") # Test with default value (no environment set) if existing_environment: @@ -129,7 +129,7 @@ def test_sentry_environment(): assert call_kwargs["environment"] == "production" # Test with custom environment value - os.environ["SENTRY_ENVIRONMENT"] = "development" + monkeypatch.setenv("SENTRY_ENVIRONMENT", "development") mock_init.reset_mock() set_callbacks(["sentry"]) @@ -139,7 +139,7 @@ def test_sentry_environment(): assert call_kwargs["environment"] == "development" # Test with staging environment - os.environ["SENTRY_ENVIRONMENT"] = "staging" + monkeypatch.setenv("SENTRY_ENVIRONMENT", "staging") mock_init.reset_mock() set_callbacks(["sentry"]) @@ -154,13 +154,13 @@ def test_sentry_environment(): finally: # Restore the original environment variables if existing_environment: - os.environ["SENTRY_ENVIRONMENT"] = existing_environment + monkeypatch.setenv("SENTRY_ENVIRONMENT", existing_environment) else: if "SENTRY_ENVIRONMENT" in os.environ: del os.environ["SENTRY_ENVIRONMENT"] if existing_dsn: - os.environ["SENTRY_DSN"] = existing_dsn + monkeypatch.setenv("SENTRY_DSN", existing_dsn) else: if "SENTRY_DSN" in os.environ: del os.environ["SENTRY_DSN"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 15ac73ed352..1ce683d76fc 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -528,7 +528,7 @@ class TestThinkingSummaryPreservation: finally: litellm.reasoning_auto_summary = original - def test_summary_added_when_env_var_set(self): + def test_summary_added_when_env_var_set(self, monkeypatch): """When LITELLM_REASONING_AUTO_SUMMARY env var is true, summary is added.""" import litellm from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( @@ -538,7 +538,7 @@ class TestThinkingSummaryPreservation: original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = False - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") completion_kwargs = { "model": "responses/gpt-5.2", "custom_llm_provider": "openai", diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 03cbfbb8609..17bab9bf6a5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -845,14 +845,14 @@ class TestTranslateThinkingToReasoning: finally: litellm.reasoning_auto_summary = original - def test_summary_added_when_env_var_set(self): + def test_summary_added_when_env_var_set(self, monkeypatch): """When LITELLM_REASONING_AUTO_SUMMARY env var is true, summary is included.""" import litellm original = litellm.reasoning_auto_summary try: litellm.reasoning_auto_summary = False - os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") result = _ADAPTER.translate_thinking_to_reasoning( { "type": "enabled", diff --git a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py index bc26268ee92..c925bd7de45 100644 --- a/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py +++ b/tests/test_litellm/llms/apiserpent/test_apiserpent_search.py @@ -239,8 +239,8 @@ class TestAPISerpentSearchIntegration: return mock_response @pytest.mark.asyncio - async def test_asearch_quick_default(self): - os.environ["APISERPENT_API_KEY"] = "test-api-key" + async def test_asearch_quick_default(self, monkeypatch): + monkeypatch.setenv("APISERPENT_API_KEY", "test-api-key") with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", new_callable=AsyncMock, @@ -269,8 +269,8 @@ class TestAPISerpentSearchIntegration: assert response.results[0].title == "Test Result" @pytest.mark.asyncio - async def test_asearch_deep(self): - os.environ["APISERPENT_API_KEY"] = "test-api-key" + async def test_asearch_deep(self, monkeypatch): + monkeypatch.setenv("APISERPENT_API_KEY", "test-api-key") with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", new_callable=AsyncMock, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index f7ad333293c..30f479bd7ff 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -40,8 +40,8 @@ class TestAzureMAIImageGeneration: assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") - def test_mai_flash_and_2e_model_pricing_in_cost_map(self): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") flash_info = litellm.get_model_info( @@ -328,8 +328,8 @@ class TestAzureMAIImageGeneration: assert image_response.usage.total_tokens == 1046 assert image_response.size == "1792x1024" - def test_mai_image_cost_calculator_token_based(self): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + def test_mai_image_cost_calculator_token_based(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") @@ -360,8 +360,8 @@ class TestAzureMAIImageGeneration: ) assert round(cost, 10) == round(expected_cost, 10) - def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self, monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index a2f88138eaa..2d6e938ea1f 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -678,10 +678,10 @@ def test_transform_request_helper_includes_anthropic_beta_and_tools(): assert fields["tools"][0]["type"] == "computer_20250124" -def test_parallel_tool_calls_config_kept_for_sonnet_5(): +def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: config = AmazonConverseConfig() @@ -708,7 +708,7 @@ def test_parallel_tool_calls_config_kept_for_sonnet_5(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_parallel_tool_calls_config_dropped_for_ttl_only_model( @@ -3575,7 +3575,7 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): litellm.modify_params = original_modify_params -def test_supports_native_structured_outputs(): +def test_supports_native_structured_outputs(monkeypatch): """Test model detection for native structured outputs support. Support is driven by the ``supports_native_structured_output`` flag in the @@ -3583,7 +3583,7 @@ def test_supports_native_structured_outputs(): """ old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: config = AmazonConverseConfig() @@ -3645,7 +3645,7 @@ def test_supports_native_structured_outputs(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_create_output_config_for_response_format(): @@ -3683,11 +3683,11 @@ def test_create_output_config_for_response_format(): assert parsed_schema == expected -def test_translate_response_format_native_output_config(): +def test_translate_response_format_native_output_config(monkeypatch): """For supported models, _translate_response_format_param should produce outputConfig.""" old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: config = AmazonConverseConfig() @@ -3743,7 +3743,7 @@ def test_translate_response_format_native_output_config(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_translate_response_format_fallback_tool_call(): @@ -3778,11 +3778,11 @@ def test_translate_response_format_fallback_tool_call(): assert result["json_mode"] is True -def test_native_structured_output_no_fake_stream(): +def test_native_structured_output_no_fake_stream(monkeypatch): """When using native structured outputs with streaming, fake_stream should NOT be set.""" old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: config = AmazonConverseConfig() @@ -3828,7 +3828,7 @@ def test_native_structured_output_no_fake_stream(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_transform_request_with_output_config(): @@ -4116,7 +4116,7 @@ def test_add_additional_properties_definitions(): ) -def test_json_object_no_schema_skips_tool_injection(): +def test_json_object_no_schema_skips_tool_injection(monkeypatch): """response_format: {type: json_object} with no schema should NOT inject the synthetic json_tool_call tool. @@ -4126,7 +4126,7 @@ def test_json_object_no_schema_skips_tool_injection(): the model respond naturally with the JSON the caller asked for.""" old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: config = AmazonConverseConfig() @@ -4152,7 +4152,7 @@ def test_json_object_no_schema_skips_tool_injection(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_output_config_applies_additional_properties(): @@ -4805,7 +4805,7 @@ def test_cache_control_injection_tool_config_not_added_without_injection_point() assert all("cachePoint" not in tool for tool in tools) -def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(): +def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(monkeypatch): """ Regression test: cache_control_injection_points with location=tool_config must honor the requested `control.ttl`, mirroring the message/system @@ -4819,7 +4819,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(): """ old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: config = AmazonConverseConfig() @@ -4858,10 +4858,10 @@ def test_cache_control_injection_tool_config_honors_ttl_for_supported_model(): if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) -def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacking_own_pricing(): +def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacking_own_pricing(monkeypatch): """ Regression test: a regional pricing entry that omits `cache_creation_input_token_cost_above_1hr` (e.g. `jp.anthropic.claude-opus-4-7`) @@ -4870,7 +4870,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacki """ old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") try: assert "cache_creation_input_token_cost_above_1hr" not in litellm.model_cost["jp.anthropic.claude-opus-4-7"] @@ -4911,7 +4911,7 @@ def test_cache_control_injection_tool_config_honors_ttl_for_regional_model_lacki if old_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = old_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) def test_cache_control_injection_tool_config_drops_ttl_for_unsupported_model(): diff --git a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py index 950336c7ad0..20bf65ee385 100644 --- a/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/test_litellm/llms/bedrock/search/test_agentcore_search_transformation.py @@ -60,9 +60,9 @@ class TestAgentCoreSearch: """ @pytest.mark.asyncio - async def test_agentcore_search_request_payload(self): + async def test_agentcore_search_request_payload(self, monkeypatch): """Validates the MCP tools/call payload and SigV4 signing without real AWS calls.""" - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) mock_response = _make_mock_response(_mcp_response_body()) @@ -321,11 +321,11 @@ class TestAgentCoreSearch: assert headers["Authorization"] == "Bearer test-jwt-token" assert signed_body == json.dumps(request_data).encode() - def test_sign_request_uses_bearer_token_from_env(self): + def test_sign_request_uses_bearer_token_from_env(self, monkeypatch): """Server token is attached when the request targets the configured gateway host.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) try: headers, _ = config.sign_request( headers={}, @@ -338,11 +338,11 @@ class TestAgentCoreSearch: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) os.environ.pop("AGENTCORE_GATEWAY_URL", None) - def test_sign_request_refuses_server_token_to_untrusted_host(self): + def test_sign_request_refuses_server_token_to_untrusted_host(self, monkeypatch): """Server-managed token must not be sent to a caller-chosen api_base.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) try: with pytest.raises(ValueError, match="Refusing to send"): config.sign_request( @@ -355,11 +355,11 @@ class TestAgentCoreSearch: os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) os.environ.pop("AGENTCORE_GATEWAY_URL", None) - def test_sign_request_uses_env_token_for_gateway_api_base_without_gateway_url(self): + def test_sign_request_uses_env_token_for_gateway_api_base_without_gateway_url(self, monkeypatch): """api_base pointing at a real gateway is a trusted destination for the env token, so operators configuring api_base in yaml don't also need AGENTCORE_GATEWAY_URL.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") os.environ.pop("AGENTCORE_GATEWAY_URL", None) try: headers, _ = config.sign_request( @@ -380,12 +380,12 @@ class TestAgentCoreSearch: "https://attacker.example.com/gw.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", ], ) - def test_sign_request_refuses_sigv4_to_untrusted_host(self, untrusted_api_base): + def test_sign_request_refuses_sigv4_to_untrusted_host(self, untrusted_api_base, monkeypatch): """A SigV4 signature carries the proxy's credential scope and session token, so it must never be sent to a host that is not the operator's gateway.""" config = AgentCoreSearchConfig() os.environ.pop("AGENTCORE_GATEWAY_TOKEN", None) - os.environ["AGENTCORE_GATEWAY_URL"] = GATEWAY_URL + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", GATEWAY_URL) try: with patch.object( AgentCoreSearchConfig.__mro__[2], # BaseAWSLLM @@ -410,12 +410,12 @@ class TestAgentCoreSearch: "http://internal-gateway.corp/mcp", ], ) - def test_sign_request_refuses_server_token_over_plaintext_http(self, plaintext_api_base): + def test_sign_request_refuses_server_token_over_plaintext_http(self, plaintext_api_base, monkeypatch): """A trusted hostname over plain http would expose the bearer token to network observers, so credentials only ride https (or localhost).""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = plaintext_api_base + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", plaintext_api_base) try: with pytest.raises(ValueError, match="plaintext"): config.sign_request( @@ -446,11 +446,11 @@ class TestAgentCoreSearch: ) mock_base_sign.assert_not_called() - def test_sign_request_allows_plain_http_for_localhost(self): + def test_sign_request_allows_plain_http_for_localhost(self, monkeypatch): """Local development against an MCP stub on 127.0.0.1 keeps working.""" config = AgentCoreSearchConfig() - os.environ["AGENTCORE_GATEWAY_TOKEN"] = "env-jwt-token" - os.environ["AGENTCORE_GATEWAY_URL"] = "http://127.0.0.1:8931/mcp" + monkeypatch.setenv("AGENTCORE_GATEWAY_TOKEN", "env-jwt-token") + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", "http://127.0.0.1:8931/mcp") try: headers, _ = config.sign_request( headers={}, @@ -483,11 +483,11 @@ class TestAgentCoreSearch: # AWS_BEARER_TOKEN_BEDROCK env fallback. assert mock_base_sign.call_args.kwargs["api_key"] == "" - def test_sign_request_custom_hostname_requires_region(self): + def test_sign_request_custom_hostname_requires_region(self, monkeypatch): """Custom hostname + empty AWS config chain → clear error, no guessed region.""" config = AgentCoreSearchConfig() custom_url = "https://gateway.internal.example.com/mcp" - os.environ["AGENTCORE_GATEWAY_URL"] = custom_url + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", custom_url) mock_session = MagicMock() mock_session.region_name = None # nothing configured anywhere @@ -503,11 +503,11 @@ class TestAgentCoreSearch: finally: os.environ.pop("AGENTCORE_GATEWAY_URL", None) - def test_sign_request_custom_hostname_uses_shared_config_region(self): + def test_sign_request_custom_hostname_uses_shared_config_region(self, monkeypatch): """Custom hostname + region from AWS shared config (profile) must be honored.""" config = AgentCoreSearchConfig() custom_url = "https://gateway.internal.example.com/mcp" - os.environ["AGENTCORE_GATEWAY_URL"] = custom_url + monkeypatch.setenv("AGENTCORE_GATEWAY_URL", custom_url) mock_session = MagicMock() mock_session.region_name = "eu-west-1" # e.g. from ~/.aws/config profile diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py b/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py index daedbe5052c..962933aba28 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_ssl_verify.py @@ -40,12 +40,12 @@ class TestBedrockSSLVerify: ssl_verify = base_aws._get_ssl_verify() assert ssl_verify is True - def test_base_aws_llm_get_ssl_verify_false(self): + def test_base_aws_llm_get_ssl_verify_false(self, monkeypatch): """Test that _get_ssl_verify returns False when SSL verification is disabled.""" base_aws = BaseAWSLLM() # Set SSL_VERIFY to False via environment - os.environ["SSL_VERIFY"] = "False" + monkeypatch.setenv("SSL_VERIFY", "False") ssl_verify = base_aws._get_ssl_verify() assert ssl_verify is False @@ -53,7 +53,7 @@ class TestBedrockSSLVerify: # Clean up os.environ.pop("SSL_VERIFY", None) - def test_base_aws_llm_get_ssl_verify_custom_ca_bundle(self): + def test_base_aws_llm_get_ssl_verify_custom_ca_bundle(self, monkeypatch): """Test that _get_ssl_verify returns custom CA bundle path when SSL_CERT_FILE is set.""" base_aws = BaseAWSLLM() @@ -66,7 +66,7 @@ class TestBedrockSSLVerify: try: # Set SSL_CERT_FILE environment variable - os.environ["SSL_CERT_FILE"] = ca_bundle_path + monkeypatch.setenv("SSL_CERT_FILE", ca_bundle_path) os.environ.pop("SSL_VERIFY", None) litellm.ssl_verify = True @@ -327,7 +327,7 @@ class TestBedrockSSLVerify: os.environ.pop("SSL_CERT_FILE", None) os.unlink(ca_bundle_path) - def test_ssl_verify_priority_env_over_litellm_config(self): + def test_ssl_verify_priority_env_over_litellm_config(self, monkeypatch): """Test that SSL_VERIFY environment variable takes priority over litellm.ssl_verify.""" base_aws = BaseAWSLLM() @@ -335,7 +335,7 @@ class TestBedrockSSLVerify: litellm.ssl_verify = True # Set SSL_VERIFY environment variable to False - os.environ["SSL_VERIFY"] = "False" + monkeypatch.setenv("SSL_VERIFY", "False") try: ssl_verify = base_aws._get_ssl_verify() @@ -345,7 +345,7 @@ class TestBedrockSSLVerify: os.environ.pop("SSL_VERIFY", None) litellm.ssl_verify = True - def test_ssl_cert_file_priority_over_default(self): + def test_ssl_cert_file_priority_over_default(self, monkeypatch): """Test that SSL_CERT_FILE takes priority when ssl_verify is True.""" base_aws = BaseAWSLLM() @@ -358,7 +358,7 @@ class TestBedrockSSLVerify: try: # Set SSL_CERT_FILE environment variable - os.environ["SSL_CERT_FILE"] = ca_bundle_path + monkeypatch.setenv("SSL_CERT_FILE", ca_bundle_path) os.environ.pop("SSL_VERIFY", None) litellm.ssl_verify = True diff --git a/tests/test_litellm/llms/crusoe/test_crusoe.py b/tests/test_litellm/llms/crusoe/test_crusoe.py index 0a05126919a..34a6d37663b 100644 --- a/tests/test_litellm/llms/crusoe/test_crusoe.py +++ b/tests/test_litellm/llms/crusoe/test_crusoe.py @@ -105,14 +105,14 @@ def test_crusoe_provider_detection_by_prefix(): assert model == "meta-llama/Llama-3.3-70B-Instruct" -def test_crusoe_model_list_populated(): +def test_crusoe_model_list_populated(monkeypatch): """Test Crusoe models are present in model_prices_and_context_window.json""" import litellm original_model_cost = litellm.model_cost original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") try: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") expected = [ @@ -132,4 +132,4 @@ def test_crusoe_model_list_populated(): if original_env is None: os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", original_env) diff --git a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py index 3f772b263fd..153d37d549c 100644 --- a/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py +++ b/tests/test_litellm/llms/datarobot/chat/test_datarobot_chat_transformation.py @@ -83,8 +83,8 @@ class TestDataRobotConfig: == api_base ) - def test_resolve_api_base_with_environment_variable(self, handler): - os.environ["DATAROBOT_ENDPOINT"] = "https://env.datarobot.com" + def test_resolve_api_base_with_environment_variable(self, handler, monkeypatch): + monkeypatch.setenv("DATAROBOT_ENDPOINT", "https://env.datarobot.com") assert ( handler._resolve_api_base(None) == "https://env.datarobot.com/api/v2/genai/llmgw/chat/completions/" @@ -101,7 +101,7 @@ class TestDataRobotConfig: def test_resolve_api_key(self, api_key, expected_api_key, handler): assert handler._resolve_api_key(api_key) == expected_api_key - def test_resolve_api_key_with_environment_variable(self, handler): - os.environ["DATAROBOT_API_TOKEN"] = "env_key" + def test_resolve_api_key_with_environment_variable(self, handler, monkeypatch): + monkeypatch.setenv("DATAROBOT_API_TOKEN", "env_key") assert handler._resolve_api_key(None) == "env_key" del os.environ["DATAROBOT_API_TOKEN"] diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py index a5eb836e71d..ff309bc44ed 100644 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_chat_transformation.py @@ -11,14 +11,14 @@ sys.path.insert(0, os.path.abspath("../../../..")) import litellm -def test_deepseek_supported_openai_params(): +def test_deepseek_supported_openai_params(monkeypatch): """ Test "reasoning_effort" is an openai param supported for the DeepSeek model on deepinfra """ from litellm.llms.deepinfra.chat.transformation import DeepInfraConfig # Ensure we're using the local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") supported_openai_params = DeepInfraConfig().get_supported_openai_params( diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 6917092966b..fc8d71afaa9 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -81,8 +81,8 @@ def test_no_usage_details(): assert cost == 0.0 -def test_gemini_image_edit_cost_prefers_token_usage_metadata(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_edit_cost_prefers_token_usage_metadata(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -120,8 +120,8 @@ def test_gemini_image_edit_cost_prefers_token_usage_metadata(): assert cost != flat_image_cost -def test_gemini_image_edit_cost_uses_output_token_details(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_edit_cost_uses_output_token_details(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -176,8 +176,8 @@ def test_gemini_image_edit_cost_uses_output_token_details(): assert cost != all_output_as_image_cost -def test_gemini_image_generation_cost_uses_output_token_details(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_generation_cost_uses_output_token_details(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -232,8 +232,8 @@ def test_gemini_image_generation_cost_uses_output_token_details(): assert cost != all_output_as_image_cost -def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -264,8 +264,8 @@ def _image_response_with_web_search(web_search_requests): return ImageResponse(data=[ImageObject(b64_json="img1")], usage=usage) -def test_gemini_image_generation_cost_adds_web_search_grounding(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_generation_cost_adds_web_search_grounding(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -286,8 +286,8 @@ def test_gemini_image_generation_cost_adds_web_search_grounding(): assert round(grounded - ungrounded, 10) == round(expected_web_search_cost, 10) -def test_gemini_image_generation_cost_no_web_search_when_absent(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_gemini_image_generation_cost_no_web_search_when_absent(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py index 0750fb9e405..cff3c6be940 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py @@ -231,10 +231,10 @@ def test_inception_in_provider_lists(): assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints -def test_inception_model_configuration(): +def test_inception_model_configuration(monkeypatch): from litellm import get_model_info - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") litellm.inception_models = set() litellm.add_known_models() @@ -251,8 +251,8 @@ def test_inception_model_configuration(): assert info.get("supports_response_schema") is True -def test_inception_model_list_populated(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_inception_model_list_populated(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") litellm.inception_models = set() litellm.add_known_models() diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py index 9b7c8dd3742..62688a13c35 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py @@ -143,10 +143,10 @@ async def test_inception_fim_async(): assert r.choices[0].text == "a + b" -def test_inception_fim_model_configuration(): +def test_inception_fim_model_configuration(monkeypatch): from litellm import get_model_info - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") litellm.text_completion_inception_models = set() litellm.add_known_models() diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py index cd866187166..e54e25cbd18 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_cost_calculator.py @@ -30,8 +30,8 @@ def _image_response_with_web_search(web_search_requests): return ImageResponse(data=[ImageObject(b64_json="img1")], usage=usage) -def test_vertex_image_generation_cost_adds_web_search_grounding(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_vertex_image_generation_cost_adds_web_search_grounding(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") @@ -55,8 +55,8 @@ def test_vertex_image_generation_cost_adds_web_search_grounding(): assert round(grounded - ungrounded, 10) == round(expected_web_search_cost, 10) -def test_vertex_image_generation_cost_no_web_search_when_absent(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_vertex_image_generation_cost_no_web_search_when_absent(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3-pro-image-preview" diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index e8374f92a19..61e1121257c 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -51,11 +51,11 @@ def test_zai_in_provider_lists(): assert "zai" in litellm.provider_list -def test_zai_models_in_model_cost(): +def test_zai_models_in_model_cost(monkeypatch): """Test that ZAI models are in the model cost map""" import os - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") zai_models = [ @@ -75,11 +75,11 @@ def test_zai_models_in_model_cost(): assert litellm.model_cost[model]["litellm_provider"] == "zai" -def test_zai_glm46_cost_calculation(): +def test_zai_glm46_cost_calculation(monkeypatch): """Test the cost calculation for glm-4.6""" import os - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") key = "zai/glm-4.6" @@ -96,11 +96,11 @@ def test_zai_glm46_cost_calculation(): assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) -def test_zai_flash_model_is_free(): +def test_zai_flash_model_is_free(monkeypatch): """Test that glm-4.5-flash has zero cost""" import os - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") key = "zai/glm-4.5-flash" @@ -110,11 +110,11 @@ def test_zai_flash_model_is_free(): assert info["output_cost_per_token"] == 0 -def test_glm47_supports_reasoning(): +def test_glm47_supports_reasoning(monkeypatch): """Test that GLM-4.7 supports reasoning""" import os - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") key = "zai/glm-4.7" @@ -124,11 +124,11 @@ def test_glm47_supports_reasoning(): assert info["supports_reasoning"] is True -def test_glm47_cost_calculation(): +def test_glm47_cost_calculation(monkeypatch): """Test cost calculation for GLM-4.7""" import os - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") prompt_cost, completion_cost = cost_per_token( diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index c589014f276..1c66acf8678 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -109,7 +109,7 @@ async def test_authenticate_user_admin_login_with_ui_credentials(): @pytest.mark.asyncio -async def test_authenticate_user_admin_login_with_master_key_as_password(): +async def test_authenticate_user_admin_login_with_master_key_as_password(monkeypatch): """Test admin login when UI_PASSWORD is not set, should use master_key""" master_key = "sk-1234" ui_username = "admin" @@ -131,39 +131,35 @@ async def test_authenticate_user_admin_login_with_master_key_as_password(): with patch.dict(os.environ, env_vars, clear=False): # Explicitly remove UI_PASSWORD if it exists - original_ui_password = os.environ.pop("UI_PASSWORD", None) - try: + monkeypatch.delenv("UI_PASSWORD", raising=False) + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = { + "token": "test-token-123", + "user_id": LITELLM_PROXY_ADMIN_NAME, + } + with patch( - "litellm.proxy.auth.login_utils.generate_key_helper_fn", + "litellm.proxy.auth.login_utils.user_update", new_callable=AsyncMock, - ) as mock_generate_key: - mock_generate_key.return_value = { - "token": "test-token-123", - "user_id": LITELLM_PROXY_ADMIN_NAME, - } - + return_value=None, + ) as mock_user_update: with patch( - "litellm.proxy.auth.login_utils.user_update", - new_callable=AsyncMock, - return_value=None, - ) as mock_user_update: - with patch( - "litellm.proxy.auth.login_utils.get_secret_bool", - return_value=False, - ): - result = await authenticate_user( - username=ui_username, - password=master_key, - master_key=master_key, - prisma_client=mock_prisma_client, - ) + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ): + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + ) - assert isinstance(result, LoginResult) - assert result.user_id == LITELLM_PROXY_ADMIN_NAME - assert result.user_role == LitellmUserRoles.PROXY_ADMIN - finally: - if original_ui_password: - os.environ["UI_PASSWORD"] = original_ui_password + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + assert result.user_role == LitellmUserRoles.PROXY_ADMIN @pytest.mark.asyncio @@ -319,7 +315,7 @@ async def test_authenticate_user_email_case_insensitive_login(): @pytest.mark.asyncio -async def test_authenticate_user_database_required_for_admin(): +async def test_authenticate_user_database_required_for_admin(monkeypatch): """Test that database is required for admin login""" master_key = "sk-1234" ui_username = "admin" @@ -353,7 +349,7 @@ async def test_authenticate_user_database_required_for_admin(): assert "No Database connected" in exc_info.value.message finally: if original_db_url: - os.environ["DATABASE_URL"] = original_db_url + monkeypatch.setenv("DATABASE_URL", original_db_url) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py index a2b8894910c..af0686fcc59 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py @@ -17,14 +17,14 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.exceptions import GuardrailRaisedException -def test_deepkeep_guard_config(): +def test_deepkeep_guard_config(monkeypatch): """Test DeepKeep guard configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} - os.environ["DEEPKEEP_API_KEY"] = "test-key" - os.environ["DEEPKEEP_API_BASE"] = "https://test.deepkeep.ai" - os.environ["DEEPKEEP_FIREWALL_ID"] = "fw-123" + monkeypatch.setenv("DEEPKEEP_API_KEY", "test-key") + monkeypatch.setenv("DEEPKEEP_API_BASE", "https://test.deepkeep.ai") + monkeypatch.setenv("DEEPKEEP_FIREWALL_ID", "fw-123") init_guardrails_v2( all_guardrails=[ @@ -108,11 +108,11 @@ class TestDeepKeepGuardrail: == "https://test.deepkeep.ai/v3/openai/beta/litellm_basic_guardrail_api" ) - def test_initialization_with_env_vars(self): + def test_initialization_with_env_vars(self, monkeypatch): """should initialize successfully using environment variables.""" - os.environ["DEEPKEEP_API_KEY"] = "env-key" - os.environ["DEEPKEEP_API_BASE"] = "https://env.deepkeep.ai" - os.environ["DEEPKEEP_FIREWALL_ID"] = "fw-env-456" + monkeypatch.setenv("DEEPKEEP_API_KEY", "env-key") + monkeypatch.setenv("DEEPKEEP_API_BASE", "https://env.deepkeep.ai") + monkeypatch.setenv("DEEPKEEP_FIREWALL_ID", "fw-env-456") guardrail = DeepKeepGuardrail( guardrail_name="deepkeep-env-test", diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index c5b182a00ab..57adf85b3d9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -26,13 +26,13 @@ from litellm.types.utils import ( ) -def test_hiddenlayer_config_saas(): +def test_hiddenlayer_config_saas(monkeypatch): """Test Hiddenlayer SaaS configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variables for testing - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") init_guardrails_v2( all_guardrails=[ @@ -71,9 +71,9 @@ class TestHiddenlayerGuardrail: if key in os.environ: del os.environ[key] - def test_initialization(self): + def test_initialization(self, monkeypatch): """Test successful initialization with default values.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -94,9 +94,9 @@ class TestHiddenlayerGuardrail: HiddenlayerGuardrail(guardrail_name="hiddenlayer", event_hook="pre_call") @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self): + async def test_apply_guardrail_request_no_violations(self, monkeypatch): """Test apply_guardrail for request with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -151,9 +151,9 @@ class TestHiddenlayerGuardrail: assert call_args.args[0] == f"{guardrail.api_base}/detection/v1/interactions" @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self): + async def test_apply_guardrail_request_with_violations(self, monkeypatch): """Test apply_guardrail for request with violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -209,9 +209,9 @@ class TestHiddenlayerGuardrail: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self): + async def test_apply_guardrail_response_no_violations(self, monkeypatch): """Test apply_guardrail for response with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -279,10 +279,10 @@ class TestHiddenlayerGuardrail: mock_post.assert_called_once() @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self): + async def test_apply_guardrail_response_with_violations(self, monkeypatch): """Test apply_guardrail for response with violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") # Setup guardrail guardrail = HiddenlayerGuardrail( @@ -348,10 +348,10 @@ class TestHiddenlayerGuardrail: assert exc_info.value.status_code == 400 @pytest.mark.asyncio - async def test_apply_guardrail_api_error_handling(self): + async def test_apply_guardrail_api_error_handling(self, monkeypatch): """Test handling of API errors in apply_guardrail.""" # Set required API key - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -391,10 +391,10 @@ class TestHiddenlayerGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_validate_with_call_hiddenlayer_method(self): + async def test_validate_with_call_hiddenlayer_method(self, monkeypatch): """Test the _validate_with_guard_server internal method.""" # Set required API key - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -433,9 +433,9 @@ class TestHiddenlayerGuardrail: ) @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image(self): + async def test_apply_guardrail_request_with_image(self, monkeypatch): """Test apply_guardrail sends multimodal content (image) to HiddenLayer v1.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -498,9 +498,9 @@ class TestHiddenlayerGuardrail: assert result is not None @pytest.mark.asyncio - async def test_apply_guardrail_redact_with_image_content(self): + async def test_apply_guardrail_redact_with_image_content(self, monkeypatch): """Test that REDACT action with multimodal content extracts text properly into inputs['texts'].""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -570,12 +570,12 @@ class TestHiddenlayerGuardrail: assert config_model.__name__ == "HiddenlayerGuardrailConfigModel" -def test_hiddenlayer_config_v2(): +def test_hiddenlayer_config_v2(monkeypatch): """Test HiddenLayer V2 configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") init_guardrails_v2( all_guardrails=[ @@ -612,9 +612,9 @@ class TestHiddenlayerGuardrailV2: if key in os.environ: del os.environ[key] - def test_initialization(self): + def test_initialization(self, monkeypatch): """Test successful initialization with default values.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -633,9 +633,9 @@ class TestHiddenlayerGuardrailV2: HiddenlayerGuardrailV2(guardrail_name="hiddenlayer", event_hook="pre_call") @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self): + async def test_apply_guardrail_request_no_violations(self, monkeypatch): """Test apply_guardrail for request with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -691,9 +691,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/request-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self): + async def test_apply_guardrail_request_with_violations(self, monkeypatch): """Test apply_guardrail for request with violations detected (block via header).""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -751,9 +751,9 @@ class TestHiddenlayerGuardrailV2: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self): + async def test_apply_guardrail_response_no_violations(self, monkeypatch): """Test apply_guardrail for response with no violations detected.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="post_call", default_on=True @@ -816,9 +816,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self): + async def test_apply_guardrail_response_with_violations(self, monkeypatch): """Test apply_guardrail for response with violations detected (block via header).""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="post_call", default_on=True @@ -863,9 +863,9 @@ class TestHiddenlayerGuardrailV2: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_with_tool_calls(self): + async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch): """Test apply_guardrail for response containing tool calls.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="post_call", default_on=True @@ -924,9 +924,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_call_hiddenlayer_uses_correct_endpoints(self): + async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch): """Test that _call_hiddenlayer uses the v2 request/response evaluation endpoints.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -959,9 +959,9 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in mock_post.call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image(self): + async def test_apply_guardrail_request_with_image(self, monkeypatch): """Test apply_guardrail sends multimodal content (image) to HiddenLayer v2.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True @@ -1030,9 +1030,9 @@ class TestHiddenlayerGuardrailV2: assert texts == ["how much is on this receipt?"] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image_multimodal_response(self): + async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch): """Test that new_texts extraction handles multimodal content (list) returned by HiddenLayer v2.""" - os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer" + monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrailV2( guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py index 16185cadbdf..dcb004e5422 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -19,13 +19,13 @@ from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import ( from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -def test_lasso_guard_config(): +def test_lasso_guard_config(monkeypatch): """Test Lasso guard configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variable for testing - os.environ["LASSO_API_KEY"] = "test-key" + monkeypatch.setenv("LASSO_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py index c7a6df1361e..fa4624eac99 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py @@ -18,14 +18,14 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message -def test_onyx_guard_config(): +def test_onyx_guard_config(monkeypatch): """Test Onyx guard configuration with init_guardrails_v2.""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") init_guardrails_v2( all_guardrails=[ @@ -48,11 +48,11 @@ def test_onyx_guard_config(): del os.environ["ONYX_API_KEY"] -def test_onyx_guard_with_custom_timeout_from_kwargs(): +def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch): """Test Onyx guard instantiation with custom timeout passed via kwargs.""" # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -81,16 +81,16 @@ def test_onyx_guard_with_custom_timeout_from_kwargs(): del os.environ["ONYX_API_KEY"] -def test_onyx_guard_with_timeout_none_uses_env_var(): +def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch): """Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var. When timeout=None is passed (as it would be from config model with default None), the ONYX_TIMEOUT environment variable should be used. """ # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" - os.environ["ONYX_TIMEOUT"] = "60" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") + monkeypatch.setenv("ONYX_TIMEOUT", "60") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -121,11 +121,11 @@ def test_onyx_guard_with_timeout_none_uses_env_var(): del os.environ["ONYX_TIMEOUT"] -def test_onyx_guard_with_timeout_none_defaults_to_10(): +def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch): """Test Onyx guard with timeout=None and no env var defaults to 10 seconds.""" # Set environment variables for testing - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Ensure ONYX_TIMEOUT is not set if "ONYX_TIMEOUT" in os.environ: del os.environ["ONYX_TIMEOUT"] @@ -174,10 +174,10 @@ class TestOnyxGuardrail: if key in os.environ: del os.environ[key] - def test_initialization_with_defaults(self): + def test_initialization_with_defaults(self, monkeypatch): """Test successful initialization with default values.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -189,10 +189,10 @@ class TestOnyxGuardrail: assert guardrail.guardrail_name == "test-guard" assert guardrail.event_hook == "pre_call" - def test_initialization_with_env_vars(self): + def test_initialization_with_env_vars(self, monkeypatch): """Test initialization with environment variables.""" - os.environ["ONYX_API_BASE"] = "https://custom.onyx.security" - os.environ["ONYX_API_KEY"] = "custom-api-key" + monkeypatch.setenv("ONYX_API_BASE", "https://custom.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "custom-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -213,9 +213,9 @@ class TestOnyxGuardrail: ): OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call") - def test_initialization_with_default_timeout(self): + def test_initialization_with_default_timeout(self, monkeypatch): """Test that default timeout is 10.0 seconds.""" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -232,9 +232,9 @@ class TestOnyxGuardrail: assert timeout_param.read == 10.0 assert timeout_param.connect == 5.0 - def test_initialization_with_custom_timeout_parameter(self): + def test_initialization_with_custom_timeout_parameter(self, monkeypatch): """Test initialization with custom timeout parameter.""" - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -254,14 +254,14 @@ class TestOnyxGuardrail: assert timeout_param.read == 30.0 assert timeout_param.connect == 5.0 - def test_initialization_with_timeout_from_env_var(self): + def test_initialization_with_timeout_from_env_var(self, monkeypatch): """Test initialization with timeout from ONYX_TIMEOUT environment variable. Note: The env var is only used when timeout=None is explicitly passed, since the default parameter value is 10.0 (not None). """ - os.environ["ONYX_API_KEY"] = "test-api-key" - os.environ["ONYX_TIMEOUT"] = "25" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") + monkeypatch.setenv("ONYX_TIMEOUT", "25") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -282,10 +282,10 @@ class TestOnyxGuardrail: assert timeout_param.read == 25.0 assert timeout_param.connect == 5.0 - def test_initialization_timeout_parameter_overrides_env_var(self): + def test_initialization_timeout_parameter_overrides_env_var(self, monkeypatch): """Test that timeout parameter overrides ONYX_TIMEOUT environment variable.""" - os.environ["ONYX_API_KEY"] = "test-api-key" - os.environ["ONYX_TIMEOUT"] = "25" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") + monkeypatch.setenv("ONYX_TIMEOUT", "25") with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -306,10 +306,10 @@ class TestOnyxGuardrail: assert timeout_param.connect == 5.0 @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self): + async def test_apply_guardrail_request_no_violations(self, monkeypatch): """Test apply_guardrail for request with no violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -372,10 +372,10 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-call-id" @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self): + async def test_apply_guardrail_request_with_violations(self, monkeypatch): """Test apply_guardrail for request with violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -423,10 +423,10 @@ class TestOnyxGuardrail: assert "prompt_injection" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self): + async def test_apply_guardrail_response_no_violations(self, monkeypatch): """Test apply_guardrail for response with no violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -497,10 +497,10 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-call-id-2" @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self): + async def test_apply_guardrail_response_with_violations(self, monkeypatch): """Test apply_guardrail for response with violations detected.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail guardrail = OnyxGuardrail( @@ -558,10 +558,10 @@ class TestOnyxGuardrail: assert "illegal_instructions" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_api_error_handling(self): + async def test_apply_guardrail_api_error_handling(self, monkeypatch): """Test handling of API errors in apply_guardrail.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -591,10 +591,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_timeout_error_handling(self): + async def test_apply_guardrail_timeout_error_handling(self, monkeypatch): """Test handling of timeout errors in apply_guardrail (graceful degradation).""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", @@ -629,10 +629,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_read_timeout_error_handling(self): + async def test_apply_guardrail_read_timeout_error_handling(self, monkeypatch): """Test handling of read timeout errors in apply_guardrail.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", @@ -667,10 +667,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_connect_timeout_error_handling(self): + async def test_apply_guardrail_connect_timeout_error_handling(self, monkeypatch): """Test handling of connect timeout errors in apply_guardrail.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", @@ -705,10 +705,10 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_no_logging_obj(self): + async def test_apply_guardrail_no_logging_obj(self, monkeypatch): """Test apply_guardrail without logging object (uses UUID).""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -747,10 +747,10 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-uuid" @pytest.mark.asyncio - async def test_validate_with_guard_server_method(self): + async def test_validate_with_guard_server_method(self, monkeypatch): """Test the _validate_with_guard_server internal method.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -788,10 +788,10 @@ class TestOnyxGuardrail: ) @pytest.mark.asyncio - async def test_validate_with_guard_server_blocked(self): + async def test_validate_with_guard_server_blocked(self, monkeypatch): """Test _validate_with_guard_server when request is blocked.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -825,10 +825,10 @@ class TestOnyxGuardrail: assert config_model.__name__ == "OnyxGuardrailConfigModel" @pytest.mark.asyncio - async def test_apply_guardrail_with_modelresponse(self): + async def test_apply_guardrail_with_modelresponse(self, monkeypatch): """Test apply_guardrail with ModelResponse object for response type.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -880,10 +880,10 @@ class TestOnyxGuardrail: assert "payload" in call_args.kwargs["json"] @pytest.mark.asyncio - async def test_apply_guardrail_response_error_handling(self): + async def test_apply_guardrail_response_error_handling(self, monkeypatch): """Test error handling when processing response data.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -925,11 +925,11 @@ class TestOnyxIntegration: """Test integration scenarios.""" @pytest.mark.asyncio - async def test_full_guardrail_flow(self): + async def test_full_guardrail_flow(self, monkeypatch): """Test full guardrail flow with multiple hooks.""" # Set environment variables - os.environ["ONYX_API_BASE"] = "https://test.onyx.security" - os.environ["ONYX_API_KEY"] = "test-key" + monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") + monkeypatch.setenv("ONYX_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ @@ -973,10 +973,10 @@ class TestOnyxIntegration: del os.environ["ONYX_API_KEY"] @pytest.mark.asyncio - async def test_apply_guardrail_empty_request_data(self): + async def test_apply_guardrail_empty_request_data(self, monkeypatch): """Test apply_guardrail with empty request data.""" # Set required API key - os.environ["ONYX_API_KEY"] = "test-api-key" + monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py index 55f01ebddfd..1322d93ce70 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py @@ -93,24 +93,24 @@ class TestRepelloAIInitialization: with pytest.raises(ValueError, match="asset_id"): RepelloAIGuardrail(api_key="test-api-key", guardrail_name="t") - def test_api_key_from_env(self): - os.environ["REPELLOAI_API_KEY"] = "env-key" + def test_api_key_from_env(self, monkeypatch): + monkeypatch.setenv("REPELLOAI_API_KEY", "env-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "env-key" - def test_api_key_from_argus_env(self): - os.environ["ARGUS_API_KEY"] = "argus-key" + def test_api_key_from_argus_env(self, monkeypatch): + monkeypatch.setenv("ARGUS_API_KEY", "argus-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "argus-key" - def test_argus_env_preferred_over_legacy(self): - os.environ["ARGUS_API_KEY"] = "argus-key" - os.environ["REPELLOAI_API_KEY"] = "legacy-key" + def test_argus_env_preferred_over_legacy(self, monkeypatch): + monkeypatch.setenv("ARGUS_API_KEY", "argus-key") + monkeypatch.setenv("REPELLOAI_API_KEY", "legacy-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "argus-key" - def test_explicit_api_key_preferred_over_env(self): - os.environ["ARGUS_API_KEY"] = "argus-key" + def test_explicit_api_key_preferred_over_env(self, monkeypatch): + monkeypatch.setenv("ARGUS_API_KEY", "argus-key") guardrail = RepelloAIGuardrail( api_key="explicit-key", asset_id="asset-123", guardrail_name="t" ) @@ -145,10 +145,10 @@ class TestRepelloAIInitialization: assert guardrail.api_base == DEFAULT_REPELLOAI_API_BASE assert guardrail.unreachable_fallback == "fail_closed" - def test_init_guardrails_v2_wiring(self): + def test_init_guardrails_v2_wiring(self, monkeypatch): """The guardrail registers and constructs via the config.yaml path.""" litellm.guardrail_name_config_map = {} - os.environ["REPELLOAI_API_KEY"] = "test-key" + monkeypatch.setenv("REPELLOAI_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ { diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index c8f22e6c15e..996a3ff0824 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -19,14 +19,14 @@ import litellm from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -def test_prompt_security_guard_config(): +def test_prompt_security_guard_config(monkeypatch): """Test guardrail initialization with proper configuration""" litellm.set_verbose = True litellm.guardrail_name_config_map = {} # Set environment variables for testing - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") init_guardrails_v2( all_guardrails=[ @@ -78,10 +78,10 @@ def test_prompt_security_guard_config_no_api_key(): @pytest.mark.asyncio -async def test_apply_guardrail_block_request(): +async def test_apply_guardrail_block_request(monkeypatch): """Test that apply_guardrail blocks malicious prompts""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -132,10 +132,10 @@ async def test_apply_guardrail_block_request(): @pytest.mark.asyncio -async def test_apply_guardrail_modify_request(): +async def test_apply_guardrail_modify_request(monkeypatch): """Test that apply_guardrail modifies prompts when needed""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -183,10 +183,10 @@ async def test_apply_guardrail_modify_request(): @pytest.mark.asyncio -async def test_apply_guardrail_allow_request(): +async def test_apply_guardrail_allow_request(monkeypatch): """Test that apply_guardrail allows safe prompts""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -226,10 +226,10 @@ async def test_apply_guardrail_allow_request(): @pytest.mark.asyncio -async def test_apply_guardrail_block_response(): +async def test_apply_guardrail_block_response(monkeypatch): """Test that apply_guardrail blocks malicious responses""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -273,10 +273,10 @@ async def test_apply_guardrail_block_response(): @pytest.mark.asyncio -async def test_apply_guardrail_modify_response(): +async def test_apply_guardrail_modify_response(monkeypatch): """Test that apply_guardrail modifies responses when needed""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="post_call", default_on=True @@ -317,10 +317,10 @@ async def test_apply_guardrail_modify_response(): @pytest.mark.asyncio -async def test_file_sanitization(): +async def test_file_sanitization(monkeypatch): """Test file sanitization for images""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -407,10 +407,10 @@ async def test_file_sanitization(): @pytest.mark.asyncio -async def test_file_sanitization_block(): +async def test_file_sanitization_block(monkeypatch): """Test that file sanitization blocks malicious files""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -491,10 +491,10 @@ async def test_file_sanitization_block(): @pytest.mark.asyncio -async def test_user_api_key_alias_forwarding(): +async def test_user_api_key_alias_forwarding(monkeypatch): """Test that user API key alias is properly sent via headers and payload""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -535,10 +535,10 @@ async def test_user_api_key_alias_forwarding(): @pytest.mark.asyncio -async def test_role_filtering(): +async def test_role_filtering(monkeypatch): """Test that tool/function messages are filtered out by default""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True @@ -600,11 +600,11 @@ async def test_role_filtering(): @pytest.mark.asyncio -async def test_check_tool_results_enabled(): +async def test_check_tool_results_enabled(monkeypatch): """Test with check_tool_results=True: transforms tool/function to 'other' role""" - os.environ["PROMPT_SECURITY_API_KEY"] = "test-key" - os.environ["PROMPT_SECURITY_API_BASE"] = "https://test.prompt.security" - os.environ["PROMPT_SECURITY_CHECK_TOOL_RESULTS"] = "true" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true") guardrail = PromptSecurityGuardrail( guardrail_name="test-guard", event_hook="pre_call", default_on=True diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 6c717d6f71c..13997fc4cd1 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -42,7 +42,7 @@ def time_controller(monkeypatch): @pytest.mark.asyncio -async def test_priority_weight_allocation(): +async def test_priority_weight_allocation(monkeypatch): """ Test that priority weights are correctly applied instead of equal splitting. @@ -53,7 +53,7 @@ async def test_priority_weight_allocation(): This validates the core fix where before it would split 50/50. """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up priority reservations litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -128,7 +128,7 @@ async def test_priority_weight_allocation(): @pytest.mark.asyncio -async def test_concurrent_priority_requests(): +async def test_concurrent_priority_requests(monkeypatch): """ Test the core issue: 5 concurrent requests with different priorities should get proper allocation based on priority weights, not equal splitting. @@ -136,7 +136,7 @@ async def test_concurrent_priority_requests(): This tests the exact scenario mentioned: priorities 0.9 and 0.1 should be 0.9/0.1, not 0.5/0.5. """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up the exact scenario from the issue litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -214,7 +214,7 @@ async def test_concurrent_priority_requests(): @pytest.mark.asyncio -async def test_100_concurrent_priority_requests(time_controller): +async def test_100_concurrent_priority_requests(time_controller, monkeypatch): """ Stress test: 100 concurrent requests with mixed priorities over 10 seconds. @@ -224,7 +224,7 @@ async def test_100_concurrent_priority_requests(time_controller): - Spread across 10 seconds to simulate real-world load """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up priority reservations litellm.priority_reservation = {"high": 0.9, "low": 0.1} @@ -384,7 +384,7 @@ async def test_100_concurrent_priority_requests(time_controller): @pytest.mark.asyncio -async def test_concurrent_pre_call_hooks_stress(): +async def test_concurrent_pre_call_hooks_stress(monkeypatch): """ Stress test: 50 concurrent pre-call hooks with saturation-aware priority enforcement. @@ -394,7 +394,7 @@ async def test_concurrent_pre_call_hooks_stress(): Standard users (20% allocation) should have ~70% success rate with 30% random limiting. """ # Set up environment for premium feature - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"premium": 0.8, "standard": 0.2} @@ -634,7 +634,7 @@ async def test_concurrent_pre_call_hooks_stress(): @pytest.mark.asyncio -async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): +async def test_fake_calls_case_1_no_rate_limiting_at_capacity(monkeypatch): """ Test Case 1: Saturation-Aware Rate Limiting at 50% Threshold @@ -650,7 +650,7 @@ async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): Once saturation hits 50%, strict mode enforces priority-based limits. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set up priority reservations litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} @@ -759,7 +759,7 @@ async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): @pytest.mark.asyncio -async def test_fake_calls_case_2_priority_queue_during_saturation(): +async def test_fake_calls_case_2_priority_queue_during_saturation(monkeypatch): """ Test Case 2: Priority Queue Behavior During Saturation @@ -773,7 +773,7 @@ async def test_fake_calls_case_2_priority_queue_during_saturation(): When total traffic exceeds capacity, rate limiting enforces priority reservations. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} @@ -886,7 +886,7 @@ async def test_fake_calls_case_2_priority_queue_during_saturation(): @pytest.mark.asyncio -async def test_fake_calls_case_3_spillover_capacity_default_keys(): +async def test_fake_calls_case_3_spillover_capacity_default_keys(monkeypatch): """ Test Case 3: Spillover Capacity for Default Keys @@ -906,7 +906,7 @@ async def test_fake_calls_case_3_spillover_capacity_default_keys(): Tests spillover behavior where default keys share remaining capacity. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.75} litellm.priority_reservation_settings.default_priority = 0.25 @@ -1025,7 +1025,7 @@ async def test_fake_calls_case_3_spillover_capacity_default_keys(): @pytest.mark.asyncio -async def test_fake_calls_case_4_over_allocated_with_normalization(): +async def test_fake_calls_case_4_over_allocated_with_normalization(monkeypatch): """ Test Case 4: Over-Allocated Priority reservations with Normalization @@ -1042,7 +1042,7 @@ async def test_fake_calls_case_4_over_allocated_with_normalization(): - Due to concurrent burst, total successful may exceed 100 RPM in the test window - This test verifies normalization works and total capacity is reasonably bounded """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.60, "key_b": 0.80} @@ -1156,7 +1156,7 @@ async def test_fake_calls_case_4_over_allocated_with_normalization(): @pytest.mark.asyncio -async def test_fake_calls_case_5_default_value_priority_reservation(): +async def test_fake_calls_case_5_default_value_priority_reservation(monkeypatch): """ Test Case 5: Default value for priority reservation @@ -1176,7 +1176,7 @@ async def test_fake_calls_case_5_default_value_priority_reservation(): Tests complex scenario with explicit priorities and default priority. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"key_a": 0.50, "key_b": 0.20, "key_c": 0.05} litellm.priority_reservation_settings.default_priority = 0.05 @@ -1296,7 +1296,7 @@ async def test_fake_calls_case_5_default_value_priority_reservation(): @pytest.mark.asyncio -async def test_default_priority_shared_pool(): +async def test_default_priority_shared_pool(monkeypatch): """ Test that keys without explicit priority share ONE default pool, not get individual allocations. @@ -1304,7 +1304,7 @@ async def test_default_priority_shared_pool(): - Key A, B, C (no priority) should share ONE 25 RPM pool - NOT get 25 RPM each (which would be 75 RPM total) """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"prod": 0.75} litellm.priority_reservation_settings.default_priority = 0.25 @@ -1382,7 +1382,7 @@ async def test_default_priority_shared_pool(): @pytest.mark.asyncio -async def test_async_log_success_event_increments_by_actual_tokens(): +async def test_async_log_success_event_increments_by_actual_tokens(monkeypatch): """ Test that async_log_success_event increments token counters by actual token usage. @@ -1394,7 +1394,7 @@ async def test_async_log_success_event_increments_by_actual_tokens(): from litellm.types.utils import ModelResponse, Usage - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"dev": 0.1, "prod": 0.9} dual_cache = DualCache() @@ -1483,7 +1483,7 @@ async def test_async_log_success_event_increments_by_actual_tokens(): @pytest.mark.asyncio -async def test_saturation_check_cache_ttl_configuration(): +async def test_saturation_check_cache_ttl_configuration(monkeypatch): """ Test that saturation_check_cache_ttl controls how long saturation values are cached locally. @@ -1492,7 +1492,7 @@ async def test_saturation_check_cache_ttl_configuration(): - After expiration, fresh values should be fetched from Redis - This prevents nodes from having stale saturation data in multi-node deployments """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") # Set a short TTL for testing (5 seconds) original_ttl = litellm.priority_reservation_settings.saturation_check_cache_ttl @@ -1587,7 +1587,7 @@ async def test_saturation_check_cache_ttl_configuration(): @pytest.mark.asyncio -async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): +async def test_async_log_success_event_uses_team_priority_from_auth_metadata(monkeypatch): """ Test that async_log_success_event correctly retrieves priority from user_api_key_auth_metadata. @@ -1598,7 +1598,7 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): from litellm.types.utils import ModelResponse, Usage - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"team_priority": 0.8, "default": 0.2} dual_cache = DualCache() @@ -1680,7 +1680,7 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(): @pytest.mark.asyncio -async def test_priority_429_includes_model_name_and_configured_limits(): +async def test_priority_429_includes_model_name_and_configured_limits(monkeypatch): """ The priority-based 429 should tell operators which model was hit and what the model's configured TPM/RPM are, so they can decide whether to tune the @@ -1694,7 +1694,7 @@ async def test_priority_429_includes_model_name_and_configured_limits(): """ from fastapi import HTTPException - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"prod": 0.5} dual_cache = DualCache() @@ -1774,7 +1774,7 @@ async def test_priority_429_includes_model_name_and_configured_limits(): @pytest.mark.asyncio -async def test_tpm_only_model_enforces_priority_and_model_capacity(): +async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch): """Regression: a model configured with ONLY tpm (no rpm) must still be rate limited. @@ -1789,7 +1789,7 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(): from litellm.types.utils import ModelResponse, Usage - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"dev": 0.25, "prod": 0.5} dual_cache = DualCache() diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py index 1c1e8eee145..97a986d1ade 100644 --- a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -189,7 +189,7 @@ async def test_batch_limiter_uses_atomic_check_and_increment(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): +async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(monkeypatch): """ DynamicRateLimitHandler PHASE 1 (read_only check) → PHASE 3 (increment) is non-atomic: dynamic_rate_limiter_v3.py:463-548. @@ -209,7 +209,7 @@ async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): # RPM + 1 successes before the next sees counter > RPM. MAX_SEQUENTIAL_SUCCESSES = MODEL_RPM + 1 - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"high": 0.9, "low": 0.1} dual_cache = DualCache() @@ -273,7 +273,7 @@ async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): +async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(monkeypatch): """ Regression test: dynamic limiter's enforced descriptors flow through `atomic_check_and_increment_by_n`, not the legacy @@ -283,7 +283,7 @@ async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): bundled into the atomic call alongside model_saturation_check. When not enforced, priority counter is incremented for tracking only. """ - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"high": 0.9, "low": 0.1} dual_cache = DualCache() @@ -413,7 +413,7 @@ async def test_batch_zero_token_consumes_rpm_only(): @pytest.mark.asyncio -async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): +async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(monkeypatch): """ Fail-closed guard: when atomic_check_and_increment_by_n returns overall_code=OVER_LIMIT but with a descriptor_key the dispatcher does @@ -425,7 +425,7 @@ async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): """ from fastapi import HTTPException - os.environ["LITELLM_LICENSE"] = "test-license-key" + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") litellm.priority_reservation = {"high": 0.9, "low": 0.1} dual_cache = DualCache() diff --git a/tests/test_litellm/test_add_deployment_no_master_key.py b/tests/test_litellm/test_add_deployment_no_master_key.py index 6db20d7d422..f7a0e90dad0 100644 --- a/tests/test_litellm/test_add_deployment_no_master_key.py +++ b/tests/test_litellm/test_add_deployment_no_master_key.py @@ -62,7 +62,7 @@ async def test_add_deployment_without_master_key(): @pytest.mark.asyncio -async def test_add_deployment_without_salt_key_or_master_key(): +async def test_add_deployment_without_salt_key_or_master_key(monkeypatch): """ Test that add_deployment() works when both master_key and LITELLM_SALT_KEY are None. @@ -70,55 +70,50 @@ async def test_add_deployment_without_salt_key_or_master_key(): such as in a local/dev environment or when just saving spend logs. """ # Remove LITELLM_SALT_KEY from environment - old_salt_key = os.environ.pop("LITELLM_SALT_KEY", None) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) - try: - # Set master_key to None - with patch("litellm.proxy.proxy_server.master_key", None): - # Mock the required dependencies - mock_prisma_client = MagicMock(spec=PrismaClient) - mock_prisma_client.db = MagicMock() - mock_prisma_client.db.litellm_config = MagicMock() - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=None + # Set master_key to None + with patch("litellm.proxy.proxy_server.master_key", None): + # Mock the required dependencies + mock_prisma_client = MagicMock(spec=PrismaClient) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_config = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock( + return_value=None + ) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + + # Create ProxyConfig instance + proxy_config = ProxyConfig() + + # Mock the internal methods + proxy_config._should_load_db_object = MagicMock(return_value=False) + proxy_config._init_non_llm_objects_in_db = AsyncMock() + + # This should NOT raise an exception + try: + await proxy_config.add_deployment( + prisma_client=mock_prisma_client, + proxy_logging_obj=mock_proxy_logging, ) - - mock_proxy_logging = MagicMock(spec=ProxyLogging) - - # Create ProxyConfig instance - proxy_config = ProxyConfig() - - # Mock the internal methods - proxy_config._should_load_db_object = MagicMock(return_value=False) - proxy_config._init_non_llm_objects_in_db = AsyncMock() - - # This should NOT raise an exception - try: - await proxy_config.add_deployment( - prisma_client=mock_prisma_client, - proxy_logging_obj=mock_proxy_logging, + assert True + except ValueError as e: + if "Master key is not initialized" in str( + e + ) or "Encryption key is not initialized" in str(e): + pytest.fail( + f"add_deployment raised ValueError about encryption key: {e}" ) - assert True - except ValueError as e: - if "Master key is not initialized" in str( - e - ) or "Encryption key is not initialized" in str(e): - pytest.fail( - f"add_deployment raised ValueError about encryption key: {e}" - ) - raise - except Exception as e: - if "Master key is not initialized" in str( - e - ) or "Encryption key is not initialized" in str(e): - pytest.fail( - f"add_deployment raised exception about encryption key: {e}" - ) - raise - finally: - # Restore LITELLM_SALT_KEY if it was set - if old_salt_key: - os.environ["LITELLM_SALT_KEY"] = old_salt_key + raise + except Exception as e: + if "Master key is not initialized" in str( + e + ) or "Encryption key is not initialized" in str(e): + pytest.fail( + f"add_deployment raised exception about encryption key: {e}" + ) + raise def test_add_deployment_sync_without_master_key(): diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/test_litellm/test_count_tokens_public_api.py index 1e2cf83dec0..ebd9c0c9edb 100644 --- a/tests/test_litellm/test_count_tokens_public_api.py +++ b/tests/test_litellm/test_count_tokens_public_api.py @@ -144,20 +144,16 @@ def test_acount_tokens_api_error_falls_back(): assert result.total_tokens > 0 -def test_acount_tokens_no_api_key_falls_back(): +def test_acount_tokens_no_api_key_falls_back(monkeypatch): """Test that missing API key falls back to local counting.""" - env_backup = os.environ.pop("OPENAI_API_KEY", None) - try: - result = asyncio.run( - litellm.acount_tokens( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - ) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + result = asyncio.run( + litellm.acount_tokens( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], ) + ) - # Should fall back to local tokenizer since no API key - assert result.total_tokens > 0 - assert result.tokenizer_type == "local_tokenizer" - finally: - if env_backup: - os.environ["OPENAI_API_KEY"] = env_backup + # Should fall back to local tokenizer since no API key + assert result.total_tokens > 0 + assert result.tokenizer_type == "local_tokenizer" diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index ba82bfaadc6..dd19334724d 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -318,7 +318,7 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp litellm.model_cost.pop(model_key, None) -def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(): +def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypatch): """Registering a custom override under a key shape that ``get_model_info`` cannot resolve (e.g. a triple provider prefix like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6``; a double @@ -338,7 +338,7 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(): from litellm.types.utils import PromptTokensDetailsWrapper, Usage original_model_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") builtin_key = "us.anthropic.claude-sonnet-4-6" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index bd23ca11fbe..cd8dad39ad5 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -672,8 +672,8 @@ def test_all_model_configs(): ) == {"max_output_tokens": 10} -def test_anthropic_web_search_in_model_info(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_anthropic_web_search_in_model_info(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") supported_models = [ @@ -1193,11 +1193,11 @@ def test_max_tokens_consistency(): raise AssertionError(error_msg) -def test_get_model_info_gemini(): +def test_get_model_info_gemini(monkeypatch): """ Tests if ALL gemini models have 'tpm' and 'rpm' in the model info """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model_map = litellm.model_cost @@ -1252,8 +1252,8 @@ def test_get_model_info_bedrock_double_provider_prefix_resolves(local_model_cost assert info["key"] == "us.anthropic.claude-sonnet-4-6" -def test_openai_models_in_model_info(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" +def test_openai_models_in_model_info(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") model_map = litellm.model_cost @@ -1408,7 +1408,7 @@ for commitment in BEDROCK_COMMITMENTS: print("block_list", block_list) -def test_supports_computer_use_utility(): +def test_supports_computer_use_utility(monkeypatch): """ Tests the litellm.utils.supports_computer_use utility function. """ @@ -1420,7 +1420,7 @@ def test_supports_computer_use_utility(): original_env_var = os.getenv("LITELLM_LOCAL_MODEL_COST_MAP") original_model_cost = getattr(litellm, "model_cost", None) - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") # Load with local/backup try: @@ -1438,7 +1438,7 @@ def test_supports_computer_use_utility(): if original_env_var is None: del os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env_var + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", original_env_var) if original_model_cost is not None: litellm.model_cost = original_model_cost @@ -1446,13 +1446,13 @@ def test_supports_computer_use_utility(): delattr(litellm, "model_cost") -def test_get_model_info_shows_supports_computer_use(): +def test_get_model_info_shows_supports_computer_use(monkeypatch): """ Tests if 'supports_computer_use' is correctly retrieved by get_model_info. We'll use 'claude-4-sonnet-20250514' as it's configured in the backup JSON to have supports_computer_use: True. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") # Ensure litellm.model_cost is loaded, relying on the backup mechanism if primary fails # as per previous debugging. litellm.model_cost = litellm.get_model_cost_map(url="") From 74816498303d18451d7452fbcabe3042e1ba7de2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 20:39:07 -0700 Subject: [PATCH 30/43] test(datadog): restore an empty DD_API_KEY instead of unsetting it (#37832) Both datadog test files hand-roll what monkeypatch.setenv already does: read the old value, write the test value, put the old one back on the way out. The cost management fixture checks the old value for truthiness rather than for None, so an operator running the suite with DD_API_KEY set to the empty string gets it deleted rather than restored. Starting from DD_API_KEY="" and running test_init leaves it None on the current file, and "" after this. 13 raw os.environ writes become monkeypatch.setenv, the two fixtures stop being yield fixtures because there is nothing left to do on the way out, and the now unused os import goes with them. 27 tests pass across the two files, 88 across tests/test_litellm/integrations/datadog. --- test-quality-budget.json | 2 +- .../datadog/test_datadog_cost_management.py | 36 ++++--------------- .../datadog/test_datadog_metrics.py | 30 ++++++---------- 3 files changed, 18 insertions(+), 50 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 143378efec7..68244dce319 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -9,7 +9,7 @@ "limit": 1078 }, "TQ004": { - "limit": 557 + "limit": 544 }, "TQ005": { "limit": 2810 diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py index cb786d9c292..1a50a6991da 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py @@ -1,4 +1,3 @@ -import os import time from unittest.mock import AsyncMock @@ -12,34 +11,13 @@ from litellm.types.utils import StandardLoggingPayload @pytest.fixture -def clean_env(): - # Save original env - original_api_key = os.environ.get("DD_API_KEY") - original_app_key = os.environ.get("DD_APP_KEY") - original_site = os.environ.get("DD_SITE") - - # Set test env - os.environ["DD_API_KEY"] = "test_api_key" - os.environ["DD_APP_KEY"] = "test_app_key" - os.environ["DD_SITE"] = "test.datadoghq.com" - - yield - - # Restore original env - if original_api_key: - os.environ["DD_API_KEY"] = original_api_key - else: - del os.environ["DD_API_KEY"] - - if original_app_key: - os.environ["DD_APP_KEY"] = original_app_key - else: - del os.environ["DD_APP_KEY"] - - if original_site: - os.environ["DD_SITE"] = original_site - else: - del os.environ["DD_SITE"] +def clean_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key, value in ( + ("DD_API_KEY", "test_api_key"), + ("DD_APP_KEY", "test_app_key"), + ("DD_SITE", "test.datadoghq.com"), + ): + monkeypatch.setenv(key, value) @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py index a4a4ca334b0..eade92d6672 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_metrics.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_metrics.py @@ -1,4 +1,3 @@ -import os import time from datetime import datetime, timedelta from unittest.mock import AsyncMock @@ -11,25 +10,16 @@ from litellm.types.utils import StandardLoggingPayload @pytest.fixture -def clean_env(): - """Set test env vars and restore originals after test.""" - keys = ["DD_API_KEY", "DD_APP_KEY", "DD_SITE", "DD_ENV", "DD_SERVICE", "DD_VERSION"] - originals = {k: os.environ.get(k) for k in keys} - - os.environ["DD_API_KEY"] = "test_api_key" - os.environ["DD_APP_KEY"] = "test_app_key" - os.environ["DD_SITE"] = "test.datadoghq.com" - os.environ["DD_ENV"] = "test-env" - os.environ["DD_SERVICE"] = "test-service" - os.environ["DD_VERSION"] = "1.0.0" - - yield - - for k, v in originals.items(): - if v is not None: - os.environ[k] = v - elif k in os.environ: - del os.environ[k] +def clean_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key, value in ( + ("DD_API_KEY", "test_api_key"), + ("DD_APP_KEY", "test_app_key"), + ("DD_SITE", "test.datadoghq.com"), + ("DD_ENV", "test-env"), + ("DD_SERVICE", "test-service"), + ("DD_VERSION", "1.0.0"), + ): + monkeypatch.setenv(key, value) @pytest.mark.asyncio From 0c97eea66073231f37c8ae2bac4f4f18210d454d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 21:00:04 -0700 Subject: [PATCH 31/43] test(cost-calc): stop 182 global writes leaking out of the cost-calc suites (#37815) * test(cost-calc): stop 182 global writes leaking out of the cost-calc suites Across test_cost_calculator.py and llm_cost_calc/test_llm_cost_calc_utils.py, 58 tests opened by setting LITELLM_LOCAL_MODEL_COST_MAP in os.environ and replacing litellm.model_cost, and none of them put the env var back. The second file already had a _local_model_cost_map fixture doing it by hand with a try/finally, so both idioms sat in the same file. Keep that fixture, give it monkeypatch, and have every one of those tests ask for it. The margin and discount tests drop their hand-rolled copy-then-restore in favour of monkeypatch.setattr, which also puts the global back when an assertion fails part way through. Both files also drop a sys.path.insert whose argument resolves outside the repo, so it was never what made the imports work. TQ003 1077 -> 1075, TQ004 768 -> 693, TQ005 2836 -> 2731, and the budget ceilings come down with them. * fix(test): make the streamed-cost tests load the map they assert against The local_cost_map fixture set LITELLM_LOCAL_MODEL_COST_MAP but never reloaded litellm.model_cost, and reading the variable is not what loads the map. So the three streaming-cost tests billed against whatever map the process happened to be holding, and their hardcoded prices only held when something else had already swapped in the checked-in one. This branch stops the cost-calc tests leaking that map, which left test_main billing at the ambient prices instead. The fixture now loads the map it names, so the prices these tests assert hold on their own. --- test-quality-budget.json | 6 +- .../llm_cost_calc/test_llm_cost_calc_utils.py | 209 +++++------------ tests/test_litellm/test_cost_calculator.py | 218 +++++------------- tests/test_litellm/test_main.py | 3 + 4 files changed, 121 insertions(+), 315 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 68244dce319..6428a55ba78 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -6,13 +6,13 @@ "limit": 742 }, "TQ003": { - "limit": 1078 + "limit": 1075 }, "TQ004": { - "limit": 544 + "limit": 469 }, "TQ005": { - "limit": 2810 + "limit": 2661 }, "TQ006": { "limit": 34 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 e6170d47a6c..c8c36032793 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 @@ -1,6 +1,4 @@ import json -import os -import sys import pytest from fastapi.testclient import TestClient @@ -28,10 +26,6 @@ from litellm.types.utils import ( StandardBuiltInToolsParams, ) -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - from litellm.litellm_core_utils.llm_cost_calc.utils import ( PromptTokensDetailsResult, TokenTypeCostBreakdown, @@ -44,13 +38,17 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( from litellm.types.utils import CacheCreationTokenDetails, Usage -def test_reasoning_tokens_no_price_set(): +@pytest.fixture +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + +def test_reasoning_tokens_no_price_set(_local_model_cost_map): # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token) model = "o1" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] usage = Usage( completion_tokens=1578, @@ -87,11 +85,9 @@ def test_reasoning_tokens_no_price_set(): ) -def test_reasoning_tokens_gemini(): +def test_reasoning_tokens_gemini(_local_model_cost_map): model = "gemini-2.5-flash" custom_llm_provider = "gemini" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( completion_tokens=1578, @@ -132,12 +128,10 @@ def test_reasoning_tokens_gemini(): ) -def test_reasoning_tokens_gemini_3_1_flash_lite(): +def test_reasoning_tokens_gemini_3_1_flash_lite(_local_model_cost_map): """Test cost calculation for gemini-3.1-flash-lite-preview with reasoning tokens""" model = "gemini-3.1-flash-lite-preview" custom_llm_provider = "gemini" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( completion_tokens=1000, @@ -270,11 +264,9 @@ def test_image_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round(expected_completion_cost, 12) -def test_video_output_tokens_gemini_omni_flash_preview(): +def test_video_output_tokens_gemini_omni_flash_preview(_local_model_cost_map): """Video output tokens are billed at output_cost_per_video_token, not the text rate and not zero.""" model = "gemini-omni-flash-preview" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") text_tokens = 100 video_tokens = 46336 @@ -310,11 +302,9 @@ def test_video_output_tokens_gemini_omni_flash_preview(): ) -def test_video_input_tokens_gemini_omni_flash_preview(): +def test_video_input_tokens_gemini_omni_flash_preview(_local_model_cost_map): """Video input tokens are billed at the standard input rate instead of being dropped.""" model = "gemini-omni-flash-preview" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( completion_tokens=10, @@ -369,12 +359,10 @@ def test_video_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12) -def test_generic_cost_per_token_above_200k_tokens(): +def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map): # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing model = "gemini-2.5-pro" custom_llm_provider = "vertex_ai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] prompt_tokens = 220 * 1e6 @@ -420,12 +408,10 @@ def test_get_token_base_cost_picks_highest_crossed_tier(): assert prompt_base_cost == 9e-6 -def test_generic_cost_per_token_gpt54_above_272k_tokens(): +def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] prompt_tokens = 273000 # Above 272K threshold @@ -450,12 +436,10 @@ def test_generic_cost_per_token_gpt54_above_272k_tokens(): assert round(completion_cost, 10) == round(expected_completion, 10) -def test_generic_cost_per_token_minimax_m3_above_512k_tokens(): +def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_map): """MiniMax-M3: prompts >512K input tokens priced at 2x input, output, and cache read.""" model = "minimax/MiniMax-M3" custom_llm_provider = "minimax" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] prompt_tokens = 600000 @@ -493,10 +477,8 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(): "bedrock_mantle/openai.gpt-5.6-luna", ], ) -def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(model): +def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model): """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] assert model_cost_map["max_input_tokens"] == 1000000 @@ -827,12 +809,10 @@ def test_generic_cost_per_token_tiered_pricing_bills_reasoning_at_tier_rate(): litellm.model_cost.pop(model, None) -def test_generic_cost_per_token_gpt55(): +def test_generic_cost_per_token_gpt55(_local_model_cost_map): """gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input.""" model = "gpt-5.5" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] @@ -867,12 +847,10 @@ def test_generic_cost_per_token_gpt55(): ) -def test_generic_cost_per_token_gpt55_pro(): +def test_generic_cost_per_token_gpt55_pro(_local_model_cost_map): """gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input.""" model = "gpt-5.5-pro" custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] @@ -919,7 +897,7 @@ def test_generic_cost_per_token_gpt55_pro(): ("gpt-5.6-luna", 2e-7, 1.2e-6, 2e-8, 2.5e-7), ], ) -def test_generic_cost_per_token_gpt56( +def test_generic_cost_per_token_gpt56(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost, cache_write_cost ): """gpt-5.6 (sol/terra/luna): base pricing + new cache-write cost. @@ -927,8 +905,6 @@ def test_generic_cost_per_token_gpt56( Cache writes are billed at 1.25x the uncached input rate for this family. """ custom_llm_provider = "openai" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] @@ -989,7 +965,7 @@ def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map): ("gpt-5.6-luna", 2e-7, 9e-7), ], ) -def test_generic_cost_per_token_gpt56_flex_above_272k( +def test_generic_cost_per_token_gpt56_flex_above_272k(_local_model_cost_map, model, flex_long_input_cost, flex_long_output_cost ): """A >272K flex request bills the flex long-context rate, not the standard one. @@ -998,8 +974,6 @@ def test_generic_cost_per_token_gpt56_flex_above_272k( ``*_above_272k_tokens_flex`` keys these requests silently fell back to the standard long-context price, billing 2x what OpenAI charges. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") prompt_tokens = 300000 completion_tokens = 1000 @@ -1038,11 +1012,9 @@ def test_generic_cost_per_token_gpt56_flex_above_272k( ("flex", 300000, 2e-6, 2.5e-6, 2e-7), ], ) -def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context( +def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context(_local_model_cost_map, service_tier, prompt_tokens, input_rate, cache_write_rate, cache_read_rate ): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") cached_tokens = 50000 cache_write_tokens = 40000 @@ -1130,7 +1102,7 @@ def test_generic_cost_per_token_gpt56_cyber( ("azure/eu/gpt-5.6-luna", 2.2e-7, 1.32e-6, 2.2e-8), ], ) -def test_generic_cost_per_token_azure_gpt56( +def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost ): """Azure gpt-5.6 (global + us/eu regional): Azure prices this family on its own @@ -1138,8 +1110,6 @@ def test_generic_cost_per_token_azure_gpt56( promotional cut OpenAI applied to gpt-5.6-sol, so these rates deliberately sit above the openai ones and must not be lowered to match them. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_cost_map = litellm.model_cost[model] assert model_cost_map["litellm_provider"] == "azure" @@ -1180,7 +1150,7 @@ def test_generic_cost_per_token_azure_gpt56( ("gpt-5.5-pro-2026-04-23", False, True, False), ], ) -def test_gpt55_reasoning_effort_flags_match_live_openai_api( +def test_gpt55_reasoning_effort_flags_match_live_openai_api(_local_model_cost_map, model, expected_none, expected_xhigh, expected_minimal ): """Pin reasoning_effort capability flags to OpenAI's actual API contract. @@ -1189,8 +1159,6 @@ def test_gpt55_reasoning_effort_flags_match_live_openai_api( ``Unsupported value: 'reasoning_effort' does not support 'minimal' with this model``. gpt-5.5-pro additionally rejects 'none' and 'low'. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") m = litellm.model_cost[model] assert ( @@ -1211,7 +1179,7 @@ def test_gpt55_reasoning_effort_flags_match_live_openai_api( ("gpt-5.5-pro", "gpt-5.5-pro-2026-04-23"), ], ) -def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( +def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(_local_model_cost_map, base_model, dated_model ): """Dated snapshots must carry the same reasoning_effort capability flags as @@ -1223,8 +1191,6 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( behavior between ``gpt-5.5`` and ``gpt-5.5-2026-04-23``. Pinning to a dated variant must never lose capabilities relative to the base alias. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") base = litellm.model_cost[base_model] dated = litellm.model_cost[dated_model] @@ -1251,7 +1217,7 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( ("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6), ], ) -def test_azure_gpt55_entries_present_with_correct_pricing( +def test_azure_gpt55_entries_present_with_correct_pricing(_local_model_cost_map, model, expected_mode, expected_input, expected_output, expected_cache_read ): """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure. @@ -1260,8 +1226,6 @@ def test_azure_gpt55_entries_present_with_correct_pricing( on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro. Cache discount is 10% of input. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") m = litellm.model_cost[model] assert m["litellm_provider"] == "azure" @@ -1286,12 +1250,10 @@ def test_azure_gpt55_entries_present_with_correct_pricing( ("azure/gpt-5.5-pro", False, False, True), ], ) -def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api( +def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api(_local_model_cost_map, model, expected_none, expected_minimal, expected_xhigh ): """Azure entries pin reasoning_effort flags to OpenAI's actual API contract.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") m = litellm.model_cost[model] assert m.get("supports_none_reasoning_effort") is expected_none @@ -1671,11 +1633,9 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): assert round(result, 6) == round(expected, 6) -def test_service_tier_flex_pricing(): +def test_service_tier_flex_pricing(_local_model_cost_map): """Test that flex service tier uses correct pricing (approximately 50% of standard).""" # Set up environment for local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -1728,11 +1688,9 @@ def test_service_tier_flex_pricing(): ), f"Flex total cost mismatch: {flex_total} vs {expected_flex_total}" -def test_service_tier_default_pricing(): +def test_service_tier_default_pricing(_local_model_cost_map): """Test that when no service tier is provided, standard pricing is used.""" # Set up environment for local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano model = "gpt-5-nano" @@ -1779,11 +1737,9 @@ def test_service_tier_default_pricing(): ), f"Standard completion cost mismatch: {default_cost[1]} vs {expected_standard_completion}" -def test_service_tier_fallback_pricing(): +def test_service_tier_fallback_pricing(_local_model_cost_map): """Test that when service tier is provided but model doesn't have those keys, it falls back to standard pricing.""" # Set up environment for local model cost map - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-4 which doesn't have flex pricing keys model = "gpt-4" @@ -1891,15 +1847,13 @@ def test_service_tier_ultrafast_pricing(): assert completion_cost == pytest.approx(400 * 3e-04) -def test_service_tier_ultrafast_fallback_pricing(): +def test_service_tier_ultrafast_fallback_pricing(_local_model_cost_map): """Without *_ultrafast keys an ultrafast request bills the standard rate, not zero. Guards the suffix fallback in _get_cost_per_unit: "_fast" is a substring of "_ultrafast", so a shortest-first suffix match would strip the wrong suffix and price the request at 0. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) @@ -1929,7 +1883,7 @@ def test_service_tier_ultrafast_fallback_pricing(): "gemini-3.1-flash-lite-image", ], ) -def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): +def test_gemini_image_generation_cost_with_zero_text_tokens(_local_model_cost_map, model: str): """ Test that image_tokens are correctly costed when text_tokens=0. @@ -1939,8 +1893,6 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): https://github.com/BerriAI/litellm/issues/17410 """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") custom_llm_provider = "vertex_ai" @@ -1995,13 +1947,11 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str): ), f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}" -def test_vertex_image_generation_cost_prefers_token_usage_metadata(): +def test_vertex_image_generation_cost_prefers_token_usage_metadata(_local_model_cost_map): """ When usage metadata exists on image responses, Vertex image generation cost should be calculated from token pricing, not flat output_cost_per_image. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3.1-flash-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") @@ -2040,13 +1990,11 @@ def test_vertex_image_generation_cost_prefers_token_usage_metadata(): assert cost != len(image_response.data) * model_info["output_cost_per_image"] -def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(): +def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(_local_model_cost_map): """ Without usage metadata, Vertex image generation cost should fall back to output_cost_per_image * number_of_images. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-3.1-flash-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") @@ -2064,13 +2012,11 @@ def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing(): assert round(cost, 10) == round(expected_cost, 10) -def test_gemini_image_generation_cost_prefers_token_usage_metadata(): +def test_gemini_image_generation_cost_prefers_token_usage_metadata(_local_model_cost_map): """ When usage metadata exists on image responses, Gemini image generation cost should be calculated from token pricing, not flat output_cost_per_image. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -2109,13 +2055,11 @@ def test_gemini_image_generation_cost_prefers_token_usage_metadata(): assert cost != len(image_response.data) * model_info["output_cost_per_image"] -def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(): +def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing(_local_model_cost_map): """ Without usage metadata, Gemini image generation cost should fall back to output_cost_per_image * number_of_images. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini/gemini-3-pro-image-preview" model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") @@ -2212,7 +2156,7 @@ def test_reasoning_tokens_without_text_tokens_gpt5_nano(): ), "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!" -def test_image_count_prevents_text_tokens_fallback(): +def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map): """ Test that the text_tokens fallback in generic_cost_per_token does not override text_tokens=0 when image_count > 0. @@ -2221,8 +2165,6 @@ def test_image_count_prevents_text_tokens_fallback(): When image_count > 0, text_tokens=0 is intentional (image-only request), not "text_tokens not set by provider." """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Simulate Nova image-only embedding: prompt_tokens estimated from # embedding dimensions (768 for 3072-dim), image_count=1 @@ -2256,20 +2198,6 @@ def test_image_count_prevents_text_tokens_fallback(): # --------------------------------------------------------------------------- -@pytest.fixture -def _local_model_cost_map(): - prev_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - prev_model_cost = litellm.model_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - try: - yield - finally: - litellm.model_cost = prev_model_cost - if prev_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = prev_env @pytest.mark.parametrize("model", ["gpt-5.4", "gpt-realtime-2.1", "gpt-realtime-2.1-mini"]) @@ -2603,7 +2531,7 @@ def test_threshold_keys_exclude_service_tier_variants(): ("cerebras/qwen-3-32b", "cerebras", 250, 0), ], ) -def test_token_type_cost_breakdown_is_provider_agnostic( +def test_token_type_cost_breakdown_is_provider_agnostic(_local_model_cost_map, model, custom_llm_provider, reasoning_tokens, cached_tokens ): """ @@ -2615,8 +2543,6 @@ def test_token_type_cost_breakdown_is_provider_agnostic( there - not the top-level cache_read_input_tokens attribute the old breakdown code relied on - is what makes Vertex/OpenAI/Azure cache costs show up at all. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( prompt_tokens=1000, @@ -2647,10 +2573,8 @@ def test_token_type_cost_breakdown_is_provider_agnostic( assert breakdown.cache_read_cost == pytest.approx(cached_tokens * cache_read_rate) -def test_token_type_cost_breakdown_matches_real_gemini_numbers(): +def test_token_type_cost_breakdown_matches_real_gemini_numbers(_local_model_cost_map): """Hard-coded against the exact gemini-2.5-flash response that exposed the gap.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( prompt_tokens=209, @@ -2673,9 +2597,7 @@ def test_token_type_cost_breakdown_matches_real_gemini_numbers(): assert breakdown.cache_creation_cost == 0.0 -def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(_local_model_cost_map): usage = Usage( prompt_tokens=200_000, @@ -2697,9 +2619,7 @@ def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(): assert breakdown.cache_read_cost == pytest.approx(50_000 * 4e-07) -def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(_local_model_cost_map): usage = Usage( prompt_tokens=199_999, @@ -2721,14 +2641,12 @@ def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(): assert breakdown.cache_read_cost == pytest.approx(50_000 * 2e-07) -def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(): +def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(_local_model_cost_map): """ Bedrock/Anthropic report cache tokens as top-level usage fields; the Usage constructor maps them onto prompt_tokens_details, so the breakdown must still pick up both cache-read and cache-creation costs. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "anthropic.claude-3-5-haiku-20241022-v1:0" usage = Usage( @@ -2752,14 +2670,12 @@ def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage( ) -def test_token_type_cost_breakdown_reads_cache_write_tokens(): +def test_token_type_cost_breakdown_reads_cache_write_tokens(_local_model_cost_map): """ Some OpenAI-compatible providers (e.g. kimi-k2) report cache-write tokens under `cache_write_tokens` rather than `cache_creation_tokens`. The breakdown must read it the same way the total-cost normalization does, so the two agree. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "anthropic.claude-3-5-haiku-20241022-v1:0" usage = Usage( @@ -2780,7 +2696,7 @@ def test_token_type_cost_breakdown_reads_cache_write_tokens(): ) -def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): +def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(_local_model_cost_map): """ Regression: OpenAI gpt-5.6 reports cache-write tokens under prompt_tokens_details.cache_write_tokens (not the Anthropic cache_creation_tokens @@ -2788,8 +2704,6 @@ def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): input rate. Customer report: cache creation tokens were never counted for the GPT-5.6 series, so cost was undercounted on cache-write requests. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.6" usage = Usage( @@ -2811,14 +2725,12 @@ def test_generic_cost_per_token_openai_cache_write_tokens_gpt_5_6(): assert prompt_cost > 1000 * info["input_cost_per_token"] -def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(): +def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(_local_model_cost_map): """ Regression for #34801: when a provider reports text_tokens covering the whole prompt alongside cache-write tokens (and no cache reads), the cache-write tokens must be backed out of the text total instead of being billed twice. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.6" usage = Usage( @@ -2837,15 +2749,13 @@ def test_generic_cost_per_token_backs_out_cache_write_tokens_from_text_tokens(): assert prompt_cost == pytest.approx(expected_prompt) -def test_token_type_cost_breakdown_reconciles_with_generic_total(): +def test_token_type_cost_breakdown_reconciles_with_generic_total(_local_model_cost_map): """ Both-ways check: the reasoning subset must sum with the remaining (text) output cost to exactly the completion total, and the cache-read subset with the remaining input cost to exactly the prompt total, as computed by generic_cost_per_token. A mismatch here would mean the breakdown misrepresents what was actually billed. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gemini-2.5-flash" custom_llm_provider = "vertex_ai" @@ -2878,9 +2788,7 @@ def test_token_type_cost_breakdown_reconciles_with_generic_total(): assert text_input_cost + breakdown.cache_read_cost == pytest.approx(prompt_cost) -def test_token_type_cost_breakdown_zero_without_special_tokens(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_token_type_cost_breakdown_zero_without_special_tokens(_local_model_cost_map): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) breakdown = get_token_type_cost_breakdown( @@ -2917,7 +2825,7 @@ def test_token_type_cost_breakdown_zero_without_special_tokens(): ), ], ) -def test_token_type_cost_breakdown_openai_responses_api_cache_write_read( +def test_token_type_cost_breakdown_openai_responses_api_cache_write_read(_local_model_cost_map, raw_usage, expect_read, expect_write ): """Regression for #34309: OpenAI Responses API reports cache tokens under @@ -2926,8 +2834,6 @@ def test_token_type_cost_breakdown_openai_responses_api_cache_write_read( cache_read_cost / cache_creation_cost from the transformed usage.""" from litellm.responses.utils import ResponseAPILoggingUtils - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.6" usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage) @@ -2968,15 +2874,13 @@ def test_token_type_cost_breakdown_handles_unknown_model_gracefully(): ) -def test_token_type_cost_breakdown_applies_regional_uplift(): +def test_token_type_cost_breakdown_applies_regional_uplift(_local_model_cost_map): """ Regional OpenAI hosts (eu./us.) apply a flat uplift to every token cost. The per-type breakdown must apply the same uplift via data_residency so it stays reconciled with the uplifted input_cost/output_cost totals, instead of being logged at the base rate. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "gpt-5.4" custom_llm_provider = "openai" @@ -3024,15 +2928,13 @@ def test_token_type_cost_breakdown_applies_regional_uplift(): assert text_input_cost + eu.cache_read_cost == pytest.approx(prompt_cost) -def test_token_type_cost_breakdown_applies_vertex_regional_uplift(): +def test_token_type_cost_breakdown_applies_vertex_regional_uplift(_local_model_cost_map): """ Non-global Vertex endpoints apply a flat 1.1x uplift to every token cost. The per-type breakdown must apply the same uplift via vertex_location so it stays reconciled with the uplifted input_cost/output_cost totals, instead of being logged at the global rate. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-haiku-4-5@20251001" custom_llm_provider = "vertex_ai" @@ -3075,7 +2977,7 @@ def test_token_type_cost_breakdown_applies_vertex_regional_uplift(): assert text_input_cost + regional.cache_read_cost == pytest.approx(prompt_cost) -def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch): +def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(_local_model_cost_map, monkeypatch): """ Anthropic's regional (geo) uplift lives in provider_specific_entry and is applied to every token type in the totals, so the per-type breakdown must @@ -3088,7 +2990,6 @@ def test_token_type_cost_breakdown_applies_anthropic_geo_multiplier(monkeypatch) ) monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-geo-breakdown-model" litellm.register_model( @@ -3209,9 +3110,7 @@ GEMINI_DAY0_LAUNCH_PRICING = [ @pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_DAY0_LAUNCH_PRICING) -def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, output_cost, cache_read_cost): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_gemini_36_flash_and_35_flash_lite_launch_pricing(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost): model_cost_map = litellm.model_cost[model] assert model_cost_map["input_cost_per_token"] == input_cost @@ -3224,9 +3123,7 @@ def test_gemini_36_flash_and_35_flash_lite_launch_pricing(model, input_cost, out assert model_cost_map["max_input_tokens"] == 1048576 -def test_generic_cost_per_token_gemini_36_flash(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_generic_cost_per_token_gemini_36_flash(_local_model_cost_map): usage = Usage( prompt_tokens=1000, @@ -3292,9 +3189,7 @@ def test_gemini_36_flash_batch_introductory_pricing(model, _local_model_cost_map assert model_cost_map["output_cost_per_token_batches"] == 1.875e-06 -def test_generic_cost_per_token_gemini_35_flash_lite(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map): usage = Usage( prompt_tokens=1000, diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 2b30138faa2..8dad4bef07b 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,12 +1,6 @@ -import os -import sys import pytest -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path - from pydantic import BaseModel @@ -24,6 +18,12 @@ from litellm.types.utils import ModelInfo, ModelResponse, PromptTokensDetailsWra from litellm.utils import TranscriptionResponse +@pytest.fixture +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch): """ Router/proxy configs may use deployment ids like openai/openai/. Cost lookup must @@ -93,14 +93,12 @@ def test_cost_per_token_non_string_model_does_not_hang(): assert result.get("status") in ("returned", "raised") -def test_completion_cost_uses_response_model_for_dynamic_routing(): +def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_cost_map): """ Test that completion_cost uses the model from the response object when the input model (e.g., azure-model-router) is not in model_cost. This supports Azure Model Router and similar dynamic routing scenarios. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Simulate Azure Model Router: input is generic router, response has actual model response = ModelResponse( @@ -139,9 +137,7 @@ def test_cost_calculator_with_response_cost_in_additional_headers(): assert result == 1000 -def test_baseten_model_api_pricing_entries(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_baseten_model_api_pricing_entries(_local_model_cost_map): expected_pricing = { "baseten/nvidia/Nemotron-120B-A12B": (3e-07, 7.5e-07), @@ -165,9 +161,7 @@ def test_baseten_model_api_pricing_entries(): assert model_info["output_cost_per_token"] == output_cost -def test_wandb_model_api_pricing_entries(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_wandb_model_api_pricing_entries(_local_model_cost_map): expected_pricing = { "wandb/moonshotai/Kimi-K2.5": (6e-07, 3e-06), @@ -182,9 +176,7 @@ def test_wandb_model_api_pricing_entries(): assert model_info["output_cost_per_token"] == output_cost -def test_openrouter_qwen36_plus_model_info(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_openrouter_qwen36_plus_model_info(_local_model_cost_map): model_info = litellm.model_cost.get("openrouter/qwen/qwen3.6-plus") @@ -208,9 +200,7 @@ def test_openrouter_qwen36_plus_model_info(): "github_copilot/mai-code-1-flash-internal", ], ) -def test_github_copilot_mai_code_1_flash_pricing(model): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_github_copilot_mai_code_1_flash_pricing(_local_model_cost_map, model): model_info = litellm.model_cost.get(model) @@ -238,9 +228,7 @@ def test_github_copilot_mai_code_1_flash_pricing(model): assert completion_usd == pytest.approx(500 * 4.5e-06) -def test_cost_calculator_with_usage(monkeypatch): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): usage = Usage( prompt_tokens=120, @@ -320,11 +308,9 @@ def test_cost_calculator_with_usage(monkeypatch): assert result == expected_cost, f"Got {result}, Expected {expected_cost}" -def test_transcription_cost_uses_token_pricing(): +def test_transcription_cost_uses_token_pricing(_local_model_cost_map): from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") usage = Usage( prompt_tokens=14, @@ -348,11 +334,9 @@ def test_transcription_cost_uses_token_pricing(): assert pytest.approx(cost, rel=1e-6) == expected_cost -def test_transcription_cost_falls_back_to_duration(): +def test_transcription_cost_falls_back_to_duration(_local_model_cost_map): from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") response = TranscriptionResponse(text="demo text") response.duration = 10.0 @@ -368,14 +352,12 @@ def test_transcription_cost_falls_back_to_duration(): assert pytest.approx(cost, rel=1e-6) == expected_cost -def test_vertex_chirp_3_transcription_cost_from_duration(): +def test_vertex_chirp_3_transcription_cost_from_duration(_local_model_cost_map): """Regression: the chirp_3 cost map entry shipped with output_cost_per_second 0.0, and cost_per_second prefers output_cost_per_second whenever it is not None, so every transcription priced to $0.00 instead of using input_cost_per_second.""" from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") response = TranscriptionResponse(text="demo text") response.duration = 18.0 @@ -1127,9 +1109,7 @@ def test_tiered_pricing_only_deployment_completion_cost_is_nonzero(): assert cost > 0 -def test_azure_realtime_cost_calculator(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_azure_realtime_cost_calculator(_local_model_cost_map): cost = handle_realtime_stream_cost_calculation( results=[ @@ -1152,7 +1132,7 @@ def test_azure_realtime_cost_calculator(): assert cost > 0 -def test_azure_audio_output_cost_calculation(): +def test_azure_audio_output_cost_calculation(_local_model_cost_map): """ Test that Azure audio models correctly calculate costs for audio output tokens. @@ -1162,8 +1142,6 @@ def test_azure_audio_output_cost_calculation(): """ from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Scenario from issue #19764: # Input: 17 text tokens, 0 audio tokens @@ -1672,7 +1650,7 @@ def test_gemini_25_explicit_caching_cost_direct_usage(): assert expected_actual_cost == total_cost -def test_azure_ai_cache_cost_calculation(): +def test_azure_ai_cache_cost_calculation(_local_model_cost_map): """ Test that azure_ai provider correctly calculates cache costs using generic_cost_per_token. @@ -1683,8 +1661,6 @@ def test_azure_ai_cache_cost_calculation(): from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import PromptTokensDetailsWrapper, Usage - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Register a custom azure_ai model with cache pricing test_model_id = "test-azure-ai-claude-model" @@ -1817,15 +1793,13 @@ def test_vertex_uplift_composes_with_above_128k_pricing(monkeypatch): assert regional_completion == pytest.approx(global_completion * 1.10, rel=1e-9) -def test_cost_discount_vertex_ai(): +def test_cost_discount_vertex_ai(monkeypatch): """ Test that cost discount is applied correctly for Vertex AI provider """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_discount_config = litellm.cost_discount_config.copy() # Create mock response (use a model that exists in model_prices_and_context_window.json) response = ModelResponse( @@ -1838,7 +1812,7 @@ def test_cost_discount_vertex_ai(): ) # Calculate cost without discount - litellm.cost_discount_config = {} + monkeypatch.setattr(litellm, "cost_discount_config", {}) cost_without_discount = completion_cost( completion_response=response, model="vertex_ai/gemini-3-pro-preview", @@ -1846,7 +1820,7 @@ def test_cost_discount_vertex_ai(): ) # Set 5% discount for vertex_ai - litellm.cost_discount_config = {"vertex_ai": 0.05} + monkeypatch.setattr(litellm, "cost_discount_config", {"vertex_ai": 0.05}) # Calculate cost with discount cost_with_discount = completion_cost( @@ -1855,8 +1829,6 @@ def test_cost_discount_vertex_ai(): custom_llm_provider="vertex_ai", ) - # Restore original config - litellm.cost_discount_config = original_discount_config # Verify discount is applied (5% off means 95% of original cost) expected_cost = cost_without_discount * 0.95 @@ -1868,15 +1840,13 @@ def test_cost_discount_vertex_ai(): print(f" - Savings: ${cost_without_discount - cost_with_discount:.6f}") -def test_cost_discount_not_applied_to_other_providers(): +def test_cost_discount_not_applied_to_other_providers(monkeypatch): """ Test that cost discount only applies to configured providers """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_discount_config = litellm.cost_discount_config.copy() # Create mock response for OpenAI response = ModelResponse( @@ -1889,7 +1859,7 @@ def test_cost_discount_not_applied_to_other_providers(): ) # Set discount only for vertex_ai (not openai) - litellm.cost_discount_config = {"vertex_ai": 0.05} + monkeypatch.setattr(litellm, "cost_discount_config", {"vertex_ai": 0.05}) # Calculate cost for OpenAI - should NOT have discount applied cost_with_selective_discount = completion_cost( @@ -1899,15 +1869,13 @@ def test_cost_discount_not_applied_to_other_providers(): ) # Clear discount config - litellm.cost_discount_config = {} + monkeypatch.setattr(litellm, "cost_discount_config", {}) cost_without_discount = completion_cost( completion_response=response, model="gpt-4", custom_llm_provider="openai", ) - # Restore original config - litellm.cost_discount_config = original_discount_config # Costs should be the same (no discount applied to OpenAI) assert cost_with_selective_discount == cost_without_discount @@ -1917,15 +1885,13 @@ def test_cost_discount_not_applied_to_other_providers(): print(f" - Cost remains unchanged: ${cost_with_selective_discount:.6f}") -def test_cost_margin_percentage(): +def test_cost_margin_percentage(monkeypatch): """ Test that percentage-based cost margin is applied correctly """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_margin_config = litellm.cost_margin_config.copy() # Create mock response response = ModelResponse( @@ -1938,7 +1904,7 @@ def test_cost_margin_percentage(): ) # Calculate cost without margin - litellm.cost_margin_config = {} + monkeypatch.setattr(litellm, "cost_margin_config", {}) cost_without_margin = completion_cost( completion_response=response, model="gpt-4", @@ -1946,7 +1912,7 @@ def test_cost_margin_percentage(): ) # Set 10% margin for openai - litellm.cost_margin_config = {"openai": 0.10} + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": 0.10}) # Calculate cost with margin cost_with_margin = completion_cost( @@ -1955,8 +1921,6 @@ def test_cost_margin_percentage(): custom_llm_provider="openai", ) - # Restore original config - litellm.cost_margin_config = original_margin_config # Verify margin is applied (10% margin means 110% of original cost) expected_cost = cost_without_margin * 1.10 @@ -1968,15 +1932,13 @@ def test_cost_margin_percentage(): print(f" - Margin added: ${cost_with_margin - cost_without_margin:.6f}") -def test_cost_margin_fixed_amount(): +def test_cost_margin_fixed_amount(monkeypatch): """ Test that fixed amount cost margin is applied correctly """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_margin_config = litellm.cost_margin_config.copy() # Create mock response response = ModelResponse( @@ -1989,7 +1951,7 @@ def test_cost_margin_fixed_amount(): ) # Calculate cost without margin - litellm.cost_margin_config = {} + monkeypatch.setattr(litellm, "cost_margin_config", {}) cost_without_margin = completion_cost( completion_response=response, model="gpt-4", @@ -1997,7 +1959,7 @@ def test_cost_margin_fixed_amount(): ) # Set $0.001 fixed margin for openai - litellm.cost_margin_config = {"openai": {"fixed_amount": 0.001}} + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": 0.001}}) # Calculate cost with margin cost_with_margin = completion_cost( @@ -2006,8 +1968,6 @@ def test_cost_margin_fixed_amount(): custom_llm_provider="openai", ) - # Restore original config - litellm.cost_margin_config = original_margin_config # Verify fixed margin is applied expected_cost = cost_without_margin + 0.001 @@ -2019,15 +1979,13 @@ def test_cost_margin_fixed_amount(): print(f" - Margin added: ${cost_with_margin - cost_without_margin:.6f}") -def test_cost_margin_combined(): +def test_cost_margin_combined(monkeypatch): """ Test that combined percentage and fixed amount margin is applied correctly """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_margin_config = litellm.cost_margin_config.copy() # Create mock response response = ModelResponse( @@ -2040,7 +1998,7 @@ def test_cost_margin_combined(): ) # Calculate cost without margin - litellm.cost_margin_config = {} + monkeypatch.setattr(litellm, "cost_margin_config", {}) cost_without_margin = completion_cost( completion_response=response, model="gpt-4", @@ -2048,9 +2006,9 @@ def test_cost_margin_combined(): ) # Set 8% margin + $0.0005 fixed for openai - litellm.cost_margin_config = { + monkeypatch.setattr(litellm, "cost_margin_config", { "openai": {"percentage": 0.08, "fixed_amount": 0.0005} - } + }) # Calculate cost with margin cost_with_margin = completion_cost( @@ -2059,8 +2017,6 @@ def test_cost_margin_combined(): custom_llm_provider="openai", ) - # Restore original config - litellm.cost_margin_config = original_margin_config # Verify combined margin is applied expected_cost = cost_without_margin * 1.08 + 0.0005 @@ -2072,15 +2028,13 @@ def test_cost_margin_combined(): print(f" - Margin added: ${cost_with_margin - cost_without_margin:.6f}") -def test_cost_margin_global(): +def test_cost_margin_global(monkeypatch): """ Test that global margin is applied when no provider-specific margin is configured """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_margin_config = litellm.cost_margin_config.copy() # Create mock response response = ModelResponse( @@ -2093,7 +2047,7 @@ def test_cost_margin_global(): ) # Calculate cost without margin - litellm.cost_margin_config = {} + monkeypatch.setattr(litellm, "cost_margin_config", {}) cost_without_margin = completion_cost( completion_response=response, model="gpt-4", @@ -2101,7 +2055,7 @@ def test_cost_margin_global(): ) # Set 5% global margin (no provider-specific margin) - litellm.cost_margin_config = {"global": 0.05} + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.05}) # Calculate cost with global margin cost_with_global_margin = completion_cost( @@ -2110,8 +2064,6 @@ def test_cost_margin_global(): custom_llm_provider="openai", ) - # Restore original config - litellm.cost_margin_config = original_margin_config # Verify global margin is applied expected_cost = cost_without_margin * 1.05 @@ -2123,15 +2075,13 @@ def test_cost_margin_global(): print(f" - Margin added: ${cost_with_global_margin - cost_without_margin:.6f}") -def test_cost_margin_provider_overrides_global(): +def test_cost_margin_provider_overrides_global(monkeypatch): """ Test that provider-specific margin overrides global margin """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original config - original_margin_config = litellm.cost_margin_config.copy() # Create mock response response = ModelResponse( @@ -2144,7 +2094,7 @@ def test_cost_margin_provider_overrides_global(): ) # Calculate cost without margin - litellm.cost_margin_config = {} + monkeypatch.setattr(litellm, "cost_margin_config", {}) cost_without_margin = completion_cost( completion_response=response, model="gpt-4", @@ -2152,7 +2102,7 @@ def test_cost_margin_provider_overrides_global(): ) # Set 5% global margin and 10% provider-specific margin - litellm.cost_margin_config = {"global": 0.05, "openai": 0.10} + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.05, "openai": 0.10}) # Calculate cost - should use provider-specific margin (10%), not global (5%) cost_with_provider_margin = completion_cost( @@ -2161,8 +2111,6 @@ def test_cost_margin_provider_overrides_global(): custom_llm_provider="openai", ) - # Restore original config - litellm.cost_margin_config = original_margin_config # Verify provider-specific margin is used (not global) expected_cost = cost_without_margin * 1.10 # 10% from provider, not 5% from global @@ -2176,16 +2124,13 @@ def test_cost_margin_provider_overrides_global(): print(f" - Margin added: ${cost_with_provider_margin - cost_without_margin:.6f}") -def test_cost_margin_with_discount(): +def test_cost_margin_with_discount(monkeypatch): """ Test that margin is applied after discount (independent calculation) """ from litellm import completion_cost from litellm.types.utils import Usage - # Save original configs - original_margin_config = litellm.cost_margin_config.copy() - original_discount_config = litellm.cost_discount_config.copy() # Create mock response response = ModelResponse( @@ -2198,8 +2143,8 @@ def test_cost_margin_with_discount(): ) # Calculate base cost - litellm.cost_margin_config = {} - litellm.cost_discount_config = {} + monkeypatch.setattr(litellm, "cost_margin_config", {}) + monkeypatch.setattr(litellm, "cost_discount_config", {}) base_cost = completion_cost( completion_response=response, model="gpt-4", @@ -2207,8 +2152,8 @@ def test_cost_margin_with_discount(): ) # Set 5% discount and 10% margin - litellm.cost_discount_config = {"openai": 0.05} - litellm.cost_margin_config = {"openai": 0.10} + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.05}) + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": 0.10}) # Calculate cost with both discount and margin cost_with_both = completion_cost( @@ -2217,9 +2162,6 @@ def test_cost_margin_with_discount(): custom_llm_provider="openai", ) - # Restore original configs - litellm.cost_margin_config = original_margin_config - litellm.cost_discount_config = original_discount_config # Verify: discount applied first, then margin # Base cost -> discount: base * 0.95 -> margin: (base * 0.95) * 1.10 @@ -2286,12 +2228,10 @@ def test_azure_image_generation_cost_calculator(): assert cost > 0.079 -def test_completion_cost_extracts_service_tier_from_response(): +def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_map): """Test that completion_cost extracts service_tier from completion_response object.""" from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2338,12 +2278,10 @@ def test_completion_cost_extracts_service_tier_from_response(): ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" -def test_completion_cost_extracts_service_tier_from_usage(): +def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map): """Test that completion_cost extracts service_tier from usage object.""" from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2397,12 +2335,10 @@ def test_completion_cost_extracts_service_tier_from_usage(): ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" -def test_completion_cost_service_tier_priority(): +def test_completion_cost_service_tier_priority(_local_model_cost_map): """Test that service_tier extraction follows priority: optional_params > completion_response > usage.""" from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2457,12 +2393,10 @@ def test_completion_cost_service_tier_priority(): ), "Costs from params and usage should be similar (both flex)" -def test_completion_cost_service_tier_for_bedrock(): +def test_completion_cost_service_tier_for_bedrock(_local_model_cost_map): """Test that Bedrock cost calculation applies service_tier-specific pricing.""" from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model" litellm.register_model( @@ -2507,7 +2441,7 @@ def test_completion_cost_service_tier_for_bedrock(): assert priority_cost > default_cost > flex_cost > 0 -def test_completion_cost_service_tier_for_anthropic(): +def test_completion_cost_service_tier_for_anthropic(_local_model_cost_map): """ Anthropic priority-tier requests must be priced at the priority rate. @@ -2519,8 +2453,6 @@ def test_completion_cost_service_tier_for_anthropic(): from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-service-tier-cost-model" litellm.register_model( @@ -2561,7 +2493,7 @@ def test_completion_cost_service_tier_for_anthropic(): assert priority_cost == pytest.approx(2 * standard_cost) -def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(): +def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(_local_model_cost_map): """ Proxy billing path regression for LIT-3771. @@ -2574,8 +2506,6 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(): from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-auto-tier-cost-model" litellm.register_model( @@ -2613,7 +2543,7 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(): assert cost == pytest.approx(expected_priority) -def test_completion_cost_non_string_service_tier_defers_to_served_tier(): +def test_completion_cost_non_string_service_tier_defers_to_served_tier(_local_model_cost_map): """ Regression: a non-string request-level ``service_tier`` (reachable via ``allowed_openai_params``/``drop_params``) must not crash cost tracking. @@ -2627,8 +2557,6 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier(): from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-non-string-tier-cost-model" litellm.register_model( @@ -2665,7 +2593,7 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier(): assert cost == pytest.approx(expected_priority) -def test_completion_cost_non_string_response_service_tier_defers_to_served_tier(): +def test_completion_cost_non_string_response_service_tier_defers_to_served_tier(_local_model_cost_map): """ Regression: a non-string ``service_tier`` on the response object must not crash cost tracking. @@ -2679,8 +2607,6 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-response-non-string-tier-cost-model" litellm.register_model( @@ -2718,7 +2644,7 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( assert cost == pytest.approx(expected_priority) -def test_completion_cost_non_string_usage_service_tier_prices_standard(): +def test_completion_cost_non_string_usage_service_tier_prices_standard(_local_model_cost_map): """ Regression: a non-string ``service_tier`` on the usage object must not crash cost tracking. @@ -2729,8 +2655,6 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard(): """ from litellm import completion_cost - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-usage-non-string-tier-cost-model" litellm.register_model( @@ -2764,7 +2688,7 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard(): assert cost == pytest.approx(expected_standard) -def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(): +def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_local_model_cost_map): """ Regression for the cache/tier interaction in the Anthropic geo/speed path. @@ -2780,8 +2704,6 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(): ) from litellm.types.utils import PromptTokensDetailsWrapper, Usage - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-priority-cache-fast-model" litellm.register_model( @@ -2837,7 +2759,7 @@ def _register_anthropic_geo_cache_model(model: str) -> None: ) -def test_anthropic_geo_multiplier_applies_to_cache_tokens(monkeypatch): +def test_anthropic_geo_multiplier_applies_to_cache_tokens(_local_model_cost_map, monkeypatch): """ Regression: the regional (geo) uplift must scale cache read and cache write cost too, not just non-cache input and output. @@ -2853,7 +2775,6 @@ def test_anthropic_geo_multiplier_applies_to_cache_tokens(monkeypatch): from litellm.types.utils import PromptTokensDetailsWrapper, Usage monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-geo-cache-model" _register_anthropic_geo_cache_model(model) @@ -2882,7 +2803,7 @@ def test_anthropic_geo_multiplier_applies_to_cache_tokens(monkeypatch): assert geo_completion_cost == pytest.approx(base_completion_cost * 1.1) -def test_anthropic_geo_and_fast_multipliers_compose(monkeypatch): +def test_anthropic_geo_and_fast_multipliers_compose(_local_model_cost_map, monkeypatch): """ The ``fast`` speed multiplier stays cache-exclusive (the old explicit ``fast/`` entries kept base cache rates) while the geo multiplier scales the @@ -2895,7 +2816,6 @@ def test_anthropic_geo_and_fast_multipliers_compose(monkeypatch): from litellm.types.utils import PromptTokensDetailsWrapper, Usage monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") model = "claude-test-geo-fast-cache-model" _register_anthropic_geo_cache_model(model) @@ -3100,7 +3020,7 @@ def test_gemini_implicit_caching_cost_calculation(): ) -def test_additional_costs_only_for_azure_ai(): +def test_additional_costs_only_for_azure_ai(_local_model_cost_map): """ Test that _get_additional_costs is only called for azure_ai provider. @@ -3111,8 +3031,6 @@ def test_additional_costs_only_for_azure_ai(): """ from litellm.cost_calculator import _get_additional_costs - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") # Non-azure_ai providers should return None result = _get_additional_costs( @@ -3140,7 +3058,7 @@ def test_additional_costs_only_for_azure_ai(): assert result is None, "Vertex AI should have no additional costs" -def test_openrouter_gemini_3_1_flash_lite_preview_pricing(): +def test_openrouter_gemini_3_1_flash_lite_preview_pricing(_local_model_cost_map): """ Test that openrouter/google/gemini-3.1-flash-lite-preview has a pricing entry. @@ -3150,8 +3068,6 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing(): model_prices_and_context_window.json when other Gemini 3.x variants were present. This caused ValueError: This model isn't mapped yet during router pre-call checks. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_name = "openrouter/google/gemini-3.1-flash-lite-preview" model_info = litellm.model_cost.get(model_name) @@ -3164,9 +3080,7 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing(): assert model_info["max_output_tokens"] == 65536 -def test_gemini_3_1_flash_lite_pricing(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") +def test_gemini_3_1_flash_lite_pricing(_local_model_cost_map): for model_name in ( "gemini-3.1-flash-lite", @@ -3489,7 +3403,7 @@ def test_custom_pricing_without_cache_keys_preserves_legacy_behavior(): assert cost == pytest.approx(expected) -def test_openrouter_gemini_3_1_flash_lite_stable_pricing(): +def test_openrouter_gemini_3_1_flash_lite_stable_pricing(_local_model_cost_map): """ Test that openrouter/google/gemini-3.1-flash-lite (stable, no -preview suffix) has a pricing entry. @@ -3505,8 +3419,6 @@ def test_openrouter_gemini_3_1_flash_lite_stable_pricing(): Pricing matches the existing -preview entry one-for-one (input $0.25/M, output $1.50/M, cache-read $0.025/M) — Google did not change costs at the GA cutover. """ - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") model_name = "openrouter/google/gemini-3.1-flash-lite" model_info = litellm.model_cost.get(model_name) @@ -3520,7 +3432,7 @@ def test_openrouter_gemini_3_1_flash_lite_stable_pricing(): assert model_info["max_output_tokens"] == 65536 -def test_completion_cost_logs_reasoning_and_cache_breakdown(): +def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_map): """ completion_cost must surface explicit reasoning and cache-read costs into the cost_breakdown stored on the logging object, so they end up in the spend logs @@ -3531,8 +3443,6 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(): from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") logging_obj = Logging( model="gemini-2.5-flash", @@ -3750,13 +3660,11 @@ def test_combine_usage_objects_sums_mirrored_cache_write_fields_once(): assert combined_pair.prompt_tokens_details.cache_creation_tokens == 100 -def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(): +def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_cost_map): """Regression: an Anthropic /v1/messages response reports cache reads as top-level cache_read_input_tokens with input_tokens excluding them. Reading that usage as Responses API usage dropped the cache tokens and billed the whole prompt at the uncached input rate, overstating spend on cache hits.""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") response = { "id": "msg_1", diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 4ab09d9d85b..28762e61861 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2789,7 +2789,10 @@ def _priced_at(prompt_tokens, completion_tokens): @pytest.fixture def local_cost_map(monkeypatch): + """The prices these tests assert are the checked-in ones. Setting the environment + variable alone does not reload the map, so pin the map itself.""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): From 73307070c2c67c92c7d59100a3fdaf155de66c0a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 21:09:27 -0700 Subject: [PATCH 32/43] test(key-management): unwind the global writes the key tests scaffold around (#37822) Seventeen tests in this file save a litellm module global, open a try, write it, and restore it in a finally. Four more sit behind autouse fixtures that reset the flag to a hard-coded False rather than to whatever it was. monkeypatch.setattr does all of that, so the capture, the try and the finally come out and the test body loses a level of indentation. The alias-format fixtures stop guessing the value they are restoring to. Also drops the sys.path.insert, whose argument resolves four levels above the repo, so it was never what made the imports work. TQ003 1077 -> 1076 and TQ005 2836 -> 2796, and the budget ceilings come down with them. 443 tests pass either way; the conftest snapshot was already catching these globals, so this is about not needing it. --- test-quality-budget.json | 4 +- .../test_key_management_endpoints.py | 828 +++++++++--------- 2 files changed, 405 insertions(+), 427 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 6428a55ba78..91e8f39d195 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -6,13 +6,13 @@ "limit": 742 }, "TQ003": { - "limit": 1075 + "limit": 1074 }, "TQ004": { "limit": 469 }, "TQ005": { - "limit": 2661 + "limit": 2621 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index fff6368cfc6..0c615cbaa32 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1,16 +1,10 @@ import json -import os -import sys import litellm import pytest import yaml from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../../..") -) # Adds the parent directory to the system path - from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException @@ -8125,26 +8119,22 @@ async def test_default_key_generate_params_duration(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Set default_key_generate_params with duration - original_value = litellm.default_key_generate_params - litellm.default_key_generate_params = {"duration": "180d"} + monkeypatch.setattr(litellm, "default_key_generate_params", {"duration": "180d"}) - try: - request = GenerateKeyRequest() # No duration specified - response = await _common_key_generation_helper( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-1234", - user_id="1234", - ), - litellm_changed_by=None, - team_table=None, - ) + request = GenerateKeyRequest() # No duration specified + response = await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) - # Verify duration was applied from defaults - assert request.duration == "180d" - finally: - litellm.default_key_generate_params = original_value + # Verify duration was applied from defaults + assert request.duration == "180d" async def test_default_key_generate_params_object_permission_applied_when_absent( @@ -8184,28 +8174,28 @@ async def test_default_key_generate_params_object_permission_applied_when_absent monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - original_value = litellm.default_key_generate_params - litellm.default_key_generate_params = { - "object_permission": {"vector_stores": ["default-vs"]} - } + monkeypatch.setattr( + litellm, + "default_key_generate_params", + { + "object_permission": {"vector_stores": ["default-vs"]} + }, + ) - try: - request = GenerateKeyRequest() # No object_permission specified - await _common_key_generation_helper( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-1234", - user_id="1234", - ), - litellm_changed_by=None, - team_table=None, - ) + request = GenerateKeyRequest() # No object_permission specified + await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) - created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] - assert created_data["vector_stores"] == ["default-vs"] - finally: - litellm.default_key_generate_params = original_value + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["vector_stores"] == ["default-vs"] async def test_default_key_generate_params_object_permission_merges_partial( @@ -8247,31 +8237,31 @@ async def test_default_key_generate_params_object_permission_merges_partial( monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - original_value = litellm.default_key_generate_params - litellm.default_key_generate_params = { - "object_permission": {"vector_stores": ["default-vs"]} - } + monkeypatch.setattr( + litellm, + "default_key_generate_params", + { + "object_permission": {"vector_stores": ["default-vs"]} + }, + ) - try: - request = GenerateKeyRequest( - object_permission=LiteLLM_ObjectPermissionBase(agents=["agent-1"]) - ) - await _common_key_generation_helper( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-1234", - user_id="1234", - ), - litellm_changed_by=None, - team_table=None, - ) + request = GenerateKeyRequest( + object_permission=LiteLLM_ObjectPermissionBase(agents=["agent-1"]) + ) + await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) - created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] - assert created_data["agents"] == ["agent-1"] - assert created_data["vector_stores"] == ["default-vs"] - finally: - litellm.default_key_generate_params = original_value + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["agents"] == ["agent-1"] + assert created_data["vector_stores"] == ["default-vs"] async def test_default_key_generate_params_object_permission_does_not_override_explicit( @@ -8312,32 +8302,32 @@ async def test_default_key_generate_params_object_permission_does_not_override_e monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - original_value = litellm.default_key_generate_params - litellm.default_key_generate_params = { - "object_permission": {"vector_stores": ["default-vs"]} - } + monkeypatch.setattr( + litellm, + "default_key_generate_params", + { + "object_permission": {"vector_stores": ["default-vs"]} + }, + ) - try: - request = GenerateKeyRequest( - object_permission=LiteLLM_ObjectPermissionBase( - vector_stores=["explicit-vs"] - ) - ) - await _common_key_generation_helper( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-1234", - user_id="1234", - ), - litellm_changed_by=None, - team_table=None, + request = GenerateKeyRequest( + object_permission=LiteLLM_ObjectPermissionBase( + vector_stores=["explicit-vs"] ) + ) + await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) - created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] - assert created_data["vector_stores"] == ["explicit-vs"] - finally: - litellm.default_key_generate_params = original_value + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["vector_stores"] == ["explicit-vs"] async def test_default_key_generate_params_object_permission_not_rejected_for_non_admin_personal_key( @@ -8380,29 +8370,29 @@ async def test_default_key_generate_params_object_permission_not_rejected_for_no monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - original_value = litellm.default_key_generate_params - litellm.default_key_generate_params = { - "object_permission": {"vector_stores": ["default-vs"]} - } + monkeypatch.setattr( + litellm, + "default_key_generate_params", + { + "object_permission": {"vector_stores": ["default-vs"]} + }, + ) - try: - request = GenerateKeyRequest(user_id="alice") # No object_permission specified - response = await _common_key_generation_helper( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - api_key="sk-alice", - user_id="alice", - ), - litellm_changed_by=None, - team_table=None, - ) + request = GenerateKeyRequest(user_id="alice") # No object_permission specified + response = await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) - assert response is not None - created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] - assert created_data["vector_stores"] == ["default-vs"] - finally: - litellm.default_key_generate_params = original_value + assert response is not None + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["vector_stores"] == ["default-vs"] @pytest.mark.asyncio @@ -9261,10 +9251,8 @@ async def test_key_aliases_admin_sees_all(): class TestValidateKeyAliasFormat: @pytest.fixture(autouse=True) - def reset_key_alias_flag(self): - litellm.enable_key_alias_format_validation = False - yield - litellm.enable_key_alias_format_validation = False + def reset_key_alias_flag(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "enable_key_alias_format_validation", False) def test_validation_skipped_when_flag_disabled(self): """When enable_key_alias_format_validation is False (default), no charset/length validation occurs.""" @@ -9305,12 +9293,12 @@ class TestValidateKeyAliasFormat: assert str(exc.value.code) == "400" assert "Invalid key_alias" in str(exc.value.message) - def test_validate_key_alias_format_valid(self): + def test_validate_key_alias_format_valid(self, monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import ( _validate_key_alias_format, ) - litellm.enable_key_alias_format_validation = True + monkeypatch.setattr(litellm, "enable_key_alias_format_validation", True) # Valid cases _validate_key_alias_format(None) # OK _validate_key_alias_format("valid-alias") @@ -9322,13 +9310,13 @@ class TestValidateKeyAliasFormat: _validate_key_alias_format("user/user@example.com") _validate_key_alias_format("team/user@example.com") - def test_validate_key_alias_format_invalid(self): + def test_validate_key_alias_format_invalid(self, monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import ( _validate_key_alias_format, ) from litellm.proxy._types import ProxyException - litellm.enable_key_alias_format_validation = True + monkeypatch.setattr(litellm, "enable_key_alias_format_validation", True) invalid_aliases = [ "", # empty " ", # whitespace @@ -10956,10 +10944,8 @@ class TestKeyAliasSkipValidationOnUnchanged: """ @pytest.fixture(autouse=True) - def enable_validation(self): - litellm.enable_key_alias_format_validation = True - yield - litellm.enable_key_alias_format_validation = False + def enable_validation(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "enable_key_alias_format_validation", True) @pytest.fixture def mock_prisma(self): @@ -11075,146 +11061,142 @@ class TestKeyAliasSkipValidationOnUnchanged: # --- Tests: _enforce_upperbound_key_params --- -def test_enforce_upperbound_rejects_over_limit_on_generate(): +def test_enforce_upperbound_rejects_over_limit_on_generate(monkeypatch): """Test that key generation is rejected when values exceed upperbound.""" import litellm from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - tpm_limit=1000, rpm_limit=100, max_budget=10.0 - ) - data = GenerateKeyRequest(tpm_limit=5000) - with pytest.raises(HTTPException) as exc_info: - _enforce_upperbound_key_params(data, fill_defaults=True) - assert exc_info.value.status_code == 400 - assert "tpm_limit" in str(exc_info.value.detail) - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + tpm_limit=1000, rpm_limit=100, max_budget=10.0 + ), + ) + data = GenerateKeyRequest(tpm_limit=5000) + with pytest.raises(HTTPException) as exc_info: + _enforce_upperbound_key_params(data, fill_defaults=True) + assert exc_info.value.status_code == 400 + assert "tpm_limit" in str(exc_info.value.detail) -def test_enforce_upperbound_fills_defaults_on_generate(): +def test_enforce_upperbound_fills_defaults_on_generate(monkeypatch): """Test that None values are filled with upperbound defaults during generation.""" import litellm from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - tpm_limit=1000, rpm_limit=100 - ) - data = GenerateKeyRequest() # tpm_limit=None, rpm_limit=None - _enforce_upperbound_key_params(data, fill_defaults=True) - assert data.tpm_limit == 1000 - assert data.rpm_limit == 100 - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + tpm_limit=1000, rpm_limit=100 + ), + ) + data = GenerateKeyRequest() # tpm_limit=None, rpm_limit=None + _enforce_upperbound_key_params(data, fill_defaults=True) + assert data.tpm_limit == 1000 + assert data.rpm_limit == 100 -def test_enforce_upperbound_skips_none_on_update(): +def test_enforce_upperbound_skips_none_on_update(monkeypatch): """Test that None values are NOT filled during update (fill_defaults=False).""" import litellm from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - tpm_limit=1000, rpm_limit=100 - ) - data = UpdateKeyRequest(key="sk-test") # tpm_limit=None, rpm_limit=None - _enforce_upperbound_key_params(data, fill_defaults=False) - assert data.tpm_limit is None # should NOT be filled - assert data.rpm_limit is None # should NOT be filled - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + tpm_limit=1000, rpm_limit=100 + ), + ) + data = UpdateKeyRequest(key="sk-test") # tpm_limit=None, rpm_limit=None + _enforce_upperbound_key_params(data, fill_defaults=False) + assert data.tpm_limit is None # should NOT be filled + assert data.rpm_limit is None # should NOT be filled -def test_enforce_upperbound_rejects_over_limit_on_update(): +def test_enforce_upperbound_rejects_over_limit_on_update(monkeypatch): """Test that key update is rejected when values exceed upperbound.""" import litellm from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - tpm_limit=1000, rpm_limit=100, max_budget=10.0 - ) - data = UpdateKeyRequest(key="sk-test", tpm_limit=5000) - with pytest.raises(HTTPException) as exc_info: - _enforce_upperbound_key_params(data, fill_defaults=False) - assert exc_info.value.status_code == 400 - assert "tpm_limit" in str(exc_info.value.detail) - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + tpm_limit=1000, rpm_limit=100, max_budget=10.0 + ), + ) + data = UpdateKeyRequest(key="sk-test", tpm_limit=5000) + with pytest.raises(HTTPException) as exc_info: + _enforce_upperbound_key_params(data, fill_defaults=False) + assert exc_info.value.status_code == 400 + assert "tpm_limit" in str(exc_info.value.detail) -def test_enforce_upperbound_allows_within_limit_on_update(): +def test_enforce_upperbound_allows_within_limit_on_update(monkeypatch): """Test that key update passes when values are within upperbound.""" import litellm from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - tpm_limit=1000, rpm_limit=100, max_budget=10.0 - ) - data = UpdateKeyRequest( - key="sk-test", tpm_limit=500, rpm_limit=50, max_budget=5.0 - ) - _enforce_upperbound_key_params(data, fill_defaults=False) - # Should not raise - assert data.tpm_limit == 500 - assert data.rpm_limit == 50 - assert data.max_budget == 5.0 - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + tpm_limit=1000, rpm_limit=100, max_budget=10.0 + ), + ) + data = UpdateKeyRequest( + key="sk-test", tpm_limit=500, rpm_limit=50, max_budget=5.0 + ) + _enforce_upperbound_key_params(data, fill_defaults=False) + # Should not raise + assert data.tpm_limit == 500 + assert data.rpm_limit == 50 + assert data.max_budget == 5.0 -def test_enforce_upperbound_duration_over_limit(): +def test_enforce_upperbound_duration_over_limit(monkeypatch): """Test that duration exceeding upperbound is rejected.""" import litellm from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - duration="7d" - ) - data = UpdateKeyRequest(key="sk-test", duration="30d") - with pytest.raises(HTTPException) as exc_info: - _enforce_upperbound_key_params(data, fill_defaults=False) - assert exc_info.value.status_code == 400 - assert "duration" in str(exc_info.value.detail) - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + duration="7d" + ), + ) + data = UpdateKeyRequest(key="sk-test", duration="30d") + with pytest.raises(HTTPException) as exc_info: + _enforce_upperbound_key_params(data, fill_defaults=False) + assert exc_info.value.status_code == 400 + assert "duration" in str(exc_info.value.detail) -def test_enforce_upperbound_no_config_is_noop(): +def test_enforce_upperbound_no_config_is_noop(monkeypatch): """Test that no enforcement happens when upperbound params are not configured.""" import litellm - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = None - data = UpdateKeyRequest(key="sk-test", tpm_limit=999999) - _enforce_upperbound_key_params(data, fill_defaults=False) - # Should not raise — no enforcement configured - assert data.tpm_limit == 999999 - finally: - litellm.upperbound_key_generate_params = original + monkeypatch.setattr(litellm, "upperbound_key_generate_params", None) + data = UpdateKeyRequest(key="sk-test", tpm_limit=999999) + _enforce_upperbound_key_params(data, fill_defaults=False) + # Should not raise — no enforcement configured + assert data.tpm_limit == 999999 # --- Tests: _execute_virtual_key_regeneration enforces upperbound --- @@ -11267,7 +11249,7 @@ def _make_regenerate_existing_key(): @pytest.mark.asyncio -async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(): +async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monkeypatch): """Regenerate must reject durations exceeding upperbound_key_generate_params.duration.""" from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -11277,91 +11259,34 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(): LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - duration="1h" - ) - existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest(duration="2h") - user_api_key_dict = _make_regenerate_user_api_key_dict() - mock_prisma_client = _make_regenerate_mock_prisma() - - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - ): - with pytest.raises(HTTPException) as exc_info: - await _execute_virtual_key_regeneration( - prisma_client=mock_prisma_client, - key_in_db=existing_key, - hashed_api_key="abc123", - key="abc123", - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - user_api_key_cache=MagicMock(), - proxy_logging_obj=MagicMock(), - ) - assert exc_info.value.status_code == 400 - assert "duration" in str(exc_info.value.detail) - # Rejected regenerate must not reach the DB update. - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 - finally: - litellm.upperbound_key_generate_params = original - - -@pytest.mark.asyncio -async def test_execute_virtual_key_regeneration_allows_within_limit_duration(): - """Regenerate must accept durations within upperbound_key_generate_params.duration.""" - from litellm.proxy._types import RegenerateKeyRequest - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _execute_virtual_key_regeneration, - ) - from litellm.types.proxy.management_endpoints.ui_sso import ( - LiteLLM_UpperboundKeyGenerateParams, + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + duration="1h" + ), ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="2h") + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - duration="1h" - ) - existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest(duration="30m") - user_api_key_dict = _make_regenerate_user_api_key_dict() - mock_prisma_client = _make_regenerate_mock_prisma() - - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", - new_callable=AsyncMock, - ), - ): + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + with pytest.raises(HTTPException) as exc_info: await _execute_virtual_key_regeneration( prisma_client=mock_prisma_client, key_in_db=existing_key, @@ -11373,13 +11298,70 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(): user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), ) - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 - finally: - litellm.upperbound_key_generate_params = original + assert exc_info.value.status_code == 400 + assert "duration" in str(exc_info.value.detail) + # Rejected regenerate must not reach the DB update. + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 @pytest.mark.asyncio -async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(): +async def test_execute_virtual_key_regeneration_allows_within_limit_duration(monkeypatch): + """Regenerate must accept durations within upperbound_key_generate_params.duration.""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + duration="1h" + ), + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="30m") + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(monkeypatch): """Regenerate must reject max_budget exceeding upperbound — proves the fix covers non-duration fields.""" from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -11389,52 +11371,52 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget(): LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - max_budget=10.0 - ) - existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest(max_budget=500.0) - user_api_key_dict = _make_regenerate_user_api_key_dict() - mock_prisma_client = _make_regenerate_mock_prisma() + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + max_budget=10.0 + ), + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(max_budget=500.0) + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - ): - with pytest.raises(HTTPException) as exc_info: - await _execute_virtual_key_regeneration( - prisma_client=mock_prisma_client, - key_in_db=existing_key, - hashed_api_key="abc123", - key="abc123", - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - user_api_key_cache=MagicMock(), - proxy_logging_obj=MagicMock(), - ) - assert exc_info.value.status_code == 400 - assert "max_budget" in str(exc_info.value.detail) - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 - finally: - litellm.upperbound_key_generate_params = original + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert exc_info.value.status_code == 400 + assert "max_budget" in str(exc_info.value.detail) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 @pytest.mark.asyncio -async def test_execute_virtual_key_regeneration_skips_none_values(): +async def test_execute_virtual_key_regeneration_skips_none_values(monkeypatch): """Regenerate with data.duration=None must not raise, even when upperbound is set (fill_defaults=False semantic — None means 'inherit from existing key').""" from litellm.proxy._types import RegenerateKeyRequest @@ -11445,100 +11427,96 @@ async def test_execute_virtual_key_regeneration_skips_none_values(): LiteLLM_UpperboundKeyGenerateParams, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams( - duration="1h" - ) - existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest() # all fields None - user_api_key_dict = _make_regenerate_user_api_key_dict() - mock_prisma_client = _make_regenerate_mock_prisma() + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams( + duration="1h" + ), + ) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest() # all fields None + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", - new_callable=AsyncMock, - ), - ): - await _execute_virtual_key_regeneration( - prisma_client=mock_prisma_client, - key_in_db=existing_key, - hashed_api_key="abc123", - key="abc123", - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - user_api_key_cache=MagicMock(), - proxy_logging_obj=MagicMock(), - ) - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 - finally: - litellm.upperbound_key_generate_params = original + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 @pytest.mark.asyncio -async def test_execute_virtual_key_regeneration_no_upperbound_config_is_noop(): +async def test_execute_virtual_key_regeneration_no_upperbound_config_is_noop(monkeypatch): """Regenerate with no upperbound config set must accept any duration.""" from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( _execute_virtual_key_regeneration, ) - original = litellm.upperbound_key_generate_params - try: - litellm.upperbound_key_generate_params = None - existing_key = _make_regenerate_existing_key() - data = RegenerateKeyRequest(duration="30d") - user_api_key_dict = _make_regenerate_user_api_key_dict() - mock_prisma_client = _make_regenerate_mock_prisma() + monkeypatch.setattr(litellm, "upperbound_key_generate_params", None) + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="30d") + user_api_key_dict = _make_regenerate_user_api_key_dict() + mock_prisma_client = _make_regenerate_mock_prisma() - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", - new_callable=AsyncMock, - ), - ): - await _execute_virtual_key_regeneration( - prisma_client=mock_prisma_client, - key_in_db=existing_key, - hashed_api_key="abc123", - key="abc123", - data=data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - user_api_key_cache=MagicMock(), - proxy_logging_obj=MagicMock(), - ) - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 - finally: - litellm.upperbound_key_generate_params = original + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 class TestAllowedRoutesCallerPermission: From 6bce3dce0dd1ebcf1e0ea9d5a7bb206568d9e3d9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 21:19:35 -0700 Subject: [PATCH 33/43] test(callbacks): unwind the callbacks global the policy engine and realtime tests scaffold around (#37826) * test(policy-engine): unwind the callback global the pipeline tests scaffold around Every one of the 16 tests in this file set litellm.callbacks by hand, each wrapping its body in a try/finally to put the old value back, and each capturing that old value with a .copy() first. That is 32 TQ005 violations and about 70 lines of scaffolding to say what monkeypatch.setattr says in one. The write also sat outside the try, so the block that restores it did not cover the statement that changed it. 16 tests pass either way, and litellm.callbacks reads restored on both sides, because the conftest snapshot already lists it. The point is that these tests stop depending on that snapshot to clean up after them. * test(realtime): unwind the same callback global in the realtime streaming tests Same global, same shape as the previous commit. 25 writes to litellm.callbacks, 2 of them wrapped in a try/finally that resets to [] rather than to the old value, and 12 tests that write it with no protection at all. monkeypatch.setattr replaces all of them, and the sys.path.insert with its now-unused os and sys imports goes too. Both sides read restored here as well, for the same reason as the previous commit: litellm.callbacks is in the conftest snapshot. What changes is that these tests no longer lean on it. 101 tests pass in this file, 16 in the policy engine one. * style(realtime): wrap the one signature the monkeypatch param pushed past 120 --- test-quality-budget.json | 4 +- .../test_realtime_streaming.py | 134 ++-- .../policy_engine/test_pipeline_executor.py | 570 ++++++++---------- 3 files changed, 310 insertions(+), 398 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 91e8f39d195..46e368a495b 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -6,13 +6,13 @@ "limit": 742 }, "TQ003": { - "limit": 1074 + "limit": 1073 }, "TQ004": { "limit": 469 }, "TQ005": { - "limit": 2621 + "limit": 2562 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index ccf353b1b6c..61b63e2b917 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -1,6 +1,4 @@ import json -import os -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -8,7 +6,6 @@ from websockets.exceptions import ConnectionClosed import litellm -sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( @@ -1326,7 +1323,7 @@ async def test_log_messages_includes_tools_in_model_call_details(): @pytest.mark.asyncio -async def test_realtime_guardrail_blocks_prompt_injection(): +async def test_realtime_guardrail_blocks_prompt_injection(monkeypatch: pytest.MonkeyPatch): """ Test that when a transcription event containing prompt injection arrives from the backend, a registered guardrail blocks it — sending a warning to the client @@ -1350,7 +1347,7 @@ async def test_realtime_guardrail_blocks_prompt_injection(): event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) # --- client websocket mock --- client_ws = MagicMock() @@ -1405,11 +1402,10 @@ async def test_realtime_guardrail_blocks_prompt_injection(): f"Expected guardrail_violation error type, got: {error_events[0]}" ) - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_guardrail_allows_clean_transcript(): +async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.MonkeyPatch): """ Test that a clean transcript passes through the guardrail and triggers response.create to the backend. @@ -1430,7 +1426,7 @@ async def test_realtime_guardrail_allows_clean_transcript(): event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1463,11 +1459,10 @@ async def test_realtime_guardrail_allows_clean_transcript(): response_creates = [e for e in sent_to_backend if e.get("type") == "response.create"] assert len(response_creates) == 1, f"Clean transcript should trigger response.create, got: {sent_to_backend}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_text_input_guardrail_blocks_and_returns_error(): +async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ Test that when conversation.item.create arrives with text that triggers a guardrail, the proxy blocks it (doesn't forward to backend) and returns an error event directly @@ -1495,7 +1490,7 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(): event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1558,11 +1553,10 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(): ] assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(): +async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ Test that a client-supplied function_call_output whose content triggers a guardrail is blocked: it is not forwarded to the backend, and an error @@ -1590,7 +1584,7 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error( event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1648,11 +1642,10 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error( assert sanitized_item["call_id"] == "call_123" assert "test@example.com" not in sanitized_item["output"] - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_function_call_output_guardrail_allows_clean_output(): +async def test_realtime_function_call_output_guardrail_allows_clean_output(monkeypatch: pytest.MonkeyPatch): """ Test that a clean function_call_output passes through and reaches the backend when guardrails are configured. @@ -1670,7 +1663,7 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output(): event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1714,11 +1707,10 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output(): ] assert len(forwarded) == 1, f"Clean function_call_output should be forwarded, got: {forwarded}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_text_input_guardrail_uses_pre_call_mode(): +async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch): """ Test that _has_realtime_guardrails returns True for a guardrail configured with pre_call mode (not just realtime_input_transcription). @@ -1736,7 +1728,7 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(): event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() backend_ws = MagicMock() @@ -1751,11 +1743,10 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(): "pre_call-only guardrail must not disable server_vad auto-response" ) - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_session_created_injects_session_update_for_audio_guardrail(): +async def test_realtime_session_created_injects_session_update_for_audio_guardrail(monkeypatch: pytest.MonkeyPatch): """ Test that when an audio transcription guardrail is configured, a session.created event from the backend triggers a session.update injection (create_response: false) @@ -1775,7 +1766,7 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1809,11 +1800,12 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra "GA session.update must nest turn_detection under audio.input" ) - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only(): +async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only( + monkeypatch: pytest.MonkeyPatch, +): """ pre_call-only guardrails must not inject create_response:false on realtime sessions — that breaks server_vad for audio-only voice agents (e.g. Model Armor). @@ -1831,7 +1823,7 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c event_hook=GuardrailEventHooks.pre_call, default_on=True, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1853,11 +1845,10 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c session_updates = [e for e in sent_to_backend if e.get("type") == "session.update"] assert len(session_updates) == 0, f"pre_call-only guardrail must not inject session.update, got: {sent_to_backend}" - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(): +async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monkeypatch: pytest.MonkeyPatch): """Model Armor-style pre_call + post_call must not gate audio VAD.""" import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -1867,18 +1858,22 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(): async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): return inputs - litellm.callbacks = [ - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_pre_call", - event_hook=GuardrailEventHooks.pre_call, - default_on=False, - ), - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_post_call", - event_hook=GuardrailEventHooks.post_call, - default_on=False, - ), - ] + monkeypatch.setattr( + litellm, + "callbacks", + [ + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_pre_call", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ), + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_post_call", + event_hook=GuardrailEventHooks.post_call, + default_on=False, + ), + ], + ) client_ws = MagicMock() backend_ws = MagicMock() @@ -1900,11 +1895,10 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(): assert streaming._has_realtime_guardrails() is True assert streaming._has_audio_transcription_guardrails() is False - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_end_session_after_n_fails_closes_connection(): +async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.MonkeyPatch): """ Test that end_session_after_n_fails=2 closes the backend websocket after the second guardrail violation in a session. @@ -1923,7 +1917,7 @@ async def test_end_session_after_n_fails_closes_connection(): default_on=True, end_session_after_n_fails=2, ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1948,11 +1942,10 @@ async def test_end_session_after_n_fails_closes_connection(): assert backend_ws.close.called, "Expected backend_ws.close() to be called after 2 violations" assert streaming._violation_count == 2 - litellm.callbacks = [] # cleanup @pytest.mark.asyncio -async def test_on_violation_end_session_closes_on_first_fail(): +async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest.MonkeyPatch): """ Test that on_violation='end_session' closes the session immediately on the first violation, regardless of end_session_after_n_fails. @@ -1971,7 +1964,7 @@ async def test_on_violation_end_session_closes_on_first_fail(): default_on=True, on_violation="end_session", ) - litellm.callbacks = [guardrail] + monkeypatch.setattr(litellm, "callbacks", [guardrail]) client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -1995,7 +1988,6 @@ async def test_on_violation_end_session_closes_on_first_fail(): assert backend_ws.close.called, "Expected session to close immediately with on_violation=end_session" assert streaming._violation_count == 1 - litellm.callbacks = [] # cleanup @pytest.mark.asyncio @@ -2898,53 +2890,47 @@ def _transcription_guardrail(): ) -def test_setup_folds_in_auto_response_disable_when_transcription_guardrail_active(): +def test_setup_folds_in_auto_response_disable_when_transcription_guardrail_active(monkeypatch: pytest.MonkeyPatch): """Gemini rejects a second setup, so a transcription guardrail's auto-response disable must be folded into the one-and-only setup; otherwise the model auto-responds and the guardrail is bypassed.""" import litellm - litellm.callbacks = [_transcription_guardrail()] - try: - streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) - setup = json.dumps( - { - "setup": { - "model": "models/gemini-3.1-flash-live-preview", - "generationConfig": {"responseModalities": ["AUDIO"]}, - "inputAudioTranscription": {}, - } + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) + setup = json.dumps( + { + "setup": { + "model": "models/gemini-3.1-flash-live-preview", + "generationConfig": {"responseModalities": ["AUDIO"]}, + "inputAudioTranscription": {}, } - ) - out = json.loads(streaming._maybe_inject_guardrail_auto_response_disable(setup)) - aad = out["setup"]["realtimeInputConfig"]["automaticActivityDetection"] - assert aad["disabled"] is True - finally: - litellm.callbacks = [] + } + ) + out = json.loads(streaming._maybe_inject_guardrail_auto_response_disable(setup)) + aad = out["setup"]["realtimeInputConfig"]["automaticActivityDetection"] + assert aad["disabled"] is True -def test_setup_unchanged_without_transcription_guardrail(): +def test_setup_unchanged_without_transcription_guardrail(monkeypatch: pytest.MonkeyPatch): import litellm - litellm.callbacks = [] + monkeypatch.setattr(litellm, "callbacks", []) streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) setup = json.dumps({"setup": {"model": "x", "generationConfig": {"responseModalities": ["AUDIO"]}}}) out = streaming._maybe_inject_guardrail_auto_response_disable(setup) assert json.loads(out) == json.loads(setup) -def test_non_bidi_setup_left_untouched_for_followup_capable_providers(): +def test_non_bidi_setup_left_untouched_for_followup_capable_providers(monkeypatch: pytest.MonkeyPatch): """OpenAI realtime accepts a follow-up session.update, so a non-bidi message (no top-level 'setup' key) must be left untouched even with a guardrail on.""" import litellm - litellm.callbacks = [_transcription_guardrail()] - try: - streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) - msg = json.dumps({"type": "session.update", "session": {"instructions": "hi"}}) - assert streaming._maybe_inject_guardrail_auto_response_disable(msg) == msg - finally: - litellm.callbacks = [] + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) + msg = json.dumps({"type": "session.update", "session": {"instructions": "hi"}}) + assert streaming._maybe_inject_guardrail_auto_response_disable(msg) == msg @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 840d93eb12c..22d212dd8ae 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -165,7 +165,7 @@ class ContentCheckGuardrail(CustomGuardrail): @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_escalation_step1_fails_step2_blocks(): +async def test_escalation_step1_fails_step2_blocks(monkeypatch): """ Pipeline: simple-filter (on_fail: next) -> advanced-filter (on_fail: block) Input: request that fails simple-filter @@ -182,36 +182,32 @@ async def test_escalation_step1_fails_step2_blocks(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [simple_guard, advanced_guard] + monkeypatch.setattr(litellm, "callbacks", [simple_guard, advanced_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "bad content"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "bad content"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert simple_guard.calls == 1 - assert advanced_guard.calls == 1 - assert result.terminal_action == "block" - assert len(result.step_results) == 2 - assert result.step_results[0].guardrail_name == "simple-filter" - assert result.step_results[0].outcome == "fail" - assert result.step_results[0].action_taken == "next" - assert result.step_results[1].guardrail_name == "advanced-filter" - assert result.step_results[1].outcome == "fail" - assert result.step_results[1].action_taken == "block" - finally: - litellm.callbacks = original_callbacks + assert simple_guard.calls == 1 + assert advanced_guard.calls == 1 + assert result.terminal_action == "block" + assert len(result.step_results) == 2 + assert result.step_results[0].guardrail_name == "simple-filter" + assert result.step_results[0].outcome == "fail" + assert result.step_results[0].action_taken == "next" + assert result.step_results[1].guardrail_name == "advanced-filter" + assert result.step_results[1].outcome == "fail" + assert result.step_results[1].action_taken == "block" @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_block_carries_original_guardrail_exception(): +async def test_block_carries_original_guardrail_exception(monkeypatch): """A blocking step must expose the guardrail's own raised exception on the result so the caller can re-raise it verbatim, giving the policy path the same response/trace as a direct guardrail attachment.""" @@ -219,67 +215,52 @@ async def test_block_carries_original_guardrail_exception(): pipeline = GuardrailPipeline( mode="pre_call", - steps=[ - PipelineStep( - guardrail="moderation-filter", on_fail="block", on_pass="allow" - ) - ], + steps=[PipelineStep(guardrail="moderation-filter", on_fail="block", on_pass="allow")], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [guard] + monkeypatch.setattr(litellm, "callbacks", [guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "bad content"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "bad content"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert result.terminal_action == "block" - assert isinstance(result.original_exception, HTTPException) - assert result.original_exception.status_code == 400 - assert result.original_exception.detail == "Content policy violation" - finally: - litellm.callbacks = original_callbacks + assert result.terminal_action == "block" + assert isinstance(result.original_exception, HTTPException) + assert result.original_exception.status_code == 400 + assert result.original_exception.detail == "Content policy violation" @pytest.mark.asyncio -async def test_unsupported_mode_yields_error_outcome_without_exception(): +async def test_unsupported_mode_yields_error_outcome_without_exception(monkeypatch): """An unexpected hook mode must surface as an error outcome (carrying no original exception), not crash or run the guardrail.""" guard = AlwaysPassGuardrail(guardrail_name="filter") - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [guard] + monkeypatch.setattr(litellm, "callbacks", [guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=[PipelineStep(guardrail="filter", on_error="block", on_fail="block")], - mode="during_call", - data={"messages": [{"role": "user", "content": "hi"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="filter", on_error="block", on_fail="block")], + mode="during_call", + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert guard.calls == 0 - assert result.terminal_action == "block" - assert result.step_results[0].outcome == "error" - assert ( - "Unsupported pipeline mode: during_call" - in result.step_results[0].error_detail - ) - assert result.original_exception is None - finally: - litellm.callbacks = original_callbacks + assert guard.calls == 0 + assert result.terminal_action == "block" + assert result.step_results[0].outcome == "error" + assert "Unsupported pipeline mode: during_call" in result.step_results[0].error_detail + assert result.original_exception is None @pytest.mark.asyncio -async def test_passthrough_guardrail_failure_can_pipeline_block(): +async def test_passthrough_guardrail_failure_can_pipeline_block(monkeypatch): """ Pipeline: passthrough guardrail (on_fail: block) Expected: passthrough ModifyResponseException is treated as policy fail, @@ -298,35 +279,31 @@ async def test_passthrough_guardrail_failure_can_pipeline_block(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [passthrough_guard] + monkeypatch.setattr(litellm, "callbacks", [passthrough_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={ - "model": "fake-model", - "messages": [{"role": "user", "content": "bad content"}], - }, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={ + "model": "fake-model", + "messages": [{"role": "user", "content": "bad content"}], + }, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert passthrough_guard.calls == 1 - assert result.terminal_action == "block" - assert len(result.step_results) == 1 - assert result.step_results[0].guardrail_name == "passthrough-filter" - assert result.step_results[0].outcome == "fail" - assert result.step_results[0].action_taken == "block" - assert result.error_message == "Content policy violation" - finally: - litellm.callbacks = original_callbacks + assert passthrough_guard.calls == 1 + assert result.terminal_action == "block" + assert len(result.step_results) == 1 + assert result.step_results[0].guardrail_name == "passthrough-filter" + assert result.step_results[0].outcome == "fail" + assert result.step_results[0].action_taken == "block" + assert result.error_message == "Content policy violation" @pytest.mark.asyncio -async def test_custom_code_guardrail_failure_can_pipeline_block(): +async def test_custom_code_guardrail_failure_can_pipeline_block(monkeypatch): """ Pipeline: custom code guardrail (on_fail: block) Expected: custom code keeps its standalone passthrough block behavior, and @@ -334,10 +311,7 @@ async def test_custom_code_guardrail_failure_can_pipeline_block(): """ custom_guard = CustomCodeGuardrail( guardrail_name="custom-code-filter", - custom_code=( - "def apply_guardrail(inputs, request_data, input_type):\n" - ' return block("SSN detected")\n' - ), + custom_code=('def apply_guardrail(inputs, request_data, input_type):\n return block("SSN detected")\n'), ) pipeline = GuardrailPipeline( @@ -351,35 +325,31 @@ async def test_custom_code_guardrail_failure_can_pipeline_block(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [custom_guard] + monkeypatch.setattr(litellm, "callbacks", [custom_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={ - "model": "fake-model", - "messages": [{"role": "user", "content": "123-45-6789"}], - }, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={ + "model": "fake-model", + "messages": [{"role": "user", "content": "123-45-6789"}], + }, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert result.terminal_action == "block" - assert len(result.step_results) == 1 - assert result.step_results[0].guardrail_name == "custom-code-filter" - assert result.step_results[0].outcome == "fail" - assert result.step_results[0].action_taken == "block" - assert result.error_message == "SSN detected" - finally: - litellm.callbacks = original_callbacks + assert result.terminal_action == "block" + assert len(result.step_results) == 1 + assert result.step_results[0].guardrail_name == "custom-code-filter" + assert result.step_results[0].outcome == "fail" + assert result.step_results[0].action_taken == "block" + assert result.error_message == "SSN detected" @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_early_allow_step1_passes_step2_skipped(): +async def test_early_allow_step1_passes_step2_skipped(monkeypatch): """ Pipeline: simple-filter (on_pass: allow) -> advanced-filter Input: clean request that passes simple-filter @@ -396,32 +366,28 @@ async def test_early_allow_step1_passes_step2_skipped(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [simple_guard, advanced_guard] + monkeypatch.setattr(litellm, "callbacks", [simple_guard, advanced_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "clean content"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "clean content"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert simple_guard.calls == 1 - assert advanced_guard.calls == 0 - assert result.terminal_action == "allow" - assert len(result.step_results) == 1 - assert result.step_results[0].outcome == "pass" - assert result.step_results[0].action_taken == "allow" - finally: - litellm.callbacks = original_callbacks + assert simple_guard.calls == 1 + assert advanced_guard.calls == 0 + assert result.terminal_action == "allow" + assert len(result.step_results) == 1 + assert result.step_results[0].outcome == "pass" + assert result.step_results[0].action_taken == "allow" @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_escalation_step1_fails_step2_passes(): +async def test_escalation_step1_fails_step2_passes(monkeypatch): """ Pipeline: simple-filter (on_fail: next) -> advanced-filter (on_pass: allow) Input: request that fails simple but passes advanced @@ -438,34 +404,30 @@ async def test_escalation_step1_fails_step2_passes(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [simple_guard, advanced_guard] + monkeypatch.setattr(litellm, "callbacks", [simple_guard, advanced_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "borderline content"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="content-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "borderline content"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="content-safety", + ) - assert simple_guard.calls == 1 - assert advanced_guard.calls == 1 - assert result.terminal_action == "allow" - assert len(result.step_results) == 2 - assert result.step_results[0].outcome == "fail" - assert result.step_results[0].action_taken == "next" - assert result.step_results[1].outcome == "pass" - assert result.step_results[1].action_taken == "allow" - finally: - litellm.callbacks = original_callbacks + assert simple_guard.calls == 1 + assert advanced_guard.calls == 1 + assert result.terminal_action == "allow" + assert len(result.step_results) == 2 + assert result.step_results[0].outcome == "fail" + assert result.step_results[0].action_taken == "next" + assert result.step_results[1].outcome == "pass" + assert result.step_results[1].action_taken == "allow" @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_data_forwarding_pii_masking(): +async def test_data_forwarding_pii_masking(monkeypatch): """ Pipeline: pii-masker (pass_data: true, on_pass: next) -> content-check (on_pass: allow) Input: "Hello John Smith" @@ -487,31 +449,27 @@ async def test_data_forwarding_pii_masking(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [pii_guard, content_guard] + monkeypatch.setattr(litellm, "callbacks", [pii_guard, content_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "Hello John Smith"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="pii-then-safety", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "Hello John Smith"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="pii-then-safety", + ) - assert pii_guard.calls == 1 - assert content_guard.calls == 1 - assert content_guard.received_messages[0]["content"] == "Hello [REDACTED]" - assert result.terminal_action == "allow" - assert result.modified_data is not None - assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]" - finally: - litellm.callbacks = original_callbacks + assert pii_guard.calls == 1 + assert content_guard.calls == 1 + assert content_guard.received_messages[0]["content"] == "Hello [REDACTED]" + assert result.terminal_action == "allow" + assert result.modified_data is not None + assert result.modified_data["messages"][0]["content"] == "Hello [REDACTED]" @pytest.mark.asyncio -async def test_guardrail_not_found_uses_on_fail(): +async def test_guardrail_not_found_uses_on_fail(monkeypatch): """ If a guardrail is not found, treat as error and use on_fail action. """ @@ -526,29 +484,25 @@ async def test_guardrail_not_found_uses_on_fail(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [] + monkeypatch.setattr(litellm, "callbacks", []) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "test"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test-policy", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test-policy", + ) - assert result.terminal_action == "block" - assert result.step_results[0].outcome == "error" - assert "not found" in result.step_results[0].error_detail - finally: - litellm.callbacks = original_callbacks + assert result.terminal_action == "block" + assert result.step_results[0].outcome == "error" + assert "not found" in result.step_results[0].error_detail @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(): +async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(monkeypatch): """ Policy intervention (400) uses on_fail; technical error (503) uses on_error. @@ -574,32 +528,28 @@ async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [primary, fallback] + monkeypatch.setattr(litellm, "callbacks", [primary, fallback]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "any"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="mod-fallback", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "any"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="mod-fallback", + ) - assert primary.calls == 1 - assert fallback.calls == 1 - assert result.terminal_action == "allow" - assert result.step_results[0].outcome == "error" - assert result.step_results[0].action_taken == "next" - assert result.step_results[1].outcome == "pass" - finally: - litellm.callbacks = original_callbacks + assert primary.calls == 1 + assert fallback.calls == 1 + assert result.terminal_action == "allow" + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "next" + assert result.step_results[1].outcome == "pass" @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_on_fail_next_on_content_on_error_block_stops_api_fallback(): +async def test_on_fail_next_on_content_on_error_block_stops_api_fallback(monkeypatch): """ Content policy fail (400) uses on_fail: next; API error uses on_error: block (no second step). """ @@ -625,48 +575,40 @@ async def test_on_fail_next_on_content_on_error_block_stops_api_fallback(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [primary_content, fallback] + monkeypatch.setattr(litellm, "callbacks", [primary_content, fallback]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline_content.steps, - mode=pipeline_content.mode, - data={"messages": [{"role": "user", "content": "bad"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test", - ) - assert result.terminal_action == "allow" - assert primary_content.calls == 1 - assert fallback.calls == 1 - finally: - litellm.callbacks = original_callbacks + result = await PipelineExecutor.execute_steps( + steps=pipeline_content.steps, + mode=pipeline_content.mode, + data={"messages": [{"role": "user", "content": "bad"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) + assert result.terminal_action == "allow" + assert primary_content.calls == 1 + assert fallback.calls == 1 # API outage: on_error block -> do not run fallback fallback.calls = 0 - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [primary_api, fallback] - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline_content.steps, - mode=pipeline_content.mode, - data={"messages": [{"role": "user", "content": "ok"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test", - ) - assert result.terminal_action == "block" - assert primary_api.calls == 1 - assert fallback.calls == 0 - assert result.step_results[0].outcome == "error" - assert result.step_results[0].action_taken == "block" - finally: - litellm.callbacks = original_callbacks + monkeypatch.setattr(litellm, "callbacks", [primary_api, fallback]) + result = await PipelineExecutor.execute_steps( + steps=pipeline_content.steps, + mode=pipeline_content.mode, + data={"messages": [{"role": "user", "content": "ok"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) + assert result.terminal_action == "block" + assert primary_api.calls == 1 + assert fallback.calls == 0 + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "block" @pytest.mark.asyncio -async def test_guardrail_not_found_with_next_continues(): +async def test_guardrail_not_found_with_next_continues(monkeypatch): """ If a guardrail is not found and on_fail is 'next', continue to next step. """ @@ -688,32 +630,28 @@ async def test_guardrail_not_found_with_next_continues(): ], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [pass_guard] + monkeypatch.setattr(litellm, "callbacks", [pass_guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "test"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test-policy", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test-policy", + ) - assert result.terminal_action == "allow" - assert len(result.step_results) == 2 - assert result.step_results[0].outcome == "error" - assert result.step_results[0].action_taken == "next" - assert result.step_results[1].outcome == "pass" - assert pass_guard.calls == 1 - finally: - litellm.callbacks = original_callbacks + assert result.terminal_action == "allow" + assert len(result.step_results) == 2 + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "next" + assert result.step_results[1].outcome == "pass" + assert pass_guard.calls == 1 @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio -async def test_single_step_pipeline_block(): +async def test_single_step_pipeline_block(monkeypatch): """Single step pipeline that blocks.""" guard = AlwaysFailGuardrail(guardrail_name="blocker") @@ -722,27 +660,23 @@ async def test_single_step_pipeline_block(): steps=[PipelineStep(guardrail="blocker", on_fail="block")], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [guard] + monkeypatch.setattr(litellm, "callbacks", [guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "test"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) - assert result.terminal_action == "block" - assert guard.calls == 1 - finally: - litellm.callbacks = original_callbacks + assert result.terminal_action == "block" + assert guard.calls == 1 @pytest.mark.asyncio -async def test_single_step_pipeline_allow(): +async def test_single_step_pipeline_allow(monkeypatch): """Single step pipeline that allows.""" guard = AlwaysPassGuardrail(guardrail_name="passer") @@ -751,27 +685,23 @@ async def test_single_step_pipeline_allow(): steps=[PipelineStep(guardrail="passer", on_pass="allow")], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [guard] + monkeypatch.setattr(litellm, "callbacks", [guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "test"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) - assert result.terminal_action == "allow" - assert guard.calls == 1 - finally: - litellm.callbacks = original_callbacks + assert result.terminal_action == "allow" + assert guard.calls == 1 @pytest.mark.asyncio -async def test_step_results_include_duration(): +async def test_step_results_include_duration(monkeypatch): """Step results should include timing information.""" guard = AlwaysPassGuardrail(guardrail_name="timed") @@ -780,23 +710,19 @@ async def test_step_results_include_duration(): steps=[PipelineStep(guardrail="timed")], ) - original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [guard] + monkeypatch.setattr(litellm, "callbacks", [guard]) - try: - result = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode=pipeline.mode, - data={"messages": [{"role": "user", "content": "test"}]}, - user_api_key_dict=MagicMock(), - call_type="completion", - policy_name="test", - ) + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) - assert result.step_results[0].duration_seconds is not None - assert result.step_results[0].duration_seconds >= 0 - finally: - litellm.callbacks = original_callbacks + assert result.step_results[0].duration_seconds is not None + assert result.step_results[0].duration_seconds >= 0 class _PolicyOptOutGuardrail(CustomGuardrail): From 39a580aa91e5cc4d2677100d44546d2a6309568c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 21:29:31 -0700 Subject: [PATCH 34/43] test(guardrails): stop five guardrail test files leaking env vars on failure (#37828) Onyx, prompt security, hiddenlayer, repelloai and deepkeep all write straight to os.environ and unset again at the bottom of each test. None of the five has a try/finally, so the moment a test fails it returns to the runner with the keys still set and whatever runs next in that worker inherits them. Raising inside test_onyx_guard_with_custom_timeout_from_kwargs on the current files leaves ONYX_API_BASE and ONYX_API_KEY behind; doing the same in test_hiddenlayer_config_saas leaves HIDDENLAYER_API_BASE. Both come back clean after this. 89 raw writes and the hand-rolled deletes become monkeypatch calls. The class-level setup_method and teardown_method pair in the onyx file, sweeping the same three keys twice, becomes one autouse fixture. The sys.path.insert lines and their now-unused imports go too, and litellm.set_verbose = True, which only turned global debug logging on for whatever ran next, is dropped rather than restored. test_onyx_guard_config and test_prompt_security_guard_config asserted nothing at all, so they could only fail by raising. Each now pins what init_guardrails_v2 produces: exactly one guardrail of the right class on litellm.callbacks, carrying the configured name, default_on and hook. The zero-assert tests in the other three are left alone; those are a judgement about each guardrail rather than a mechanical sweep. tests/test_litellm/proxy/guardrails passes at 2873. --- test-quality-budget.json | 6 +- .../guardrail_hooks/test_deepkeep.py | 12 +- .../guardrail_hooks/test_hiddenlayer.py | 62 ++++---- .../guardrails/guardrail_hooks/test_onyx.py | 137 +++++------------- .../guardrail_hooks/test_repelloai.py | 14 +- .../test_prompt_security_guardrails.py | 91 +++--------- 6 files changed, 99 insertions(+), 223 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 46e368a495b..6a62d783044 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -1,18 +1,18 @@ { "TQ001": { - "limit": 746 + "limit": 744 }, "TQ002": { "limit": 742 }, "TQ003": { - "limit": 1073 + "limit": 1068 }, "TQ004": { "limit": 469 }, "TQ005": { - "limit": 2562 + "limit": 2549 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py index af0686fcc59..03f418e6d7a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_deepkeep.py @@ -1,10 +1,8 @@ import os -import sys import pytest from unittest.mock import patch, MagicMock, AsyncMock from httpx import Response, Request -sys.path.insert(0, os.path.abspath("../..")) import litellm from litellm.proxy.guardrails.guardrail_hooks.deepkeep.deepkeep import ( @@ -17,10 +15,9 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.exceptions import GuardrailRaisedException -def test_deepkeep_guard_config(monkeypatch): +def test_deepkeep_guard_config(monkeypatch: pytest.MonkeyPatch): """Test DeepKeep guard configuration with init_guardrails_v2.""" - litellm.set_verbose = True - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) monkeypatch.setenv("DEEPKEEP_API_KEY", "test-key") monkeypatch.setenv("DEEPKEEP_API_BASE", "https://test.deepkeep.ai") @@ -42,9 +39,6 @@ def test_deepkeep_guard_config(monkeypatch): ) # Clean up - del os.environ["DEEPKEEP_API_KEY"] - del os.environ["DEEPKEEP_API_BASE"] - del os.environ["DEEPKEEP_FIREWALL_ID"] class TestDeepKeepGuardrail: @@ -108,7 +102,7 @@ class TestDeepKeepGuardrail: == "https://test.deepkeep.ai/v3/openai/beta/litellm_basic_guardrail_api" ) - def test_initialization_with_env_vars(self, monkeypatch): + def test_initialization_with_env_vars(self, monkeypatch: pytest.MonkeyPatch): """should initialize successfully using environment variables.""" monkeypatch.setenv("DEEPKEEP_API_KEY", "env-key") monkeypatch.setenv("DEEPKEEP_API_BASE", "https://env.deepkeep.ai") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index 57adf85b3d9..1b2108c837d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -1,5 +1,4 @@ import os -import sys import uuid from typing import List, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -8,7 +7,6 @@ import pytest from fastapi import HTTPException from httpx import Request, Response -sys.path.insert(0, os.path.abspath("../..")) import litellm from litellm import ModelResponse @@ -26,10 +24,9 @@ from litellm.types.utils import ( ) -def test_hiddenlayer_config_saas(monkeypatch): +def test_hiddenlayer_config_saas(monkeypatch: pytest.MonkeyPatch): """Test Hiddenlayer SaaS configuration with init_guardrails_v2.""" - litellm.set_verbose = True - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) # Set environment variables for testing monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -50,8 +47,6 @@ def test_hiddenlayer_config_saas(monkeypatch): ) # Clean up - if "HIDDENLAYER_API_BASE" in os.environ: - del os.environ["HIDDENLAYER_API_BASE"] class TestHiddenlayerGuardrail: @@ -71,7 +66,7 @@ class TestHiddenlayerGuardrail: if key in os.environ: del os.environ[key] - def test_initialization(self, monkeypatch): + def test_initialization(self, monkeypatch: pytest.MonkeyPatch): """Test successful initialization with default values.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -84,17 +79,16 @@ class TestHiddenlayerGuardrail: assert guardrail.guardrail_name == "hiddenlayer" assert guardrail.event_hook == "pre_call" - def test_initialization_fails_when_api_key_missing(self): + def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch): """Test that initialization fails when API key is not set.""" # Ensure API key is not set - if "HIDDENLAYER_CLIENT_SECRET" in os.environ: - del os.environ["HIDDENLAYER_CLIENT_SECRET"] + monkeypatch.delenv("HIDDENLAYER_CLIENT_SECRET", raising=False) with pytest.raises(RuntimeError): HiddenlayerGuardrail(guardrail_name="hiddenlayer", event_hook="pre_call") @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self, monkeypatch): + async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for request with no violations detected.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -151,7 +145,7 @@ class TestHiddenlayerGuardrail: assert call_args.args[0] == f"{guardrail.api_base}/detection/v1/interactions" @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self, monkeypatch): + async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for request with violations detected.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -209,7 +203,7 @@ class TestHiddenlayerGuardrail: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self, monkeypatch): + async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response with no violations detected.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -279,7 +273,7 @@ class TestHiddenlayerGuardrail: mock_post.assert_called_once() @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self, monkeypatch): + async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response with violations detected.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -348,7 +342,7 @@ class TestHiddenlayerGuardrail: assert exc_info.value.status_code == 400 @pytest.mark.asyncio - async def test_apply_guardrail_api_error_handling(self, monkeypatch): + async def test_apply_guardrail_api_error_handling(self, monkeypatch: pytest.MonkeyPatch): """Test handling of API errors in apply_guardrail.""" # Set required API key monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -391,7 +385,7 @@ class TestHiddenlayerGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_validate_with_call_hiddenlayer_method(self, monkeypatch): + async def test_validate_with_call_hiddenlayer_method(self, monkeypatch: pytest.MonkeyPatch): """Test the _validate_with_guard_server internal method.""" # Set required API key monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -433,7 +427,7 @@ class TestHiddenlayerGuardrail: ) @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image(self, monkeypatch): + async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail sends multimodal content (image) to HiddenLayer v1.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -498,7 +492,7 @@ class TestHiddenlayerGuardrail: assert result is not None @pytest.mark.asyncio - async def test_apply_guardrail_redact_with_image_content(self, monkeypatch): + async def test_apply_guardrail_redact_with_image_content(self, monkeypatch: pytest.MonkeyPatch): """Test that REDACT action with multimodal content extracts text properly into inputs['texts'].""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -570,10 +564,9 @@ class TestHiddenlayerGuardrail: assert config_model.__name__ == "HiddenlayerGuardrailConfigModel" -def test_hiddenlayer_config_v2(monkeypatch): +def test_hiddenlayer_config_v2(monkeypatch: pytest.MonkeyPatch): """Test HiddenLayer V2 configuration with init_guardrails_v2.""" - litellm.set_verbose = True - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -593,8 +586,6 @@ def test_hiddenlayer_config_v2(monkeypatch): config_file_path="", ) - if "HIDDENLAYER_API_BASE" in os.environ: - del os.environ["HIDDENLAYER_API_BASE"] class TestHiddenlayerGuardrailV2: @@ -612,7 +603,7 @@ class TestHiddenlayerGuardrailV2: if key in os.environ: del os.environ[key] - def test_initialization(self, monkeypatch): + def test_initialization(self, monkeypatch: pytest.MonkeyPatch): """Test successful initialization with default values.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -624,16 +615,15 @@ class TestHiddenlayerGuardrailV2: assert guardrail.guardrail_name == "hiddenlayer" assert guardrail.event_hook == "pre_call" - def test_initialization_fails_when_api_key_missing(self): + def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch): """Test that initialization fails when API key is not set for SaaS.""" - if "HIDDENLAYER_CLIENT_SECRET" in os.environ: - del os.environ["HIDDENLAYER_CLIENT_SECRET"] + monkeypatch.delenv("HIDDENLAYER_CLIENT_SECRET", raising=False) with pytest.raises(RuntimeError): HiddenlayerGuardrailV2(guardrail_name="hiddenlayer", event_hook="pre_call") @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self, monkeypatch): + async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for request with no violations detected.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -691,7 +681,7 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/request-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self, monkeypatch): + async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for request with violations detected (block via header).""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -751,7 +741,7 @@ class TestHiddenlayerGuardrailV2: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self, monkeypatch): + async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response with no violations detected.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -816,7 +806,7 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self, monkeypatch): + async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response with violations detected (block via header).""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -863,7 +853,7 @@ class TestHiddenlayerGuardrailV2: assert "Blocked by Hiddenlayer" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch): + async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response containing tool calls.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -924,7 +914,7 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in call_args.args[0] @pytest.mark.asyncio - async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch): + async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch: pytest.MonkeyPatch): """Test that _call_hiddenlayer uses the v2 request/response evaluation endpoints.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -959,7 +949,7 @@ class TestHiddenlayerGuardrailV2: assert "detection/v2/response-evaluations" in mock_post.call_args.args[0] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image(self, monkeypatch): + async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail sends multimodal content (image) to HiddenLayer v2.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") @@ -1030,7 +1020,7 @@ class TestHiddenlayerGuardrailV2: assert texts == ["how much is on this receipt?"] @pytest.mark.asyncio - async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch): + async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch: pytest.MonkeyPatch): """Test that new_texts extraction handles multimodal content (list) returned by HiddenLayer v2.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py index fa4624eac99..9208e0b3075 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_onyx.py @@ -1,5 +1,3 @@ -import os -import sys import uuid from unittest.mock import AsyncMock, MagicMock, patch @@ -8,8 +6,6 @@ import pytest from fastapi import HTTPException from httpx import Request, Response -sys.path.insert(0, os.path.abspath("../..")) - import litellm from litellm import ModelResponse from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -18,12 +14,11 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message -def test_onyx_guard_config(monkeypatch): +def test_onyx_guard_config(monkeypatch: pytest.MonkeyPatch): """Test Onyx guard configuration with init_guardrails_v2.""" - litellm.set_verbose = True - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setattr(litellm, "callbacks", []) - # Set environment variables for testing monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") monkeypatch.setenv("ONYX_API_KEY", "test-api-key") @@ -41,16 +36,15 @@ def test_onyx_guard_config(monkeypatch): config_file_path="", ) - # Clean up - if "ONYX_API_BASE" in os.environ: - del os.environ["ONYX_API_BASE"] - if "ONYX_API_KEY" in os.environ: - del os.environ["ONYX_API_KEY"] + registered = [c for c in litellm.callbacks if isinstance(c, OnyxGuardrail)] + assert len(registered) == 1 + assert registered[0].guardrail_name == "onyx-guard" + assert registered[0].default_on is True + assert registered[0].event_hook == "pre_call" -def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch): +def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch: pytest.MonkeyPatch): """Test Onyx guard instantiation with custom timeout passed via kwargs.""" - # Set environment variables for testing monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") monkeypatch.setenv("ONYX_API_KEY", "test-api-key") @@ -74,20 +68,13 @@ def test_onyx_guard_with_custom_timeout_from_kwargs(monkeypatch): assert timeout_param.read == 45.0 assert timeout_param.connect == 5.0 - # Clean up - if "ONYX_API_BASE" in os.environ: - del os.environ["ONYX_API_BASE"] - if "ONYX_API_KEY" in os.environ: - del os.environ["ONYX_API_KEY"] - -def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch): +def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch: pytest.MonkeyPatch): """Test Onyx guard with timeout=None uses ONYX_TIMEOUT env var. When timeout=None is passed (as it would be from config model with default None), the ONYX_TIMEOUT environment variable should be used. """ - # Set environment variables for testing monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") monkeypatch.setenv("ONYX_API_KEY", "test-api-key") monkeypatch.setenv("ONYX_TIMEOUT", "60") @@ -112,23 +99,13 @@ def test_onyx_guard_with_timeout_none_uses_env_var(monkeypatch): assert timeout_param.read == 60.0 assert timeout_param.connect == 5.0 - # Clean up - if "ONYX_API_BASE" in os.environ: - del os.environ["ONYX_API_BASE"] - if "ONYX_API_KEY" in os.environ: - del os.environ["ONYX_API_KEY"] - if "ONYX_TIMEOUT" in os.environ: - del os.environ["ONYX_TIMEOUT"] - -def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch): +def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch: pytest.MonkeyPatch): """Test Onyx guard with timeout=None and no env var defaults to 10 seconds.""" - # Set environment variables for testing monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Ensure ONYX_TIMEOUT is not set - if "ONYX_TIMEOUT" in os.environ: - del os.environ["ONYX_TIMEOUT"] + monkeypatch.delenv("ONYX_TIMEOUT", raising=False) with patch( "litellm.proxy.guardrails.guardrail_hooks.onyx.onyx.get_async_httpx_client" @@ -150,33 +127,17 @@ def test_onyx_guard_with_timeout_none_defaults_to_10(monkeypatch): assert timeout_param.read == 10.0 assert timeout_param.connect == 5.0 - # Clean up - if "ONYX_API_BASE" in os.environ: - del os.environ["ONYX_API_BASE"] - if "ONYX_API_KEY" in os.environ: - del os.environ["ONYX_API_KEY"] - class TestOnyxGuardrail: """Test suite for Onyx Security Guardrail integration.""" - def setup_method(self): - """Setup test environment.""" - # Clean up any existing environment variables - for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]: - if key in os.environ: - del os.environ[key] + @pytest.fixture(autouse=True) + def clear_onyx_env(self, monkeypatch: pytest.MonkeyPatch) -> None: + for key in ("ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"): + monkeypatch.delenv(key, raising=False) - def teardown_method(self): - """Clean up test environment.""" - # Clean up any environment variables set during tests - for key in ["ONYX_API_BASE", "ONYX_API_KEY", "ONYX_TIMEOUT"]: - if key in os.environ: - del os.environ[key] - - def test_initialization_with_defaults(self, monkeypatch): + def test_initialization_with_defaults(self, monkeypatch: pytest.MonkeyPatch): """Test successful initialization with default values.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -189,7 +150,7 @@ class TestOnyxGuardrail: assert guardrail.guardrail_name == "test-guard" assert guardrail.event_hook == "pre_call" - def test_initialization_with_env_vars(self, monkeypatch): + def test_initialization_with_env_vars(self, monkeypatch: pytest.MonkeyPatch): """Test initialization with environment variables.""" monkeypatch.setenv("ONYX_API_BASE", "https://custom.onyx.security") monkeypatch.setenv("ONYX_API_KEY", "custom-api-key") @@ -202,18 +163,17 @@ class TestOnyxGuardrail: assert guardrail.api_key == "custom-api-key" assert guardrail.event_hook == "post_call" - def test_initialization_fails_when_api_key_missing(self): + def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch): """Test that initialization fails when API key is not set.""" # Ensure API key is not set - if "ONYX_API_KEY" in os.environ: - del os.environ["ONYX_API_KEY"] + monkeypatch.delenv("ONYX_API_KEY", raising=False) with pytest.raises( ValueError, match="ONYX_API_KEY environment variable is not set" ): OnyxGuardrail(guardrail_name="test-guard", event_hook="pre_call") - def test_initialization_with_default_timeout(self, monkeypatch): + def test_initialization_with_default_timeout(self, monkeypatch: pytest.MonkeyPatch): """Test that default timeout is 10.0 seconds.""" monkeypatch.setenv("ONYX_API_KEY", "test-api-key") @@ -232,7 +192,7 @@ class TestOnyxGuardrail: assert timeout_param.read == 10.0 assert timeout_param.connect == 5.0 - def test_initialization_with_custom_timeout_parameter(self, monkeypatch): + def test_initialization_with_custom_timeout_parameter(self, monkeypatch: pytest.MonkeyPatch): """Test initialization with custom timeout parameter.""" monkeypatch.setenv("ONYX_API_KEY", "test-api-key") @@ -254,7 +214,7 @@ class TestOnyxGuardrail: assert timeout_param.read == 30.0 assert timeout_param.connect == 5.0 - def test_initialization_with_timeout_from_env_var(self, monkeypatch): + def test_initialization_with_timeout_from_env_var(self, monkeypatch: pytest.MonkeyPatch): """Test initialization with timeout from ONYX_TIMEOUT environment variable. Note: The env var is only used when timeout=None is explicitly passed, @@ -282,7 +242,7 @@ class TestOnyxGuardrail: assert timeout_param.read == 25.0 assert timeout_param.connect == 5.0 - def test_initialization_timeout_parameter_overrides_env_var(self, monkeypatch): + def test_initialization_timeout_parameter_overrides_env_var(self, monkeypatch: pytest.MonkeyPatch): """Test that timeout parameter overrides ONYX_TIMEOUT environment variable.""" monkeypatch.setenv("ONYX_API_KEY", "test-api-key") monkeypatch.setenv("ONYX_TIMEOUT", "25") @@ -306,9 +266,8 @@ class TestOnyxGuardrail: assert timeout_param.connect == 5.0 @pytest.mark.asyncio - async def test_apply_guardrail_request_no_violations(self, monkeypatch): + async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for request with no violations detected.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail @@ -372,9 +331,8 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-call-id" @pytest.mark.asyncio - async def test_apply_guardrail_request_with_violations(self, monkeypatch): + async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for request with violations detected.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail @@ -423,9 +381,8 @@ class TestOnyxGuardrail: assert "prompt_injection" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_response_no_violations(self, monkeypatch): + async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response with no violations detected.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail @@ -497,9 +454,8 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-call-id-2" @pytest.mark.asyncio - async def test_apply_guardrail_response_with_violations(self, monkeypatch): + async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail for response with violations detected.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") # Setup guardrail @@ -558,9 +514,8 @@ class TestOnyxGuardrail: assert "illegal_instructions" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_apply_guardrail_api_error_handling(self, monkeypatch): + async def test_apply_guardrail_api_error_handling(self, monkeypatch: pytest.MonkeyPatch): """Test handling of API errors in apply_guardrail.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -591,9 +546,8 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_timeout_error_handling(self, monkeypatch): + async def test_apply_guardrail_timeout_error_handling(self, monkeypatch: pytest.MonkeyPatch): """Test handling of timeout errors in apply_guardrail (graceful degradation).""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -629,9 +583,8 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_read_timeout_error_handling(self, monkeypatch): + async def test_apply_guardrail_read_timeout_error_handling(self, monkeypatch: pytest.MonkeyPatch): """Test handling of read timeout errors in apply_guardrail.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -667,9 +620,8 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_connect_timeout_error_handling(self, monkeypatch): + async def test_apply_guardrail_connect_timeout_error_handling(self, monkeypatch: pytest.MonkeyPatch): """Test handling of connect timeout errors in apply_guardrail.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -705,9 +657,8 @@ class TestOnyxGuardrail: assert result == inputs @pytest.mark.asyncio - async def test_apply_guardrail_no_logging_obj(self, monkeypatch): + async def test_apply_guardrail_no_logging_obj(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail without logging object (uses UUID).""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -747,9 +698,8 @@ class TestOnyxGuardrail: assert call_args.kwargs["json"]["conversation_id"] == "test-uuid" @pytest.mark.asyncio - async def test_validate_with_guard_server_method(self, monkeypatch): + async def test_validate_with_guard_server_method(self, monkeypatch: pytest.MonkeyPatch): """Test the _validate_with_guard_server internal method.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -788,9 +738,8 @@ class TestOnyxGuardrail: ) @pytest.mark.asyncio - async def test_validate_with_guard_server_blocked(self, monkeypatch): + async def test_validate_with_guard_server_blocked(self, monkeypatch: pytest.MonkeyPatch): """Test _validate_with_guard_server when request is blocked.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -825,9 +774,8 @@ class TestOnyxGuardrail: assert config_model.__name__ == "OnyxGuardrailConfigModel" @pytest.mark.asyncio - async def test_apply_guardrail_with_modelresponse(self, monkeypatch): + async def test_apply_guardrail_with_modelresponse(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail with ModelResponse object for response type.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -880,9 +828,8 @@ class TestOnyxGuardrail: assert "payload" in call_args.kwargs["json"] @pytest.mark.asyncio - async def test_apply_guardrail_response_error_handling(self, monkeypatch): + async def test_apply_guardrail_response_error_handling(self, monkeypatch: pytest.MonkeyPatch): """Test error handling when processing response data.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( @@ -925,7 +872,7 @@ class TestOnyxIntegration: """Test integration scenarios.""" @pytest.mark.asyncio - async def test_full_guardrail_flow(self, monkeypatch): + async def test_full_guardrail_flow(self, monkeypatch: pytest.MonkeyPatch): """Test full guardrail flow with multiple hooks.""" # Set environment variables monkeypatch.setenv("ONYX_API_BASE", "https://test.onyx.security") @@ -966,16 +913,10 @@ class TestOnyxIntegration: ) assert len(custom_loggers) >= 3 - # Clean up - if "ONYX_API_BASE" in os.environ: - del os.environ["ONYX_API_BASE"] - if "ONYX_API_KEY" in os.environ: - del os.environ["ONYX_API_KEY"] @pytest.mark.asyncio - async def test_apply_guardrail_empty_request_data(self, monkeypatch): + async def test_apply_guardrail_empty_request_data(self, monkeypatch: pytest.MonkeyPatch): """Test apply_guardrail with empty request data.""" - # Set required API key monkeypatch.setenv("ONYX_API_KEY", "test-api-key") guardrail = OnyxGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py index 1322d93ce70..1ef25b6e7ab 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_repelloai.py @@ -1,11 +1,9 @@ import os -import sys import pytest from fastapi import HTTPException from httpx import ConnectError, Request, Response -sys.path.insert(0, os.path.abspath("../..")) import litellm from litellm import DualCache @@ -93,23 +91,23 @@ class TestRepelloAIInitialization: with pytest.raises(ValueError, match="asset_id"): RepelloAIGuardrail(api_key="test-api-key", guardrail_name="t") - def test_api_key_from_env(self, monkeypatch): + def test_api_key_from_env(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("REPELLOAI_API_KEY", "env-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "env-key" - def test_api_key_from_argus_env(self, monkeypatch): + def test_api_key_from_argus_env(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("ARGUS_API_KEY", "argus-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "argus-key" - def test_argus_env_preferred_over_legacy(self, monkeypatch): + def test_argus_env_preferred_over_legacy(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("ARGUS_API_KEY", "argus-key") monkeypatch.setenv("REPELLOAI_API_KEY", "legacy-key") guardrail = RepelloAIGuardrail(asset_id="asset-123", guardrail_name="t") assert guardrail.repelloai_api_key == "argus-key" - def test_explicit_api_key_preferred_over_env(self, monkeypatch): + def test_explicit_api_key_preferred_over_env(self, monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("ARGUS_API_KEY", "argus-key") guardrail = RepelloAIGuardrail( api_key="explicit-key", asset_id="asset-123", guardrail_name="t" @@ -145,9 +143,9 @@ class TestRepelloAIInitialization: assert guardrail.api_base == DEFAULT_REPELLOAI_API_BASE assert guardrail.unreachable_fallback == "fail_closed" - def test_init_guardrails_v2_wiring(self, monkeypatch): + def test_init_guardrails_v2_wiring(self, monkeypatch: pytest.MonkeyPatch): """The guardrail registers and constructs via the config.yaml path.""" - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) monkeypatch.setenv("REPELLOAI_API_KEY", "test-key") init_guardrails_v2( all_guardrails=[ diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 996a3ff0824..26beaa78a46 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -1,5 +1,3 @@ -import os -import sys from fastapi.exceptions import HTTPException from unittest.mock import patch, AsyncMock from httpx import Response, Request @@ -12,19 +10,15 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im PromptSecurityGuardrail, ) -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path import litellm from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -def test_prompt_security_guard_config(monkeypatch): +def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): """Test guardrail initialization with proper configuration""" - litellm.set_verbose = True - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) + monkeypatch.setattr(litellm, "callbacks", []) - # Set environment variables for testing monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -42,21 +36,19 @@ def test_prompt_security_guard_config(monkeypatch): config_file_path="", ) - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] + registered = [c for c in litellm.callbacks if isinstance(c, PromptSecurityGuardrail)] + assert len(registered) == 1 + assert registered[0].guardrail_name == "prompt_security" + assert registered[0].default_on is True + assert registered[0].event_hook == "during_call" -def test_prompt_security_guard_config_no_api_key(): +def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch): """Test that initialization fails when API key is missing""" - litellm.set_verbose = True - litellm.guardrail_name_config_map = {} + monkeypatch.setattr(litellm, "guardrail_name_config_map", {}) - # Ensure API key is not in environment - if "PROMPT_SECURITY_API_KEY" in os.environ: - del os.environ["PROMPT_SECURITY_API_KEY"] - if "PROMPT_SECURITY_API_BASE" in os.environ: - del os.environ["PROMPT_SECURITY_API_BASE"] + monkeypatch.delenv("PROMPT_SECURITY_API_KEY", raising=False) + monkeypatch.delenv("PROMPT_SECURITY_API_BASE", raising=False) with pytest.raises( PromptSecurityGuardrailMissingSecrets, @@ -78,7 +70,7 @@ def test_prompt_security_guard_config_no_api_key(): @pytest.mark.asyncio -async def test_apply_guardrail_block_request(monkeypatch): +async def test_apply_guardrail_block_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail blocks malicious prompts""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -126,13 +118,9 @@ async def test_apply_guardrail_block_request(monkeypatch): assert "prompt_injection" in str(excinfo.value.detail) assert "jailbreak" in str(excinfo.value.detail) - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_apply_guardrail_modify_request(monkeypatch): +async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail modifies prompts when needed""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -177,13 +165,9 @@ async def test_apply_guardrail_modify_request(monkeypatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_apply_guardrail_allow_request(monkeypatch): +async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail allows safe prompts""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -220,13 +204,9 @@ async def test_apply_guardrail_allow_request(monkeypatch): assert result == inputs - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_apply_guardrail_block_response(monkeypatch): +async def test_apply_guardrail_block_response(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail blocks malicious responses""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -267,13 +247,9 @@ async def test_apply_guardrail_block_response(monkeypatch): assert "Blocked by Prompt Security" in str(excinfo.value.detail) assert "pii_exposure" in str(excinfo.value.detail) - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_apply_guardrail_modify_response(monkeypatch): +async def test_apply_guardrail_modify_response(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail modifies responses when needed""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -311,13 +287,9 @@ async def test_apply_guardrail_modify_response(monkeypatch): assert result["texts"] == ["Your SSN is [REDACTED]"] - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_file_sanitization(monkeypatch): +async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch): """Test file sanitization for images""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -401,13 +373,9 @@ async def test_file_sanitization(monkeypatch): # Should complete without errors and return the data assert result is not None - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_file_sanitization_block(monkeypatch): +async def test_file_sanitization_block(monkeypatch: pytest.MonkeyPatch): """Test that file sanitization blocks malicious files""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -485,13 +453,9 @@ async def test_file_sanitization_block(monkeypatch): assert "File blocked by Prompt Security" in str(excinfo.value.detail) assert "malware_detected" in str(excinfo.value.detail) - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_user_api_key_alias_forwarding(monkeypatch): +async def test_user_api_key_alias_forwarding(monkeypatch: pytest.MonkeyPatch): """Test that user API key alias is properly sent via headers and payload""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -530,12 +494,9 @@ async def test_user_api_key_alias_forwarding(monkeypatch): payload = call_kwargs["json"] assert payload["user"] == "vk-alias" - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_role_filtering(monkeypatch): +async def test_role_filtering(monkeypatch: pytest.MonkeyPatch): """Test that tool/function messages are filtered out by default""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -594,13 +555,9 @@ async def test_role_filtering(monkeypatch): assert len(sent_messages) == 3 assert all(msg["role"] in ["system", "user", "assistant"] for msg in sent_messages) - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - @pytest.mark.asyncio -async def test_check_tool_results_enabled(monkeypatch): +async def test_check_tool_results_enabled(monkeypatch: pytest.MonkeyPatch): """Test with check_tool_results=True: transforms tool/function to 'other' role""" monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") @@ -680,7 +637,3 @@ async def test_check_tool_results_enabled(monkeypatch): assert "indirect_prompt_injection" in str(excinfo.value.detail) - # Clean up - del os.environ["PROMPT_SECURITY_API_KEY"] - del os.environ["PROMPT_SECURITY_API_BASE"] - del os.environ["PROMPT_SECURITY_CHECK_TOOL_RESULTS"] From add2d23df22e9468252624880b0f1dbe72c0b251 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:07:16 -0700 Subject: [PATCH 35/43] test(e2e): bypass the proxy response cache in the mid-conversation system and fallback tests (#37915) The mid-conversation system tests prime the prompt cache by re-sending an identical /v1/messages body until its usage shows the full prefix read back three times in a row. The e2e stack runs with the litellm response cache on, so every resend after the first is served from redis with the first call's usage and the streak can never form; the three unflagged-model tests have failed on every litellm-e2e build since the consecutive-read check landed. Send cache: {"no-cache": true} on RichMessagesRequest, as test_cache_control already does, so each resend reaches the provider. The two fallback tests sent the same "say hi" / max_tokens=16 body to the gpt-5.5 fallback, so one empty (finish_reason=length) completion served the second test from the response cache and failed both. Give each test a unique prompt and leave gpt-5.5 enough tokens to emit text. --- tests/e2e/llm_translation/endpoints_client.py | 1 + tests/e2e/router/reliability_support.py | 2 +- tests/e2e/router/test_reliability_fallbacks_e2e.py | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 5df61247db2..fa33737467e 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -87,6 +87,7 @@ class RichMessagesRequest(BaseModel): max_tokens: int = 64 system: list[TextBlock] messages: list[RichMessage] + cache: dict[str, bool] = {"no-cache": True} class CompletionsRequest(BaseModel): diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 4dab0aaa3fa..cd70ac45da6 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -56,7 +56,7 @@ def chat_override( json=ReliabilityChatBody( model=model, messages=[ChatMessage(role="user", content=content)], - max_tokens=16, + max_tokens=64, stream=stream, router_settings_override=override, ), diff --git a/tests/e2e/router/test_reliability_fallbacks_e2e.py b/tests/e2e/router/test_reliability_fallbacks_e2e.py index 5b7d21c6ef7..fe2d924ae2c 100644 --- a/tests/e2e/router/test_reliability_fallbacks_e2e.py +++ b/tests/e2e/router/test_reliability_fallbacks_e2e.py @@ -49,7 +49,7 @@ class TestReliabilityFallbacks: resources.defer(lambda: client.proxy.delete_model(model_id)) resp = chat_override( - client.proxy, scoped_key, primary, "say hi", + client.proxy, scoped_key, primary, f"say hi {unique_marker()}", override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), ) _assert_served_by_fallback(resp) @@ -63,7 +63,7 @@ class TestReliabilityFallbacks: resources.defer(lambda: client.proxy.delete_model(model_id)) resp = chat_override( - client.proxy, scoped_key, primary, "say hi", + client.proxy, scoped_key, primary, f"say hi {unique_marker()}", override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), ) _assert_served_by_fallback(resp) From d74fc77eb113ba33c2f46c26a5cdffff65d1845a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:09:38 -0700 Subject: [PATCH 36/43] docs(terraform/provider): the provider now ships at the LiteLLM version (#37912) The provider is published in lockstep with LiteLLM: every dev, rc and stable release mirrors terraform/provider/ from the release commit and tags it with the LiteLLM version, alongside the aws/google modules. The 0.x line ends at 0.4.0, and the CHANGELOG headings no longer drive a release. RELEASING.md describes the new flow and how to recover a version whose goreleaser run failed; README gains a Versioning section with the re-pin note for anyone on `~> 0.4`; CHANGELOG records the change under Unreleased. goreleaser gets `prerelease: auto` so a v1.99.0-dev.1 / -rc.1 tag in the mirror is marked as a pre-release instead of becoming the repo's latest release. The registry ingests it either way. --- terraform/provider/.goreleaser.yml | 1 + terraform/provider/CHANGELOG.md | 15 +++- terraform/provider/README.md | 16 +++- terraform/provider/RELEASING.md | 139 ++++++----------------------- 4 files changed, 55 insertions(+), 116 deletions(-) diff --git a/terraform/provider/.goreleaser.yml b/terraform/provider/.goreleaser.yml index f41a29406b8..ba898ed9b2c 100644 --- a/terraform/provider/.goreleaser.yml +++ b/terraform/provider/.goreleaser.yml @@ -72,6 +72,7 @@ signs: - "--detach-sign" - "${artifact}" release: + prerelease: auto extra_files: - glob: 'terraform-registry-manifest.json' name_template: '{{ .ProjectName }}_{{ .Version }}_manifest.json' diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 7c744f04064..ff2f3f817f9 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -2,11 +2,22 @@ All notable changes to this project will be documented in this file. -The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), -and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). + +Up to `0.4.0` the provider had its own version line, cut from the headings in +this file. It now ships at the **LiteLLM version**, on every LiteLLM release +channel, built from the same commit as the proxy (see `RELEASING.md`). The +headings below no longer drive a release; they record what changed and which +LiteLLM line first carried it. A change that breaks existing configurations +or state must be called out loudly here, because the version number can no +longer signal it. ## [Unreleased] +### Changed + +- **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying + ## [0.4.0] - 2026-08-06 ### Fixed diff --git a/terraform/provider/README.md b/terraform/provider/README.md index 3b59edd97c6..fe67d6aa430 100644 --- a/terraform/provider/README.md +++ b/terraform/provider/README.md @@ -6,6 +6,18 @@ This Terraform provider allows you to manage LiteLLM resources through Infrastru This directory (`terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm)) is the source of truth for the provider. [BerriAI/terraform-provider-litellm](https://github.com/BerriAI/terraform-provider-litellm) is a thin release mirror that the public Terraform Registry ingests from; do not open PRs there. Changes land here, where CI builds the provider, runs its tests, and statically audits every endpoint the provider calls against the proxy's generated OpenAPI schema (`tools/endpointaudit/`), so the provider cannot drift from the LiteLLM API silently. Releases are published by mirroring this directory into the split repo and tagging it, which triggers the goreleaser workflow there (see `RELEASING.md`) +## Versioning + +The provider version **is the LiteLLM version**. Every LiteLLM release (dev, rc and stable) publishes the provider at the same version as the proxy, built from the same commit, so `1.99.0` of the provider is the one that shipped with `1.99.0` of the proxy and was audited against that proxy's API. Pin the provider to the line your proxy runs: + +```hcl +version = "~> 1.99.0" +``` + +Pre-release versions (`1.99.0-rc.1`, `1.99.0-dev.1`) are published too; Terraform only selects one when it is pinned exactly. + +Versions `0.1.0` through `0.4.0` predate this scheme and sit on their own line. They stay in the registry, but **a `~> 0.4` constraint will never pick up another release**: re-pin to the LiteLLM version to keep receiving updates. + ## Features - Manage LiteLLM model configurations @@ -32,7 +44,7 @@ terraform { required_providers { litellm = { source = "BerriAI/litellm" - version = "~> 0.1.1" #HERE UPDATE VERSION ACCORDINGLY + version = "~> 1.99.0" # the LiteLLM version your proxy runs } } } @@ -218,6 +230,6 @@ This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENS - Always use environment variables or secure secret management solutions to handle sensitive information like API keys and AWS credentials. - Refer to the comprehensive documentation in the `docs/` directory for detailed usage examples and configuration options. -- Make sure to keep your provider version updated for the latest features and bug fixes. +- Keep the provider version in step with the LiteLLM version your proxy runs; see [Versioning](#versioning). - The provider now supports AWS cross-account access with `aws_session_name` and `aws_role_name` parameters in the model resource. - All example configurations have been consolidated into the documentation for better organization and maintenance. diff --git a/terraform/provider/RELEASING.md b/terraform/provider/RELEASING.md index 7b359047e2f..59f4c5f066c 100644 --- a/terraform/provider/RELEASING.md +++ b/terraform/provider/RELEASING.md @@ -4,7 +4,16 @@ This document describes the release process for the LiteLLM Terraform Provider. ## Overview -Releases are automated via GitHub Actions when a version tag is pushed. The workflow builds the provider for multiple platforms, signs the artifacts with GPG, and publishes them to GitHub Releases. +The provider is released **in lockstep with LiteLLM**: every LiteLLM release (dev, rc and stable) publishes the provider at the LiteLLM version, built from the same commit as the proxy. There is no separate provider release to cut. + +The flow, end to end: + +1. `BerriAI/project-releaser`'s release pipeline resolves the commit to release (`main` HEAD for dev; `main` HEAD or an operator-supplied SHA for rc/stable) and passes the release approval gate +2. Its componentized terraform job rsyncs `terraform/provider/` from that commit into `BerriAI/terraform-provider-litellm`, commits, and pushes the tag `v` (for example `v1.99.0`, `v1.99.0-rc.1`, `v1.99.0-dev.1`), alongside the `terraform-aws-litellm` / `terraform-google-litellm` module mirrors which get the same tag +3. The tag push triggers the mirror's own `Release` workflow (goreleaser): multi-platform build, GPG-signed checksums, GitHub release. It runs unattended; project-releaser does not wait for it +4. The public Terraform Registry ingests the GitHub release as provider version `` + +`terraform/provider/` only exists from LiteLLM ~1.95, so a stable patch cut from an older line skips the provider and publishes only the modules. ## Prerequisites @@ -68,113 +77,26 @@ Before publishing to the Terraform Registry: **Note**: The public key fingerprint must match the key used to sign the provider releases. -## Release Steps +## What a change needs -### 1. Prepare the Release +1. **Land it in `BerriAI/litellm`.** Open a PR against `litellm_internal_staging` with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more +2. **Wait for the next LiteLLM release.** The nightly dev release carries it within a day; it reaches a stable version on the next stable cut +3. **Verify** (optional): the version appears at https://registry.terraform.io/providers/BerriAI/litellm and https://github.com/BerriAI/terraform-provider-litellm/releases. If the tag is on the mirror but there is no release, the goreleaser run failed: https://github.com/BerriAI/terraform-provider-litellm/actions -Before creating a release: +Locally, before opening the PR: -1. **Update CHANGELOG.md** - - Move items from `[Unreleased]` section to a new version section - - Follow [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) format - - Use [Semantic Versioning](https://semver.org/spec/v2.0.0.html) for version numbers - - Include all notable changes since the last release +```bash +make test +make build +``` - Example: - ```markdown - ## [0.1.2] - 2026-02-20 +## Out-of-band publish or recovery - ### Added - - New feature description +Dispatch `Build and Publish Componentized Images + Chart` in `BerriAI/project-releaser` by hand with only `publish_terraform` enabled and the `git_ref` / `tag` of the release to (re)publish. The run waits on project-releaser's release approval, then mirrors and tags exactly as the pipeline does. - ### Fixed - - Bug fix description +The mirror is push-only: do not commit or tag `BerriAI/terraform-provider-litellm` directly. The publish refuses to overwrite an existing tag; a version that failed in goreleaser is recovered by re-running the mirror's `Release` workflow for that tag, not by re-tagging. - ### Changed - - Changed behavior description - ``` - -2. **Verify tests pass** - ```bash - make test - ``` - -3. **Verify the build works locally** - ```bash - make build - ``` - -4. **Land the changes in BerriAI/litellm** - - Open a PR to `BerriAI/litellm` updating `terraform/provider/CHANGELOG.md` (and any source changes) and merge it - -### 2. Mirror and Tag via project-releaser - -The provider source lives at `terraform/provider/` in `BerriAI/litellm`; `BerriAI/terraform-provider-litellm` is a thin release mirror. Do not commit or tag the mirror directly - -Normally there is nothing to do here. `BerriAI/project-releaser`'s release pipeline runs the same check on every release except `adhoc`, nightly included: it reads the topmost released heading in `terraform/provider/CHANGELOG.md`, probes the mirror for `v`, and dispatches `Publish Terraform provider` only when the changelog has moved ahead of what the mirror carries. Cutting the version heading in step 1 is therefore what releases the provider, and the next release picks it up, so the wait is a day rather than a week - -Dispatch by hand only for an out-of-band release, or to recover a run that failed: - -1. Go to `BerriAI/project-releaser` > **Actions** > `Publish Terraform provider` -2. Click **Run workflow**: - - `git_ref`: full 40-char commit SHA from `BerriAI/litellm` to release from - - `provider_version`: the new version without the `v` prefix (e.g. `0.3.0`) - - `dry_run`: optional; validates without pushing - -Automatic or manual, the run waits on the `production-release` approval in `project-releaser`, then rsyncs `terraform/provider/` into the mirror repo, commits, and pushes tag `v`. That approval is the only one in the flow. The tag push triggers the mirror's `Release` workflow (goreleaser), which runs unattended - -**Important**: -- Tags must follow the format: `v..` (e.g., `v0.1.2`, `v1.0.0`) -- The workflow refuses to overwrite an existing tag; publish a new version instead - -### 3. Monitor the Release Workflow - -1. Go to: https://github.com/BerriAI/terraform-provider-litellm/actions -2. Find the "Release" workflow run for your tag -3. Monitor the progress and check for any errors - -The workflow will: -- Check out the code -- Set up Go -- Import the GPG key -- Run `go mod tidy` -- Build binaries for multiple platforms (Linux, macOS, Windows, FreeBSD) -- Create archives and checksums -- Sign the checksums with GPG -- Create a GitHub release -- Upload all artifacts - -### 4. Verify the Release - -After the workflow completes successfully: - -1. **Check the GitHub Release** - - Go to: https://github.com/BerriAI/terraform-provider-litellm/releases - - Verify the release was created with the correct version - - Confirm all artifacts are present: - - Binary archives for each platform - - SHA256SUMS file - - SHA256SUMS.sig (GPG signature) - - terraform-registry-manifest.json - -2. **Verify the signature** (optional) - ```bash - # Download the checksums and signature - wget https://github.com/BerriAI/terraform-provider-litellm/releases/download/v0.1.2/terraform-provider-litellm_0.1.2_SHA256SUMS - wget https://github.com/BerriAI/terraform-provider-litellm/releases/download/v0.1.2/terraform-provider-litellm_0.1.2_SHA256SUMS.sig - - # Verify the signature - gpg --verify terraform-provider-litellm_0.1.2_SHA256SUMS.sig terraform-provider-litellm_0.1.2_SHA256SUMS - ``` - -### 5. Publish to Terraform Registry (Optional) - -If this provider is published to the Terraform Registry: - -1. The registry should automatically detect the new release via the GitHub webhook -2. If not, you may need to manually trigger a sync on the Terraform Registry dashboard -3. Verify the new version appears at: https://registry.terraform.io/providers/BerriAI/litellm/latest +The mirror's `.github/` directory (the `Release` workflow) is the one thing the rsync preserves, so a change to the goreleaser *workflow* is a direct PR on the mirror; a change to `.goreleaser.yml` itself lands here like any other source change. ## Troubleshooting @@ -207,21 +129,15 @@ If this provider is published to the Terraform Registry: ### Tag Already Exists -**Error**: The publish workflow refuses to push because the tag already exists on the mirror +**Error**: The publish job refuses to push because the tag already exists on the mirror -**Solution**: Tags are immutable by design. Re-run the workflow with a new patch version instead of deleting or moving an existing tag +**Solution**: Tags are immutable by design and the version is the LiteLLM version, so this means the provider was already mirrored for this release. If the registry is missing the version, re-run the mirror's `Release` workflow for the existing tag rather than re-tagging ## Version Numbering -This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html): +The provider version is the LiteLLM version, verbatim: `X.Y.Z` for a stable release, `X.Y.Z-rc.N` for a release candidate and `X.Y.Z-dev.N` for a nightly. It says which proxy the provider shipped with and was audited against; it does not follow SemVer's break-signalling, so breaking changes are announced in `CHANGELOG.md` and the registry docs instead. -- **MAJOR** version (1.0.0): Incompatible API changes -- **MINOR** version (0.1.0): New functionality in a backward-compatible manner -- **PATCH** version (0.0.1): Backward-compatible bug fixes - -For pre-1.0 releases: -- Breaking changes may occur in minor versions -- Patch versions should only contain bug fixes +Versions `0.1.0` to `0.4.0` predate this and remain in the registry on their own line. A `~> 0.4` constraint never receives another release. ## Security Considerations @@ -237,5 +153,4 @@ For pre-1.0 releases: - [Terraform Provider Publishing](https://www.terraform.io/docs/registry/providers/publishing.html) - [HashiCorp GPG Signing Requirements](https://www.terraform.io/docs/registry/providers/publishing.html#signing-releases) - [GitHub Actions Secrets](https://docs.github.com/en/actions/security-guides/encrypted-secrets) -- [Semantic Versioning](https://semver.org/) - [Keep a Changelog](https://keepachangelog.com/) From 3ac339cfbbe199e1667942cf082f7025505fe293 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:10:34 -0700 Subject: [PATCH 37/43] test: stop the zai tests from leaking env and litellm globals into the session (#37834) test_zai_provider.py set LITELLM_LOCAL_MODEL_COST_MAP and litellm.model_cost directly and never put them back, so every test that ran after it in the same process saw a local cost map instead of the real one. The two respx tests did the same to litellm.disable_aiohttp_transport with no restore at all. Both now go through monkeypatch, which restores on teardown including when the test fails. The cost-map setup moves into a fixture requested by exactly the five tests that read the cost map. --- test-quality-budget.json | 2 +- .../llms/zai/test_zai_provider.py | 40 ++++++------------- 2 files changed, 14 insertions(+), 28 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 6a62d783044..fd984fd1e21 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -12,7 +12,7 @@ "limit": 469 }, "TQ005": { - "limit": 2549 + "limit": 2542 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index 61e1121257c..38ddac8d510 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -13,6 +13,12 @@ from litellm import completion from litellm.cost_calculator import cost_per_token +@pytest.fixture +def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + @pytest.fixture def zai_response(): """Mock response from Z.AI API""" @@ -51,12 +57,8 @@ def test_zai_in_provider_lists(): assert "zai" in litellm.provider_list -def test_zai_models_in_model_cost(monkeypatch): +def test_zai_models_in_model_cost(local_model_cost_map): """Test that ZAI models are in the model cost map""" - import os - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") zai_models = [ "zai/glm-4.7", @@ -75,12 +77,8 @@ def test_zai_models_in_model_cost(monkeypatch): assert litellm.model_cost[model]["litellm_provider"] == "zai" -def test_zai_glm46_cost_calculation(monkeypatch): +def test_zai_glm46_cost_calculation(local_model_cost_map): """Test the cost calculation for glm-4.6""" - import os - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") key = "zai/glm-4.6" info = litellm.model_cost[key] @@ -96,12 +94,8 @@ def test_zai_glm46_cost_calculation(monkeypatch): assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) -def test_zai_flash_model_is_free(monkeypatch): +def test_zai_flash_model_is_free(local_model_cost_map): """Test that glm-4.5-flash has zero cost""" - import os - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") key = "zai/glm-4.5-flash" info = litellm.model_cost[key] @@ -110,12 +104,8 @@ def test_zai_flash_model_is_free(monkeypatch): assert info["output_cost_per_token"] == 0 -def test_glm47_supports_reasoning(monkeypatch): +def test_glm47_supports_reasoning(local_model_cost_map): """Test that GLM-4.7 supports reasoning""" - import os - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") key = "zai/glm-4.7" assert key in litellm.model_cost, f"Model {key} not found in model_cost" @@ -124,12 +114,8 @@ def test_glm47_supports_reasoning(monkeypatch): assert info["supports_reasoning"] is True -def test_glm47_cost_calculation(monkeypatch): +def test_glm47_cost_calculation(local_model_cost_map): """Test cost calculation for GLM-4.7""" - import os - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") prompt_cost, completion_cost = cost_per_token( model="zai/glm-4.7", @@ -146,7 +132,7 @@ def test_glm47_cost_calculation(monkeypatch): async def test_zai_completion_call(respx_mock, zai_response, monkeypatch): """Test completion call with zai provider using mocked response""" monkeypatch.setenv("ZAI_API_KEY", "test-api-key") - litellm.disable_aiohttp_transport = True + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond( json=zai_response @@ -172,7 +158,7 @@ async def test_zai_completion_call(respx_mock, zai_response, monkeypatch): def test_zai_sync_completion(respx_mock, zai_response, monkeypatch): """Test synchronous completion call""" monkeypatch.setenv("ZAI_API_KEY", "test-api-key") - litellm.disable_aiohttp_transport = True + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond( json=zai_response From 092d97708d816c159f0fce8d7793ec0d05c5b261 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:21:29 -0700 Subject: [PATCH 38/43] test(s3): stop the logger tests leaking s3_callback_params on failure (#37831) Ten tests set litellm.s3_callback_params by hand. Four of them reset it to None on the last line of the test body, which only runs when the test passes; the other six wrap the body in try/finally to put the old value back. Raising inside test_s3_verify_false_handling on the current file leaves the whole callback config, bucket, endpoint and keys, set in the process for whatever runs next. monkeypatch.setattr covers both shapes and restores on failure, so the 28 TQ005 violations and the try/finally scaffolding come out together. 51 tests pass, and the wider tests/test_litellm/integrations tree is unchanged. The five TQ002 mock-echo tests in this file are left alone; those need a judgement about what S3 logging should assert, not a mechanical sweep. --- test-quality-budget.json | 2 +- tests/test_litellm/integrations/test_s3_v2.py | 316 +++++++++--------- 2 files changed, 151 insertions(+), 167 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index fd984fd1e21..a6bc189a7cf 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -12,7 +12,7 @@ "limit": 469 }, "TQ005": { - "limit": 2542 + "limit": 2514 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 8cccfd937e7..933e41d17a0 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -751,7 +751,7 @@ async def test_strip_base64_mixed_nested_objects(): @pytest.mark.asyncio -async def test_s3_verify_false_handling(): +async def test_s3_verify_false_handling(monkeypatch: pytest.MonkeyPatch): """ Test that s3_verify=False is properly handled and not treated as None. @@ -763,15 +763,19 @@ async def test_s3_verify_false_handling(): import litellm # Set up s3_callback_params with s3_verify=False - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, # This should NOT be ignored - "s3_use_ssl": False, # This should also NOT be ignored - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, # This should NOT be ignored + "s3_use_ssl": False, # This should also NOT be ignored + }, + ) with patch("asyncio.create_task"): with patch( @@ -801,12 +805,9 @@ async def test_s3_verify_false_handling(): "ssl_verify": False }, f"Expected ssl_verify=False in params, got {call_kwargs.get('params')}" - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio -async def test_s3_verify_none_handling(): +async def test_s3_verify_none_handling(monkeypatch: pytest.MonkeyPatch): """ Test that s3_verify=None uses default behavior. """ @@ -815,12 +816,16 @@ async def test_s3_verify_none_handling(): import litellm # Set up s3_callback_params without s3_verify - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_aws_access_key_id": "test-key", - "s3_aws_secret_access_key": "test-secret", - "s3_region_name": "us-east-1", - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_aws_access_key_id": "test-key", + "s3_aws_secret_access_key": "test-secret", + "s3_region_name": "us-east-1", + }, + ) with patch("asyncio.create_task"): with patch( @@ -846,12 +851,9 @@ async def test_s3_verify_none_handling(): assert call_kwargs["params"].get("ssl_verify") is None # Either params is None or params={'ssl_verify': None} is acceptable - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio -async def test_s3_verify_false_creates_httpx_client_with_verify_false(): +async def test_s3_verify_false_creates_httpx_client_with_verify_false(monkeypatch: pytest.MonkeyPatch): """ Test that when s3_verify=False, the actual httpx client has verify=False. @@ -862,14 +864,18 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(): import litellm # Set up s3_callback_params with s3_verify=False - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, + ) with patch("asyncio.create_task"): # Create logger - this creates the httpx client @@ -888,12 +894,9 @@ async def test_s3_verify_false_creates_httpx_client_with_verify_false(): httpx_client._verify is False ), f"Expected httpx client _verify=False, got {httpx_client._verify}" - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio -async def test_s3_verify_false_async_client(): +async def test_s3_verify_false_async_client(monkeypatch: pytest.MonkeyPatch): """ Test that the async httpx client respects s3_verify=False. """ @@ -903,14 +906,18 @@ async def test_s3_verify_false_async_client(): from litellm.types.integrations.s3_v2 import s3BatchLoggingElement # Set up s3_callback_params with s3_verify=False - litellm.s3_callback_params = { - "s3_bucket_name": "test-bucket", - "s3_endpoint_url": "https://localhost:443", - "s3_aws_access_key_id": "minioadmin", - "s3_aws_secret_access_key": "minioadmin", - "s3_region_name": "us-east-1", - "s3_verify": False, - } + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "test-bucket", + "s3_endpoint_url": "https://localhost:443", + "s3_aws_access_key_id": "minioadmin", + "s3_aws_secret_access_key": "minioadmin", + "s3_region_name": "us-east-1", + "s3_verify": False, + }, + ) with patch("asyncio.create_task"): logger = S3Logger() @@ -945,9 +952,6 @@ async def test_s3_verify_false_async_client(): httpx_client._verify is False ), f"Expected async httpx client _verify=False, got {httpx_client._verify}" - # Clean up - litellm.s3_callback_params = None - @pytest.mark.asyncio async def test_strip_base64_recursive_redaction(): @@ -1169,26 +1173,22 @@ def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # -------------------------------------------------------------- -def test_s3_callback_params_override_uses_alternate_dict(): +def test_s3_callback_params_override_uses_alternate_dict(monkeypatch): """`s3_callback_params_override` makes the logger read its config from the override dict instead of `litellm.s3_callback_params`.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "normal-bucket"} - try: - logger = S3Logger( - s3_callback_params_override={ - "s3_bucket_name": "audit-bucket", - "s3_path": "audit-prefix", - "s3_region_name": "us-west-2", - } - ) - assert logger.s3_bucket_name == "audit-bucket" - assert logger.s3_path == "audit-prefix" - assert logger.s3_region_name == "us-west-2" - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-bucket"}) + logger = S3Logger( + s3_callback_params_override={ + "s3_bucket_name": "audit-bucket", + "s3_path": "audit-prefix", + "s3_region_name": "us-west-2", + } + ) + assert logger.s3_bucket_name == "audit-bucket" + assert logger.s3_path == "audit-prefix" + assert logger.s3_region_name == "us-west-2" def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): @@ -1198,43 +1198,31 @@ def test_s3_callback_params_override_does_not_mutate_inputs(monkeypatch): monkeypatch.setenv("MY_AUDIT_BUCKET", "resolved-bucket") override = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"} - original_global = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"} - try: - logger = S3Logger(s3_callback_params_override=override) - assert logger.s3_bucket_name == "resolved-bucket" - assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - assert ( - litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" - ) - finally: - litellm.s3_callback_params = original_global + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "os.environ/MY_AUDIT_BUCKET"}) + logger = S3Logger(s3_callback_params_override=override) + assert logger.s3_bucket_name == "resolved-bucket" + assert override["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" + assert ( + litellm.s3_callback_params["s3_bucket_name"] == "os.environ/MY_AUDIT_BUCKET" + ) -def test_s3_callback_params_override_none_falls_back_to_global(): +def test_s3_callback_params_override_none_falls_back_to_global(monkeypatch): """No override → behaves exactly as today (reads `litellm.s3_callback_params`).""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "from-global"} - try: - logger = S3Logger() - assert logger.s3_bucket_name == "from-global" - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"}) + logger = S3Logger() + assert logger.s3_bucket_name == "from-global" -def test_s3_callback_params_override_empty_dict_is_opt_in(): +def test_s3_callback_params_override_empty_dict_is_opt_in(monkeypatch): """An empty override dict skips the global entirely (env/IAM-only config).""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "from-global"} - try: - logger = S3Logger(s3_callback_params_override={}) - assert logger.s3_bucket_name is None - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "from-global"}) + logger = S3Logger(s3_callback_params_override={}) + assert logger.s3_bucket_name is None def _expected_content_md5(payload: dict) -> str: @@ -1374,20 +1362,20 @@ async def test_async_upload_sets_server_side_encryption_header_when_configured() assert headers["x-amz-server-side-encryption"] == "aws:kms" -def test_s3_server_side_encryption_read_from_callback_params(): +def test_s3_server_side_encryption_read_from_callback_params(monkeypatch): """s3_server_side_encryption can be configured via s3_callback_params.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "aws:kms" - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" @pytest.mark.asyncio @@ -1505,21 +1493,21 @@ async def test_async_upload_omits_kms_key_id_header_when_not_configured(): assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers -def test_s3_sse_kms_key_id_read_from_callback_params(): +def test_s3_sse_kms_key_id_read_from_callback_params(monkeypatch): """s3_sse_kms_key_id can be configured via s3_callback_params.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - } - try: - logger = S3Logger() - assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, + ) + logger = S3Logger() + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") @pytest.mark.asyncio @@ -1561,83 +1549,79 @@ async def test_async_upload_infers_aws_kms_when_only_key_id_set(): ) -def test_s3_sse_kms_key_id_read_from_audit_override_params(): +def test_s3_sse_kms_key_id_read_from_audit_override_params(monkeypatch): """The audit-log override path must honor s3_sse_kms_key_id too.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = {"s3_bucket_name": "normal-logs-bucket"} - try: - logger = S3Logger( - s3_callback_params_override={ - "s3_bucket_name": "audit-logs-bucket", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id", - } - ) - assert logger.s3_bucket_name == "audit-logs-bucket" - assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id") - finally: - litellm.s3_callback_params = original + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_bucket_name": "normal-logs-bucket"}) + logger = S3Logger( + s3_callback_params_override={ + "s3_bucket_name": "audit-logs-bucket", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/audit-key-id", + } + ) + assert logger.s3_bucket_name == "audit-logs-bucket" + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/audit-key-id") -def test_kms_key_id_dropped_when_algorithm_is_not_kms(): +def test_kms_key_id_dropped_when_algorithm_is_not_kms(monkeypatch): """ AES256 plus a KMS key id is an invalid S3 combination; the key id must be dropped at init so uploads keep working instead of silently 400ing. """ import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "AES256", - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "AES256" - assert logger.s3_sse_kms_key_id is None - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "AES256", + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "AES256" + assert logger.s3_sse_kms_key_id is None -def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(): +def test_non_string_algorithm_is_dropped_and_valid_key_id_is_rescued(monkeypatch): """ A YAML boolean in s3_server_side_encryption must not crash logger init and must not discard the valid key id; aws:kms is inferred from the key id. """ import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": True, - "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "aws:kms" - assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": True, + "s3_sse_kms_key_id": "arn:aws:kms:us-east-1:111122223333:key/test-key-id", + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" + assert logger.s3_sse_kms_key_id == ("arn:aws:kms:us-east-1:111122223333:key/test-key-id") -def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(): +def test_non_string_key_id_is_dropped_and_valid_algorithm_is_kept(monkeypatch): """A mistyped key id (unquoted YAML number) must not disable the valid algorithm.""" import litellm - original = litellm.s3_callback_params - litellm.s3_callback_params = { - "s3_bucket_name": "from-global", - "s3_server_side_encryption": "aws:kms", - "s3_sse_kms_key_id": 12345, - } - try: - logger = S3Logger() - assert logger.s3_server_side_encryption == "aws:kms" - assert logger.s3_sse_kms_key_id is None - finally: - litellm.s3_callback_params = original + monkeypatch.setattr( + litellm, + "s3_callback_params", + { + "s3_bucket_name": "from-global", + "s3_server_side_encryption": "aws:kms", + "s3_sse_kms_key_id": 12345, + }, + ) + logger = S3Logger() + assert logger.s3_server_side_encryption == "aws:kms" + assert logger.s3_sse_kms_key_id is None _ACCESS_KEY = "AKIAIOSFODNN7EXAMPLE" From ce1321466bb0d31c05e76a5b4de554c6bf4e19cf Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:32:24 -0700 Subject: [PATCH 39/43] test(http-handler): drop the save/restore scaffolding around litellm globals (#37839) Nine tests in test_http_handler.py captured litellm.disable_aiohttp_transport, force_ipv4, ssl_ecdh_curve or the request_timeout pair, wrapped their whole body in a try, and put the value back in a finally. monkeypatch.setattr does all of that, so the captures, the try and the finally go away and the bodies lose a level of indentation. The class-scoped restore_request_timeout fixture existed only for that same bookkeeping and goes with them. litellm.in_memory_llm_clients_cache is left alone on purpose: the eviction tests assert a handler is garbage collected, and monkeypatch holds the replaced value alive until teardown, which keeps the weakref they check from clearing. --- test-quality-budget.json | 2 +- .../llms/custom_httpx/test_http_handler.py | 258 +++++++----------- 2 files changed, 107 insertions(+), 153 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index a6bc189a7cf..e5184f2049e 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -12,7 +12,7 @@ "limit": 469 }, "TQ005": { - "limit": 2514 + "limit": 2488 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index fa1c7308c6f..641fae12bc2 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -131,79 +131,62 @@ def test_sync_post_streaming_status_error_should_not_wait_forever_for_body( @pytest.mark.asyncio async def test_ssl_security_level(monkeypatch): # Ensure aiohttp transport is enabled for this test - original_disable = litellm.disable_aiohttp_transport - litellm.disable_aiohttp_transport = False + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) - try: - with patch.dict(os.environ, clear=True): - # Set environment variable for SSL security level - monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1") + with patch.dict(os.environ, clear=True): + # Set environment variable for SSL security level + monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1") - # Create async client with SSL verification disabled to isolate SSL context testing - client = AsyncHTTPHandler() + # Create async client with SSL verification disabled to isolate SSL context testing + client = AsyncHTTPHandler() - try: - # Get the transport (should be LiteLLMAiohttpTransport) - transport = client.client._transport - assert isinstance(transport, LiteLLMAiohttpTransport) + try: + # Get the transport (should be LiteLLMAiohttpTransport) + transport = client.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport) - # Get the aiohttp ClientSession - client_session = transport._get_valid_client_session() + # Get the aiohttp ClientSession + client_session = transport._get_valid_client_session() - # Get the connector from the session - connector = client_session.connector - assert isinstance(connector, TCPConnector) + # Get the connector from the session + connector = client_session.connector + assert isinstance(connector, TCPConnector) - # Get the SSL context from the connector - ssl_context = connector._ssl + # Get the SSL context from the connector + ssl_context = connector._ssl - # Verify that the SSL context exists and has the correct cipher string - assert isinstance(ssl_context, ssl.SSLContext) - finally: - await client.close() - finally: - # Restore original setting - litellm.disable_aiohttp_transport = original_disable + # Verify that the SSL context exists and has the correct cipher string + assert isinstance(ssl_context, ssl.SSLContext) + finally: + await client.close() @pytest.mark.asyncio -async def test_force_ipv4_transport(): +async def test_force_ipv4_transport(monkeypatch: pytest.MonkeyPatch): """Test transport creation with force_ipv4 enabled""" - original_force_ipv4 = litellm.force_ipv4 - original_disable = litellm.disable_aiohttp_transport - litellm.force_ipv4 = True - litellm.disable_aiohttp_transport = True + monkeypatch.setattr(litellm, "force_ipv4", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - try: - transport = AsyncHTTPHandler._create_async_transport() + transport = AsyncHTTPHandler._create_async_transport() - # Should get an AsyncHTTPTransport (no real HTTP call — avoids CI hangs) - assert isinstance(transport, httpx.AsyncHTTPTransport) - finally: - litellm.force_ipv4 = original_force_ipv4 - litellm.disable_aiohttp_transport = original_disable + # Should get an AsyncHTTPTransport (no real HTTP call — avoids CI hangs) + assert isinstance(transport, httpx.AsyncHTTPTransport) @pytest.mark.asyncio -async def test_aiohttp_disabled_transport(): +async def test_aiohttp_disabled_transport(monkeypatch: pytest.MonkeyPatch): """Test transport creation with aiohttp disabled""" - original_disable = litellm.disable_aiohttp_transport - original_force_ipv4 = litellm.force_ipv4 - litellm.disable_aiohttp_transport = True - litellm.force_ipv4 = False + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "force_ipv4", False) - try: - transport = AsyncHTTPHandler._create_async_transport() + transport = AsyncHTTPHandler._create_async_transport() - # Should get None when both aiohttp is disabled and force_ipv4 is False - assert transport is None - finally: - litellm.disable_aiohttp_transport = original_disable - litellm.force_ipv4 = original_force_ipv4 + # Should get None when both aiohttp is disabled and force_ipv4 is False + assert transport is None @pytest.mark.asyncio -async def test_ssl_verification_with_aiohttp_transport(): +async def test_ssl_verification_with_aiohttp_transport(monkeypatch: pytest.MonkeyPatch): """ Test aiohttp respects ssl_verify=False @@ -213,38 +196,33 @@ async def test_ssl_verification_with_aiohttp_transport(): import aiohttp # Ensure aiohttp transport is enabled for this test - original_disable = litellm.disable_aiohttp_transport - litellm.disable_aiohttp_transport = False + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) + + litellm_async_client = AsyncHTTPHandler(ssl_verify=False) try: - litellm_async_client = AsyncHTTPHandler(ssl_verify=False) + transport = litellm_async_client.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport) + transport_connector = transport._get_valid_client_session().connector + assert isinstance(transport_connector, TCPConnector) + aiohttp_session = aiohttp.ClientSession( + connector=aiohttp.TCPConnector(ssl=False) + ) try: - transport = litellm_async_client.client._transport - assert isinstance(transport, LiteLLMAiohttpTransport) - transport_connector = transport._get_valid_client_session().connector - assert isinstance(transport_connector, TCPConnector) + aiohttp_connector = aiohttp_session.connector + assert isinstance(aiohttp_connector, aiohttp.TCPConnector) - aiohttp_session = aiohttp.ClientSession( - connector=aiohttp.TCPConnector(ssl=False) - ) - try: - aiohttp_connector = aiohttp_session.connector - assert isinstance(aiohttp_connector, aiohttp.TCPConnector) - - # assert both litellm transport and aiohttp session have ssl_verify=False - assert transport_connector._ssl == aiohttp_connector._ssl - finally: - await aiohttp_session.close() + # assert both litellm transport and aiohttp session have ssl_verify=False + assert transport_connector._ssl == aiohttp_connector._ssl finally: - await litellm_async_client.close() + await aiohttp_session.close() finally: - # Restore original setting - litellm.disable_aiohttp_transport = original_disable + await litellm_async_client.close() @pytest.mark.asyncio -async def test_ssl_verification_with_shared_session(): +async def test_ssl_verification_with_shared_session(monkeypatch: pytest.MonkeyPatch): """ Test that ssl_verify=False is respected even with shared sessions. @@ -257,67 +235,55 @@ async def test_ssl_verification_with_shared_session(): import aiohttp # Ensure aiohttp transport is enabled for this test - original_disable = litellm.disable_aiohttp_transport - litellm.disable_aiohttp_transport = False + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) + + shared_session = aiohttp.ClientSession() try: - # Create a shared session (simulating what happens in production) - shared_session = aiohttp.ClientSession() + # Create transport with shared session and ssl_verify=False + transport = AsyncHTTPHandler._create_aiohttp_transport( + ssl_verify=False, + shared_session=shared_session, + ) - try: - # Create transport with shared session and ssl_verify=False - transport = AsyncHTTPHandler._create_aiohttp_transport( - ssl_verify=False, - shared_session=shared_session, - ) + # Verify the transport uses the shared session + assert transport.client is shared_session - # Verify the transport uses the shared session - assert transport.client is shared_session - - # Verify the SSL setting is stored in the transport for per-request use - assert transport._ssl_verify is False - finally: - await shared_session.close() + # Verify the SSL setting is stored in the transport for per-request use + assert transport._ssl_verify is False finally: - # Restore original setting - litellm.disable_aiohttp_transport = original_disable + await shared_session.close() @pytest.mark.asyncio -async def test_ssl_context_with_shared_session(): +async def test_ssl_context_with_shared_session(monkeypatch: pytest.MonkeyPatch): """ Test that ssl_context is respected even with shared sessions. """ import aiohttp # Ensure aiohttp transport is enabled for this test - original_disable = litellm.disable_aiohttp_transport - litellm.disable_aiohttp_transport = False + monkeypatch.setattr(litellm, "disable_aiohttp_transport", False) + + custom_ssl_context = ssl.create_default_context() + + # Create a shared session + shared_session = aiohttp.ClientSession() try: - # Create a custom SSL context - custom_ssl_context = ssl.create_default_context() + # Create transport with shared session and custom ssl_context + transport = AsyncHTTPHandler._create_aiohttp_transport( + ssl_context=custom_ssl_context, + shared_session=shared_session, + ) - # Create a shared session - shared_session = aiohttp.ClientSession() + # Verify the transport uses the shared session + assert transport.client is shared_session - try: - # Create transport with shared session and custom ssl_context - transport = AsyncHTTPHandler._create_aiohttp_transport( - ssl_context=custom_ssl_context, - shared_session=shared_session, - ) - - # Verify the transport uses the shared session - assert transport.client is shared_session - - # Verify the SSL context is stored in the transport for per-request use - assert transport._ssl_verify is custom_ssl_context - finally: - await shared_session.close() + # Verify the SSL context is stored in the transport for per-request use + assert transport._ssl_verify is custom_ssl_context finally: - # Restore original setting - litellm.disable_aiohttp_transport = original_disable + await shared_session.close() def test_get_ssl_configuration(): @@ -563,26 +529,22 @@ def test_ssl_ecdh_curve( if env_curve: monkeypatch.setenv("SSL_ECDH_CURVE", env_curve) - original_value = litellm.ssl_ecdh_curve - try: - litellm.ssl_ecdh_curve = litellm_curve + monkeypatch.setattr(litellm, "ssl_ecdh_curve", litellm_curve) - # Create a real SSL context and patch set_ecdh_curve on it - # We need a real SSLContext instance (not a MagicMock) because _create_ssl_context - # calls methods like set_ciphers() and minimum_version that require a real context. - # We patch set_ecdh_curve specifically to verify it's called with the correct curve. - real_ssl_context = ssl.create_default_context() - with patch("ssl.create_default_context", return_value=real_ssl_context): - with patch.object(real_ssl_context, "set_ecdh_curve") as mock_set_curve: - ssl_context = get_ssl_configuration() + # Create a real SSL context and patch set_ecdh_curve on it + # We need a real SSLContext instance (not a MagicMock) because _create_ssl_context + # calls methods like set_ciphers() and minimum_version that require a real context. + # We patch set_ecdh_curve specifically to verify it's called with the correct curve. + real_ssl_context = ssl.create_default_context() + with patch("ssl.create_default_context", return_value=real_ssl_context): + with patch.object(real_ssl_context, "set_ecdh_curve") as mock_set_curve: + ssl_context = get_ssl_configuration() - if should_call: - mock_set_curve.assert_called_once_with(expected_curve) - else: - mock_set_curve.assert_not_called() - assert isinstance(ssl_context, ssl.SSLContext) - finally: - litellm.ssl_ecdh_curve = original_value + if should_call: + mock_set_curve.assert_called_once_with(expected_curve) + else: + mock_set_curve.assert_not_called() + assert isinstance(ssl_context, ssl.SSLContext) def test_default_user_agent_is_litellm_version(monkeypatch): @@ -753,46 +715,38 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: no per-model timeout (e.g. Bedrock) hung for 600s. """ - @pytest.fixture - def restore_request_timeout(self): - original_value = litellm.request_timeout - original_flag = litellm.request_timeout_explicitly_set - try: - yield - finally: - litellm.request_timeout = original_value - litellm.request_timeout_explicitly_set = original_flag - - def test_default_when_request_timeout_unset(self, restore_request_timeout): + def test_default_when_request_timeout_unset(self, monkeypatch: pytest.MonkeyPatch): from litellm.llms.custom_httpx.http_handler import ( _DEFAULT_TIMEOUT, _default_cached_client_timeout, ) - litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS - litellm.request_timeout_explicitly_set = False + monkeypatch.setattr( + litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS + ) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT - def test_uses_explicit_request_timeout(self, restore_request_timeout): + def test_uses_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch): from litellm.llms.custom_httpx.http_handler import ( _default_cached_client_timeout, ) - litellm.request_timeout = 300 - litellm.request_timeout_explicitly_set = True + monkeypatch.setattr(litellm, "request_timeout", 300) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) resolved = _default_cached_client_timeout() assert resolved.read == 300.0 assert resolved.connect == 5.0 def test_cached_async_client_built_with_explicit_request_timeout( - self, restore_request_timeout + self, monkeypatch: pytest.MonkeyPatch ): from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.utils import LlmProviders - litellm.request_timeout = 300 - litellm.request_timeout_explicitly_set = True + monkeypatch.setattr(litellm, "request_timeout", 300) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) litellm.in_memory_llm_clients_cache = LLMClientCache() client = get_async_httpx_client(llm_provider=LlmProviders.BEDROCK) assert client.timeout.read == 300.0 From 322293ad952e88219eea8d6457e3a0c176c1d120 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:43:03 -0700 Subject: [PATCH 40/43] test(interactions): drop the save/restore scaffolding around the legacy flag (#37841) Seven tests captured litellm.use_legacy_interactions_schema, wrapped their body in a try, and put it back in a finally. monkeypatch.setattr does that, so the capture, the try and the finally go and the bodies lose an indentation level. The remaining hand-rolled restores stay. They hold the flag only across the iterator's constructor and put it back before the test iterates, so handing them to monkeypatch would widen that window to the whole test and change what the streaming assertions run against. --- test-quality-budget.json | 2 +- ...test_gemini_interactions_transformation.py | 240 ++++++++---------- 2 files changed, 108 insertions(+), 134 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index e5184f2049e..1263ed82a65 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -12,7 +12,7 @@ "limit": 469 }, "TQ005": { - "limit": 2488 + "limit": 2474 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py index 524589abf5e..05b0bde16bb 100644 --- a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py +++ b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py @@ -86,29 +86,21 @@ class TestValidateEnvironment: assert headers["X-Custom"] == "value" assert headers["x-goog-api-key"] == "test-key" - def test_api_revision_new_schema_by_default(self, config): + def test_api_revision_new_schema_by_default(self, config, monkeypatch: pytest.MonkeyPatch): # Default: use_legacy_interactions_schema=False → new steps schema - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - headers = config.validate_environment( - headers={}, model="gemini-2.5-flash", litellm_params=None - ) - assert headers["Api-Revision"] == "2026-05-20" - finally: - litellm.use_legacy_interactions_schema = original + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + headers = config.validate_environment( + headers={}, model="gemini-2.5-flash", litellm_params=None + ) + assert headers["Api-Revision"] == "2026-05-20" - def test_api_revision_legacy_schema_when_flag_set(self, config): + def test_api_revision_legacy_schema_when_flag_set(self, config, monkeypatch: pytest.MonkeyPatch): # Flag on → legacy outputs schema until June 8, 2026 - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = True - headers = config.validate_environment( - headers={}, model="gemini-2.5-flash", litellm_params=None - ) - assert headers["Api-Revision"] == "2026-05-07" - finally: - litellm.use_legacy_interactions_schema = original + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", True) + headers = config.validate_environment( + headers={}, model="gemini-2.5-flash", litellm_params=None + ) + assert headers["Api-Revision"] == "2026-05-07" class TestGetCompleteUrl: @@ -561,23 +553,19 @@ class TestInteractionOperationUrls: class TestTransformRequestSchemaCoalescing: """Test new-schema request coalescing (Api-Revision: 2026-05-20).""" - def test_response_mime_type_folded_into_response_format(self, config): - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="summarise", - optional_params={ - "response_mime_type": "application/json", - "response_format": {"type": "object", "properties": {}}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + def test_response_mime_type_folded_into_response_format(self, config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="summarise", + optional_params={ + "response_mime_type": "application/json", + "response_format": {"type": "object", "properties": {}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) # response_mime_type must not appear as a top-level body key assert "response_mime_type" not in body @@ -586,25 +574,21 @@ class TestTransformRequestSchemaCoalescing: assert rf["mime_type"] == "application/json" assert "schema" in rf - def test_image_config_moved_to_response_format(self, config): - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="draw a sunset", - optional_params={ - "generation_config": { - "temperature": 0.7, - "image_config": {"aspect_ratio": "1:1", "image_size": "1K"}, - } - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + def test_image_config_moved_to_response_format(self, config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw a sunset", + optional_params={ + "generation_config": { + "temperature": 0.7, + "image_config": {"aspect_ratio": "1:1", "image_size": "1K"}, + } + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) # image_config removed from generation_config assert "image_config" not in body.get("generation_config", {}) @@ -613,95 +597,85 @@ class TestTransformRequestSchemaCoalescing: assert rf["type"] == "image" assert rf["aspect_ratio"] == "1:1" - def test_response_mime_type_skipped_when_response_format_is_list(self, config): + def test_response_mime_type_skipped_when_response_format_is_list(self, config, monkeypatch: pytest.MonkeyPatch): """Lists are already polymorphic; do not wrap them into schema.""" - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - rf_list = [ - {"type": "text", "mime_type": "application/json"}, - {"type": "image", "aspect_ratio": "1:1"}, - ] - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="multimodal", - optional_params={ - "response_format": rf_list, - "response_mime_type": "application/json", - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + rf_list = [ + {"type": "text", "mime_type": "application/json"}, + {"type": "image", "aspect_ratio": "1:1"}, + ] + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="multimodal", + optional_params={ + "response_format": rf_list, + "response_mime_type": "application/json", + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert body["response_format"] == rf_list assert "response_mime_type" not in body def test_image_config_appended_to_response_format_list_without_mutating_input( - self, config + self, + config, + monkeypatch: pytest.MonkeyPatch, ): """When response_format is already a list, image_config must not mutate optional_params.""" - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = False - text_rf = {"type": "text", "mime_type": "application/json"} - optional_params = { - "response_format": [text_rf], - "generation_config": { - "image_config": {"aspect_ratio": "16:9", "image_size": "2K"}, - }, - } - original_rf = optional_params["response_format"] + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", False) + text_rf = {"type": "text", "mime_type": "application/json"} + optional_params = { + "response_format": [text_rf], + "generation_config": { + "image_config": {"aspect_ratio": "16:9", "image_size": "2K"}, + }, + } + original_rf = optional_params["response_format"] - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="draw and summarise", - optional_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw and summarise", + optional_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) - assert optional_params["response_format"] is original_rf - assert len(optional_params["response_format"]) == 1 - assert body["response_format"] == [ - text_rf, - {"type": "image", "aspect_ratio": "16:9", "image_size": "2K"}, - ] + assert optional_params["response_format"] is original_rf + assert len(optional_params["response_format"]) == 1 + assert body["response_format"] == [ + text_rf, + {"type": "image", "aspect_ratio": "16:9", "image_size": "2K"}, + ] - # Retry must not append a second image entry into the caller's list. - body_retry = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="draw and summarise", - optional_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - assert len(optional_params["response_format"]) == 1 - assert body_retry["response_format"] == body["response_format"] - finally: - litellm.use_legacy_interactions_schema = original + # Retry must not append a second image entry into the caller's list. + body_retry = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw and summarise", + optional_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert len(optional_params["response_format"]) == 1 + assert body_retry["response_format"] == body["response_format"] - def test_legacy_schema_passes_fields_unchanged(self, config): - original = litellm.use_legacy_interactions_schema - try: - litellm.use_legacy_interactions_schema = True - body = config.transform_request( - model="gemini/gemini-2.5-flash", - agent=None, - input="hello", - optional_params={ - "response_mime_type": "application/json", - "generation_config": {"image_config": {"aspect_ratio": "16:9"}}, - }, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) - finally: - litellm.use_legacy_interactions_schema = original + def test_legacy_schema_passes_fields_unchanged(self, config, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "use_legacy_interactions_schema", True) + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="hello", + optional_params={ + "response_mime_type": "application/json", + "generation_config": {"image_config": {"aspect_ratio": "16:9"}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert body["response_mime_type"] == "application/json" assert body["generation_config"]["image_config"]["aspect_ratio"] == "16:9" From b9bff0998c9c89034314a81000ff8f9ff9158a01 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Aug 2026 22:54:07 -0700 Subject: [PATCH 41/43] test(bedrock): drop the leftover set_verbose from the embedding tests (#37844) Fifteen tests opened with litellm.set_verbose = True and never put it back, so the flag stayed on for everything that ran after them in the same process. Nothing in the file reads the output it produces: there is no caplog, no capsys and no assertion on a log line, so the flag was left over from debugging. Deleting it beats restoring it, since restoring keeps the noise. --- test-quality-budget.json | 2 +- .../llms/bedrock/embed/test_bedrock_embedding.py | 15 --------------- 2 files changed, 1 insertion(+), 16 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 1263ed82a65..5039143eaa0 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -12,7 +12,7 @@ "limit": 469 }, "TQ005": { - "limit": 2474 + "limit": 2459 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 9955851132c..e35365cd609 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -50,7 +50,6 @@ test_image_base64 = "data:image/png,test_image_base64_data" ) def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_response): """Test embedding functionality with bearer token authentication""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" @@ -98,7 +97,6 @@ def test_bedrock_embedding_with_env_variable_bearer_token( model, input_type, embed_response ): """Test embedding functionality with bearer token from environment variable""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "env-bearer-token-12345" @@ -130,7 +128,6 @@ def test_bedrock_embedding_with_env_variable_bearer_token( @pytest.mark.asyncio async def test_async_bedrock_embedding_with_bearer_token(): """Test async embedding functionality with bearer token authentication""" - litellm.set_verbose = True client = AsyncHTTPHandler() test_api_key = "async-bearer-token-12345" model = "bedrock/amazon.titan-embed-text-v1" @@ -160,7 +157,6 @@ async def test_async_bedrock_embedding_with_bearer_token(): def test_bedrock_embedding_with_sigv4(): """Test embedding falls back to SigV4 auth when no bearer token is provided""" - litellm.set_verbose = True model = "bedrock/amazon.titan-embed-text-v1" with patch( @@ -182,7 +178,6 @@ def test_bedrock_embedding_with_sigv4(): def test_bedrock_titan_v2_encoding_format_float(): """Test amazon.titan-embed-text-v2:0 with encoding_format=float parameter""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/amazon.titan-embed-text-v2:0" @@ -220,7 +215,6 @@ def test_bedrock_titan_v2_encoding_format_float(): def test_bedrock_titan_v2_encoding_format_base64(): """Test amazon.titan-embed-text-v2:0 with encoding_format=base64 parameter (maps to binary)""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/amazon.titan-embed-text-v2:0" @@ -260,7 +254,6 @@ def test_bedrock_titan_v2_encoding_format_base64(): def test_twelvelabs_input_type_parameter_mapping(): """Test that input_type parameter is correctly mapped to inputType for TwelveLabs models""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/twelvelabs.marengo-embed-2-7-v1:0" @@ -300,7 +293,6 @@ def test_twelvelabs_input_type_parameter_mapping(): def test_twelvelabs_input_type_parameter_mapping_async_invoke(): """Test that input_type parameter is correctly mapped to inputType for TwelveLabs async invoke models""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0" @@ -343,7 +335,6 @@ def test_twelvelabs_input_type_parameter_mapping_async_invoke(): def test_twelvelabs_missing_input_type_error(): """Test that missing input_type parameter defaults to 'text' for TwelveLabs models""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" @@ -422,7 +413,6 @@ def test_bedrock_embedding_header_forwarding(model, embed_response): Relevant Issue: https://github.com/BerriAI/litellm/pull/16042 """ - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" @@ -489,7 +479,6 @@ def test_bedrock_embedding_extra_headers_and_headers_merge(): This ensures that headers from kwargs (forwarded by proxy) and extra_headers (passed explicitly) are both included in the final headers sent to the provider. """ - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/amazon.titan-embed-text-v1" @@ -557,7 +546,6 @@ def test_bedrock_cohere_v4_embedding_response_parsing(): Test parsing of Bedrock Cohere v4 embedding response which returns a dictionary of embeddings keyed by type (e.g. 'float', 'int8') instead of a direct list. """ - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/cohere.embed-v4:0" @@ -617,7 +605,6 @@ def test_bedrock_embedding_custom_headers_with_iam_role_and_custom_api_base(): Relevant Issue: Custom headers not forwarded with IAM roles + custom api_base """ - litellm.set_verbose = True client = HTTPHandler() # Simulate IAM role credentials with session token @@ -734,7 +721,6 @@ async def test_bedrock_embedding_custom_headers_with_iam_role_and_custom_api_bas This is the async version of the test above, verifying the fix works for both sync and async embedding calls. """ - litellm.set_verbose = True client = AsyncHTTPHandler() # Simulate IAM role credentials with session token @@ -977,7 +963,6 @@ def test_bedrock_cohere_embedding_types_wrapped_as_list( Malformed input request: #/embedding_types: expected type: JSONArray, found: String when `encoding_format` is passed as a string. """ - litellm.set_verbose = True client = HTTPHandler() model = "bedrock/cohere.embed-multilingual-v3" From 7dff9953cbe1b38d49ee8ffd109251e5483375f5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 22 Aug 2026 08:23:27 -0700 Subject: [PATCH 42/43] test: drop the leftover set_verbose from eleven test files (#37845) Twenty-three tests across eleven files opened with litellm.set_verbose = True and never put it back, so the flag stayed on for everything that ran after them in the same process. None of those files read the output it produces: no caplog, no capsys, no assertion on a log line, so the flag was left over from debugging. Deleting it beats restoring it, since restoring keeps the noise. Ten of the eleven stop leaving the flag on. test_volcengine_embedding.py still ends with it set, from something it exercises rather than from the test itself, which is worth its own look. --- test-quality-budget.json | 2 +- tests/test_litellm/llms/azure/test_azure_common_utils.py | 1 - .../bedrock/embed/test_bedrock_async_invoke_embedding.py | 2 -- .../llms/bedrock/image/test_bedrock_image_bearer_token.py | 4 ---- .../bedrock/rerank/test_bedrock_rerank_header_forwarding.py | 3 --- tests/test_litellm/llms/openai/test_openai_common_utils.py | 1 - tests/test_litellm/llms/vertex_ai/test_vertex.py | 6 ------ .../llms/volcengine/test_volcengine_embedding.py | 1 - .../proxy/guardrails/guardrail_hooks/test_cato_networks.py | 2 -- .../guardrail_hooks/test_cisco_ai_defense_chat.py | 1 - .../test_litellm/proxy/guardrails/test_pillar_guardrails.py | 1 - tests/test_litellm/responses/test_text_format_conversion.py | 1 - 12 files changed, 1 insertion(+), 24 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 5039143eaa0..38fd31d7275 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -12,7 +12,7 @@ "limit": 469 }, "TQ005": { - "limit": 2459 + "limit": 2436 }, "TQ006": { "limit": 34 diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 99826c14069..3cc251b6228 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -812,7 +812,6 @@ async def test_azure_client_reuse(function_name, is_async, args): """ Test that multiple Azure API calls reuse the same Azure OpenAI client """ - litellm.set_verbose = True # Determine which client class to mock based on whether the test is async client_path = ( diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 8b6034d1133..1c802ecd077 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -153,7 +153,6 @@ class TestBedrockAsyncInvokeEmbedding: def test_async_invoke_twelvelabs_embedding_with_mock(self): """Test async invoke embedding with mocked HTTP calls.""" - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0" @@ -193,7 +192,6 @@ class TestBedrockAsyncInvokeEmbedding: @pytest.mark.asyncio async def test_async_invoke_twelvelabs_embedding_async_with_mock(self): """Test async invoke embedding with async calls.""" - litellm.set_verbose = True client = AsyncHTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/async_invoke/twelvelabs.marengo-embed-2-7-v1:0" diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py index 41ac030ff07..b2b00d25051 100644 --- a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py +++ b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_bearer_token.py @@ -18,7 +18,6 @@ mock_image_response = {"images": ["base64_encoded_image_data"], "error": None} class TestBedrockImageGeneration: def test_image_generation_with_api_key_bearer_token(self): """Test image generation with bearer token authentication""" - litellm.set_verbose = True test_api_key = "test-bearer-token-12345" model = "bedrock/stability.sd3-large-v1:0" prompt = "A cute baby sea otter" @@ -53,7 +52,6 @@ class TestBedrockImageGeneration: def test_image_generation_with_env_variable_bearer_token(self, monkeypatch): """Test image generation with bearer token from environment variable""" - litellm.set_verbose = True test_api_key = "env-bearer-token-12345" model = "bedrock/stability.sd3-large-v1:0" prompt = "A cute baby sea otter" @@ -90,7 +88,6 @@ class TestBedrockImageGeneration: @pytest.mark.asyncio async def test_async_image_generation_with_bearer_token(self): """Test async image generation with bearer token authentication""" - litellm.set_verbose = True test_api_key = "async-bearer-token-12345" model = "bedrock/stability.sd3-large-v1:0" prompt = "A cute baby sea otter" @@ -125,7 +122,6 @@ class TestBedrockImageGeneration: def test_image_generation_with_sigv4(self): """Test image generation falls back to SigV4 auth when no bearer token is provided""" - litellm.set_verbose = True model = "bedrock/stability.sd3-large-v1:0" prompt = "A cute baby sea otter" diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 17443ca899e..d8259652641 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -66,7 +66,6 @@ def test_bedrock_rerank_header_forwarding_sync(model): This test verifies the fix for the issue where headers configured via forward_client_headers_to_llm_api were not being passed to Bedrock rerank provider. """ - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" @@ -160,7 +159,6 @@ async def test_bedrock_rerank_header_forwarding_async(model): This test verifies the fix for the issue where headers configured via forward_client_headers_to_llm_api were not being passed to Bedrock rerank provider. """ - litellm.set_verbose = True client = AsyncHTTPHandler() test_api_key = "test-bearer-token-12345" @@ -332,7 +330,6 @@ def test_bedrock_rerank_extra_headers_and_headers_merge(): This ensures that headers from kwargs (forwarded by proxy) and extra_headers (passed explicitly) are both included in the final headers sent to the provider. """ - litellm.set_verbose = True client = HTTPHandler() test_api_key = "test-bearer-token-12345" model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0" diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index bfd681cc06e..bef8d02b0df 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -86,7 +86,6 @@ async def test_openai_client_reuse(function_name, is_async, args): """ Test that multiple API calls reuse the same OpenAI client """ - litellm.set_verbose = True # Determine which client class to mock based on whether the test is async client_path = ( diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index ec73e5e42be..ae260a2d887 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -33,7 +33,6 @@ def test_completion_pydantic_obj_2(): from litellm.llms.custom_httpx.http_handler import HTTPHandler - litellm.set_verbose = True class CalendarEvent(BaseModel): name: str @@ -259,7 +258,6 @@ def test_vertex_tool_type_field_removal(): def test_function_calling_with_gemini(): from litellm.llms.custom_httpx.http_handler import HTTPHandler - litellm.set_verbose = True client = HTTPHandler() with patch.object(client, "post", new=MagicMock()) as mock_post: try: @@ -310,7 +308,6 @@ def test_function_calling_with_gemini(): def test_multiple_function_call(): - litellm.set_verbose = True from litellm.llms.custom_httpx.http_handler import HTTPHandler client = HTTPHandler() @@ -420,7 +417,6 @@ def test_multiple_function_call(): def test_multiple_function_call_changed_text_pos(): - litellm.set_verbose = True from litellm.llms.custom_httpx.http_handler import HTTPHandler client = HTTPHandler() @@ -528,7 +524,6 @@ def test_multiple_function_call_changed_text_pos(): def test_function_calling_with_gemini_multiple_results(): - litellm.set_verbose = True from litellm.llms.custom_httpx.http_handler import HTTPHandler client = HTTPHandler() @@ -1103,7 +1098,6 @@ def test_logprobs_unit_test(): def test_logprobs(): - litellm.set_verbose = True from litellm.llms.custom_httpx.http_handler import HTTPHandler client = HTTPHandler() diff --git a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py index 1670dac0e9d..04caecab478 100644 --- a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py +++ b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py @@ -31,7 +31,6 @@ class TestVolcEngineEmbedding(BaseLLMEmbeddingTest): @pytest.mark.parametrize("sync_mode", [True, False]) async def test_basic_embedding(self, sync_mode): """Test basic embedding functionality with realistic response""" - litellm.set_verbose = True embedding_call_args = self.get_base_embedding_call_args() # Mock the embedding functions to avoid actual API calls diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py index c23fbc0234e..caed64ef417 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -26,7 +26,6 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 def test_cato_guard_config(): - litellm.set_verbose = True litellm.guardrail_name_config_map = {} init_guardrails_v2( @@ -47,7 +46,6 @@ def test_cato_guard_config(): def test_cato_guard_config_no_api_key(monkeypatch): monkeypatch.delenv("CATO_API_KEY", raising=False) - litellm.set_verbose = True litellm.guardrail_name_config_map = {} with pytest.raises(CatoNetworksGuardrailMissingSecrets, match="Couldn't get Cato Networks api key"): init_guardrails_v2( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py index 8974a18593b..779075a40d9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py @@ -44,7 +44,6 @@ from tests.test_litellm.proxy.guardrails.guardrail_hooks._cisco_ai_defense_test_ def test_cisco_ai_defense_config_via_init_v2_chat(monkeypatch): monkeypatch.setenv("CISCO_AI_DEFENSE_API_KEY", "test-key") - litellm.set_verbose = True litellm.guardrail_name_config_map = {} init_guardrails_v2( diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index 48f6b3ba2b9..02123bc8c76 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -65,7 +65,6 @@ def setup_and_teardown(): asyncio.set_event_loop(loop) # Set up litellm state - litellm.set_verbose = True litellm.guardrail_name_config_map = {} yield diff --git a/tests/test_litellm/responses/test_text_format_conversion.py b/tests/test_litellm/responses/test_text_format_conversion.py index a48540b129b..339b73c2729 100644 --- a/tests/test_litellm/responses/test_text_format_conversion.py +++ b/tests/test_litellm/responses/test_text_format_conversion.py @@ -158,7 +158,6 @@ class TestTextFormatConversion: new=mock_handler, ): litellm._turn_on_debug() - litellm.set_verbose = True # Call aresponses with text_format parameter response = await litellm.aresponses( From 5285ae86d5005871e2c600f3adcbca65eaff45b9 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 22 Aug 2026 08:25:16 -0700 Subject: [PATCH 43/43] fix(ptu): warn when config.yaml declares PTU while attribution is off (#37898) --- litellm/litellm_core_utils/ptu_pricing.py | 8 ++ .../model_management_endpoints.py | 6 +- litellm/router.py | 17 +++ .../test_router_model_cost_isolation.py | 125 ++++++++++++++++++ 4 files changed, 153 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py index 2e73719cf52..021210d9175 100644 --- a/litellm/litellm_core_utils/ptu_pricing.py +++ b/litellm/litellm_core_utils/ptu_pricing.py @@ -124,6 +124,14 @@ def ptu_identity_error( return None +PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") + + +def declares_ptu(model_info: Mapping[str, object]) -> bool: + """Whether any PTU field is set here, including one too malformed to charge.""" + return any(model_info.get(field) is not None for field in PTU_MODEL_INFO_FIELDS) + + def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None = None) -> str | None: """Why this PTU configuration cannot be honoured, else None. diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b003daa9d79..217fc61a56c 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -27,6 +27,7 @@ from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.litellm_core_utils.ptu_pricing import ( CUSTOM_PRICING_FIELDS, PTU_EMPTIED_PRICING_FIELDS, + PTU_MODEL_INFO_FIELDS, PTU_ZEROED_PRICING_FIELDS, PTU_ZEROED_TABLE_FIELDS, SEARCH_CONTEXT_SIZES, @@ -247,7 +248,6 @@ def _raise_on_strategy_router_write_violation( ) -_PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") _PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"}) @@ -261,7 +261,7 @@ def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[st return frozenset() return frozenset( field - for field in _PTU_MODEL_INFO_FIELDS + for field in PTU_MODEL_INFO_FIELDS if field in model_info.model_fields_set and getattr(model_info, field) is None ) @@ -294,7 +294,7 @@ def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, ob """ if is_ptu_cost_attribution_enabled(): return - supplied: Final = tuple(field for field in _PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None) + supplied: Final = tuple(field for field in PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None) if not supplied: return raise HTTPException( diff --git a/litellm/router.py b/litellm/router.py index 7dedbe851d7..045fd32847c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -66,6 +66,8 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.ptu_pricing import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + declares_ptu, is_ptu_cost_attribution_enabled, ptu_config_error, ptu_identity_error, @@ -8234,6 +8236,21 @@ class Router: ) duplicate_ids: Final = frozenset(model_id for model_id in declared_ids if declared_ids.count(model_id) > 1) + ptu_declared: Final = tuple( + str(entry.get("model_name")) + for entry in original_model_list + if isinstance(entry.get("model_info"), dict) + and entry["model_info"].get("db_model") is not True + and declares_ptu(entry["model_info"]) + ) + if ptu_declared and not is_ptu_cost_attribution_enabled(): + verbose_router_logger.warning( + "PTU fields are set on config.yaml deployment(s) %s, but PTU cost attribution is disabled, so no " + "flat cost accrues and this traffic is billed per token. Set %s=True to enable it", + ", ".join(ptu_declared), + PTU_COST_ATTRIBUTION_ENV_VAR, + ) + for model in original_model_list: _model_name = model.pop("model_name") _litellm_params = model.pop("litellm_params") diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 1580ec7f437..eb454bedbd8 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -8,6 +8,7 @@ should still use the built-in pricing. """ import copy +import logging import os import re import sys @@ -2007,3 +2008,127 @@ def test_a_falsy_id_is_still_scanned_for_collisions(): }, ] ) + + +# --- a reservation declared while the feature is off says so ------------------------ + + +def _ptu_warnings(caplog): + return tuple( + record.getMessage() + for record in caplog.records + if record.name == "LiteLLM Router" and record.levelno == logging.WARNING and "PTU" in record.getMessage() + ) + + +def test_a_reservation_declared_while_the_feature_is_off_is_warned_about(caplog): + """The deployment serves and bills per token, so without this the operator believes they + reserved capacity and sees no signal anywhere that nothing accrues.""" + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + _ptu_router(ptu_enabled=False) + + warnings = _ptu_warnings(caplog) + + assert len(warnings) == 1 + assert "gpt-4o-ptu" in warnings[0] + assert "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" in warnings[0] + + +def test_a_reservation_is_not_warned_about_while_the_feature_is_on(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + _ptu_router() + + assert _ptu_warnings(caplog) == () + + +def test_a_deployment_carrying_no_ptu_field_is_not_warned_about(caplog): + """Most of every config.yaml, so warning here would fire on proxies that never asked.""" + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + _ptu_router(model_info={"team_id": "team-alpha"}, ptu_enabled=False) + + assert _ptu_warnings(caplog) == () + + +def test_a_half_written_reservation_is_warned_about(caplog): + """A count with no rate is not a chargeable reservation, but the operator still meant to + declare one, so what they wrote is what decides whether they hear about it.""" + half_written = {k: v for k, v in _PTU_MODEL_INFO.items() if k != "cost_per_ptu_per_hour"} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + _ptu_router(model_info=half_written, ptu_enabled=False) + + assert len(_ptu_warnings(caplog)) == 1 + + +@pytest.mark.parametrize( + "typo", + [ + {"ptu_count": 0}, + {"ptu_count": 0, "cost_per_ptu_per_hour": 0, "ptu_effective_from": None}, + ], + ids=["count out of range", "every value still a zero placeholder"], +) +def test_a_reservation_dropped_by_a_typo_is_warned_about(caplog, typo): + """An out-of-range value fails ModelInfo before the flag is ever consulted, so the + deployment stops serving on a proxy that never enabled PTU. The warning is what tells the + operator which feature the entry that vanished belonged to. + + Built the way proxy_server builds it, since dropping rather than raising is what + ``ignore_invalid_deployments`` does and config.yaml is loaded with it on. + """ + with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": ""}, clear=False): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + router = Router( + ignore_invalid_deployments=True, + model_list=[ + { + "model_name": "gpt-4o-ptu", + "litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://e.azure.com"}, + "model_info": {**_PTU_MODEL_INFO, **typo}, + } + ], + ) + + assert router.model_list == [] + assert len(_ptu_warnings(caplog)) == 1 + + +def test_a_db_backed_reservation_is_not_warned_about(caplog): + """/model/new already answered the caller with a 400, so repeating it on every reload + would report the operator's own rejected write back to them as a standing problem.""" + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + _ptu_router(model_info={**_PTU_MODEL_INFO, "db_model": True}, ptu_enabled=False) + + assert _ptu_warnings(caplog) == () + + +def test_every_declaring_deployment_is_named(caplog): + """One line naming all of them, so a reload does not bury the config in repeats.""" + with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": ""}, clear=False): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + Router( + model_list=[ + { + "model_name": "azure-ptu-east", + "litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://e.azure.com"}, + "model_info": dict(_PTU_MODEL_INFO), + }, + { + "model_name": "azure-ptu-west", + "litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://w.azure.com"}, + "model_info": {**_PTU_MODEL_INFO, "id": "ptu-alpha-westus"}, + }, + { + "model_name": "plain-gpt-4o", + "litellm_params": {"model": "azure/gpt-4o", "api_key": "k", "api_base": "https://p.azure.com"}, + "model_info": {"id": "plain"}, + }, + ] + ) + + warnings = _ptu_warnings(caplog) + + assert len(warnings) == 1 + assert "azure-ptu-east" in warnings[0] + assert "azure-ptu-west" in warnings[0] + assert "plain-gpt-4o" not in warnings[0]