From b93030f84e7a414d2106528114b09f1fca1ad1aa Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:20:25 +0000 Subject: [PATCH 01/74] fix(vertex_ai): surface real error/status on vertex batch create instead of IndexError 500 --- litellm/llms/vertex_ai/batches/handler.py | 44 ++++++++++++---- .../llms/vertex_ai/batches/transformation.py | 46 ++++++++++++++--- .../llms/vertex_ai/batches/test_handler.py | 50 +++++++++++++++---- .../vertex_ai/batches/test_transformation.py | 43 +++++++++++++++- 4 files changed, 152 insertions(+), 31 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ada1356fb6b..f0fd5480c75 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -13,7 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -98,7 +98,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -130,7 +132,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -242,7 +246,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -292,7 +298,9 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -365,7 +373,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -390,7 +400,9 @@ class VertexAIBatchPrediction(VertexLLM): params=params, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -475,7 +487,9 @@ class VertexAIBatchPrediction(VertexLLM): raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # HTTPHandler.get() does not accept a timeout parameter retrieve_response = sync_handler.get( @@ -488,7 +502,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -521,7 +538,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response = await client.get( @@ -534,7 +553,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index df903ba7ef0..e4299bcf2a0 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,9 @@ from typing import Any, Dict, Optional +from urllib.parse import unquote from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest @@ -199,16 +201,40 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ - from urllib.parse import unquote - - decoded_uri = unquote(gcs_file_uri) - - model_path = decoded_uri.split("publishers/")[1] - parts = model_path.split("/") - model = f"publishers/{'/'.join(parts[:3])}" + model = cls._parse_model_from_gcs_file(gcs_file_uri) + if model is None: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch creation requires the model to be part of `input_file_id`, but " + f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + "Either upload the input file through LiteLLM (POST /v1/files with " + "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " + "pass a uri of the form " + "gs:////publishers//models//" + ), + ) return model + @classmethod + def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: + """ + Returns the `publishers//models/` path from a gcs uri, or None if the uri + does not contain one. + """ + _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") + if not separator: + return None + + parts = model_path.split("/") + if len(parts) < 3 or parts[1] != "models" or not parts[2]: + return None + + return f"publishers/{'/'.join(parts[:3])}" + @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool: """ @@ -216,7 +242,11 @@ class VertexAIBatchTransformation: LiteLLM-managed unified file id) with a `publishers/` model path that `_get_model_from_gcs_file` can parse. """ - return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id + return ( + input_file_id is not None + and input_file_id.startswith("gs://") + and cls._parse_model_from_gcs_file(input_file_id) is not None + ) @classmethod def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index cacea234777..b9fb5dfe3c5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -40,6 +40,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) +from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" @@ -184,7 +185,7 @@ def test_create_batch_sync_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500") as exc_info: h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -196,6 +197,32 @@ def test_create_batch_sync_non_200_raises(): max_retries=None, ) + assert exc_info.value.status_code == 500 + assert "error text" in str(exc_info.value) + + +def test_create_batch_input_file_id_without_model_raises_400_before_post(): + """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data={"input_file_id": "gs://bucket/batch-input.jsonl"}, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "gs://bucket/batch-input.jsonl" in str(exc_info.value) + client.post.assert_not_called() + def test_create_batch_async_non_200_raises(): h = _make_handler() @@ -216,9 +243,12 @@ def test_create_batch_async_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 403"): + with pytest.raises(VertexAIError, match="Error: 403") as exc_info: _run(coro) + assert exc_info.value.status_code == 403 + assert "error text" in str(exc_info.value) + # =========================================================================== # # retrieve_batch @@ -292,7 +322,7 @@ def test_retrieve_batch_sync_non_200_raises(): patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), patch(f"{HMOD}.safe_get", return_value=_http_response(status_code=404)), ): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.retrieve_batch( _is_async=False, batch_id=BATCH_ID, @@ -438,7 +468,7 @@ def test_list_batches_sync_non_200_raises(): client.get.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.list_batches( _is_async=False, after=None, @@ -530,7 +560,7 @@ def test_cancel_batch_sync_cancel_post_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -552,7 +582,7 @@ def test_cancel_batch_sync_retrieve_non_200_raises(): client.get.return_value = _http_response(status_code=404) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -672,7 +702,7 @@ def test_async_retrieve_batch_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -726,7 +756,7 @@ def test_async_list_batches_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -779,7 +809,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) async_client_post500.get.assert_not_awaited() @@ -801,5 +831,5 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): _run(coro) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 1b37ade6b30..8352ec16389 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -25,6 +25,7 @@ from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -69,6 +70,24 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +@pytest.mark.parametrize( + "input_file_id", + [ + "gs://bucket/no-model-here.jsonl", + "gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", + "gs://bucket/publishers/google/models", + "gs://bucket/publishers/google/models//file-uuid", + ], +) +def test_transform_openai_request_unparseable_model_raises_400(input_file_id: str): + """An input_file_id with no parseable model path is a client error, not an IndexError -> 500.""" + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": input_file_id}) + + assert exc_info.value.status_code == 400 + assert input_file_id in str(exc_info.value) + + # =========================================================================== # # transform_vertex_ai_batch_response_to_openai_batch_response # =========================================================================== # @@ -299,9 +318,29 @@ def test_get_model_from_gcs_file_url_encoded(): assert T._get_model_from_gcs_file(encoded) == "publishers/google/models/gemini-1.5-flash-001" -def test_get_model_from_gcs_file_no_publishers_raises(): - with pytest.raises(IndexError): +def test_get_model_from_gcs_file_no_publishers_raises_400(): + with pytest.raises(VertexAIError) as exc_info: T._get_model_from_gcs_file("gs://bucket/no-model-here.jsonl") + assert exc_info.value.status_code == 400 + + +# =========================================================================== # +# is_unmanaged_gcs_batch_input_file_id +# =========================================================================== # + + +@pytest.mark.parametrize( + "input_file_id, expected", + [ + (INPUT_FILE, True), + (None, False), + ("file-abc123", False), + ("gs://bucket/no-model-here.jsonl", False), + ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + ], +) +def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): + assert T.is_unmanaged_gcs_batch_input_file_id(input_file_id) is expected # =========================================================================== # From 18d9c7aa21e1308c5ecf05254b29bc0715965bde Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:09:49 +0000 Subject: [PATCH 02/74] fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload --- .../llms/bedrock/batches/transformation.py | 8 ++- litellm/llms/bedrock/common_utils.py | 19 ++++- litellm/llms/bedrock/files/transformation.py | 16 ++++- litellm/types/router.py | 1 + .../bedrock/batches/test_transformation.py | 2 +- .../test_bedrock_files_transformation.py | 72 ++++++++++++++++++- tests/test_litellm/test_router.py | 31 ++++++++ 7 files changed, 141 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index a4ff1c78467..7500531b81a 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -12,7 +12,6 @@ from litellm.litellm_core_utils.cloud_storage_security import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, @@ -29,7 +28,7 @@ from litellm.types.llms.openai import ( from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM -from ..common_utils import CommonBatchFilesUtils +from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see @@ -200,7 +199,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Add optional KMS encryption key ID if provided - s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + s3_encryption_key_id = resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ) if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5114677ffc0..9d427fa6f12 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -35,7 +35,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret +from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -1313,6 +1313,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: return [] +def resolve_s3_encryption_key_id( + litellm_params: Mapping[str, Any], + optional_params: Mapping[str, Any] | None = None, +) -> str | None: + """ + Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects. + + Precedence: `s3_encryption_key_id` in litellm_params, then optional_params + (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. + """ + for source in (litellm_params, optional_params or {}): + value = source.get("s3_encryption_key_id") + if isinstance(value, str) and value: + return value + return get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + + class CommonBatchFilesUtils: """ Common utilities for Bedrock batch and file operations. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index d4865a1c87a..d1674b260b4 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -53,7 +53,7 @@ from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError +from ..common_utils import BedrockError, resolve_s3_encryption_key_id # litellm_params key used to hand the SigV4-signed GET headers from # `transform_file_content_request` to `validate_environment` (the only hook @@ -741,6 +741,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content=file_content, api_base=api_base, optional_params=optional_params, + s3_encryption_key_id=resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ), ) litellm_params["upload_url"] = api_base @@ -758,6 +762,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content: str, api_base: str, optional_params: dict, + s3_encryption_key_id: str | None = None, ) -> Tuple[dict, str]: """ Sign S3 PUT request using the same proven logic as S3Logger. @@ -790,11 +795,20 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() # Prepare headers with required S3 headers (same as s3_v2.py) + sse_headers = ( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + if s3_encryption_key_id + else {} + ) request_headers = { "Content-Type": "application/json", # JSONL files are JSON content "x-amz-content-sha256": content_hash, # REQUIRED by S3 "Content-Language": "en", "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, } # Use requests.Request to prepare the request (same pattern as s3_v2.py) diff --git a/litellm/types/router.py b/litellm/types/router.py index 28e4a8272e8..c4d679a2500 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -211,6 +211,7 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None s3_bucket_name: Optional[str] = None + s3_encryption_key_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 3681daffe5e..01420eb10df 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -172,7 +172,7 @@ def test_create_request_omits_kms_key_when_absent(config): "generate_unique_job_name", return_value="litellm-batch-1", ), patch.object(config.common_utils, "sign_aws_request") as mock_sign, patch( - "litellm.llms.bedrock.batches.transformation.get_secret_str", + "litellm.llms.bedrock.common_utils.get_secret_str", return_value=None, ): mock_sign.return_value = ({}, b"{}") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..a57e5801327 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -442,7 +442,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -498,7 +498,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -514,6 +514,74 @@ class TestBedrockFilesTransformation: captured_optional_params.get("aws_region_name") == "us-gov-west-1" ), "s3_region_name must override aws_region_name for SigV4 signing" + def _signed_upload_request(self, litellm_params: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + jsonl_content = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } + ).encode() + + request = config.transform_create_file_request( + model="amazon.nova-pro-v1:0", + create_file_data={ + "file": ("batch.jsonl", jsonl_content, "application/jsonl"), + "purpose": "batch", + }, + optional_params={ + "aws_access_key_id": "test-key-id", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + litellm_params={"s3_bucket_name": "litellm-batch-bucket", **litellm_params}, + ) + assert isinstance(request, dict) + return request + + def test_upload_signs_sse_kms_headers_when_key_configured(self, monkeypatch): + """ + Buckets whose policy requires SSE-KMS reject the batch input-file PutObject + unless the upload carries the aws:kms encryption headers; they must also be + covered by SigV4 SignedHeaders or S3 answers SignatureDoesNotMatch. + """ + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + kms_key = "arn:aws:kms:us-west-2:1234:key/abcd" + + request = self._signed_upload_request({"s3_encryption_key_id": kms_key}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == kms_key + signed_headers = headers["authorization"].split("SignedHeaders=")[1].split(",")[0] + assert "x-amz-server-side-encryption" in signed_headers + assert "x-amz-server-side-encryption-aws-kms-key-id" in signed_headers + + def test_upload_reads_sse_kms_key_from_env(self, monkeypatch): + monkeypatch.setenv("AWS_S3_ENCRYPTION_KEY_ID", "env-kms-key") + + request = self._signed_upload_request({}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == "env-kms-key" + + def test_upload_omits_sse_headers_when_no_key_configured(self, monkeypatch): + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + request = self._signed_upload_request({}) + + headers = {key.lower() for key in request["headers"]} + assert "x-amz-server-side-encryption" not in headers + assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + def test_openai_passthrough_still_works(self): """ Regression test: ensure OpenAI-compatible models (e.g. gpt-oss) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..fa047d7ee46 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3666,6 +3666,37 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" +def test_get_deployment_credentials_with_provider_includes_s3_encryption_key_id(): + """ + Regression: s3_encryption_key_id must survive the CredentialLiteLLMParams filter, + otherwise the Bedrock batch input-file upload loses the SSE-KMS key and S3 rejects + the PutObject on buckets whose policy requires aws:kms encryption. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-batch-bucket", + "s3_encryption_key_id": "arn:aws:kms:us-west-2:1234:key/abcd", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch" + ) + + assert credentials is not None + assert ( + credentials["s3_encryption_key_id"] + == "arn:aws:kms:us-west-2:1234:key/abcd" + ) + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves From 0b809cf7d68e368b7a7ee90d7134c4841c944e3c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:21:20 +0000 Subject: [PATCH 03/74] fix(anthropic adapter): stop indexing choices[0] on choiceless streaming chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../adapters/streaming_iterator.py | 30 +++++++ .../test_streaming_iterator_empty_choices.py | 87 +++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index d9bcfa19a7f..194cbcc9327 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -329,6 +329,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return self._augment_message_delta_usage(merged_chunk) + def _handle_choiceless_chunk(self, chunk: Any) -> bool: + """Consume an OpenAI-compatible chunk that carries no ``choices``. + + ``choices`` is legitimately empty on metadata-only chunks; the final + usage chunk emitted when ``stream_options.include_usage`` is set is the + common case (vLLM and other OpenAI-compatible servers do this). Such a + chunk carries no content-block information, so the caller must not run + the content-block state machine over it. + + Returns True when a merged ``message_delta`` was queued (usage folded + into the held stop-reason chunk); False when the chunk should be + skipped entirely. + """ + if self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None: + self.chunk_queue.append(self._merge_usage_into_held_stop_reason_chunk(chunk)) + self.queued_usage_chunk = True + self.holding_stop_reason_chunk = None + return True + return False + def _ensure_context_management_attached(self, message_delta_chunk: Dict[str, Any]) -> Dict[str, Any]: """Attach ``context_management`` to a ``message_delta`` chunk if ``self.applied_edits`` is non-empty and the chunk does not already @@ -490,6 +510,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): @@ -713,6 +738,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py new file mode 100644 index 00000000000..3e85872f1e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py @@ -0,0 +1,87 @@ +""" +Regression tests for OpenAI-compatible chunks with an empty ``choices`` list. + +``choices: []`` is valid OpenAI-compatible streaming: vLLM (and OpenAI itself, +when ``stream_options.include_usage`` is set) emits a final usage chunk with no +choices, and some gateways emit metadata-only chunks mid-stream. The adapter +used to index ``chunk.choices[0]`` unconditionally, so such a chunk raised +``IndexError: list index out of range`` and killed the ``/v1/messages`` stream. +""" + +import asyncio +import json +from typing import Any, AsyncIterator, Dict, List, Optional + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + +def _text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)] + ) + + +def _finish_chunk() -> ModelResponseStream: + return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]) + + +def _empty_choices_chunk(usage: Optional[Usage] = None) -> ModelResponseStream: + return ModelResponseStream(choices=[], usage=usage) + + +def _collect_async(wrapper: AnthropicStreamWrapper) -> str: + async def _run() -> str: + return "".join( + [raw.decode() if isinstance(raw, bytes) else raw async for raw in wrapper.async_anthropic_sse_wrapper()] + ) + + return asyncio.run(_run()) + + +def _message_delta(sse: str) -> Dict[str, Any]: + return next( + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"message_delta"' in line + ) + + +def test_leading_metadata_chunk_without_choices_does_not_kill_stream(): + """A metadata-only chunk before any content must be skipped, not indexed.""" + chunks: List[ModelResponseStream] = [ + _empty_choices_chunk(), + _text_chunk("Hello"), + _text_chunk(" there"), + _finish_chunk(), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="mock-model") + events = list(wrapper) + + text = "".join( + event["delta"]["text"] for event in events if event.get("type") == "content_block_delta" + ) + assert text == "Hello there" + assert events[-1]["type"] == "message_stop" + + +def test_final_usage_chunk_without_choices_is_merged_into_message_delta(): + """The vLLM/OpenAI final usage chunk carries no choices; its usage must + still land on the Anthropic ``message_delta``.""" + usage = Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13) + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in [_text_chunk("Hi"), _finish_chunk(), _empty_choices_chunk(usage)]: + yield chunk + + sse = _collect_async(AnthropicStreamWrapper(completion_stream=_aiter(), model="mock-model")) + + message_delta = _message_delta(sse) + assert message_delta["delta"]["stop_reason"] == "end_turn" + assert message_delta["usage"]["input_tokens"] == 10 + assert message_delta["usage"]["output_tokens"] == 3 + assert "Hi" in sse + assert "message_stop" in sse From c5c5a276790529e2de3378654864fd847530c5a0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:19:42 +0000 Subject: [PATCH 04/74] fix(files): enforce require_managed_files on file retrieve, content and delete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai_files_endpoints/common_utils.py | 30 +++++ .../openai_files_endpoints/files_endpoints.py | 7 ++ .../test_files_endpoint.py | 116 ++++++++++++++++++ 3 files changed, 153 insertions(+) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 87514b46dbd..3eef3868c94 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -866,6 +866,36 @@ def validate_managed_files_requirement( ) +def validate_managed_file_id_requirement(file_id: str) -> None: + """ + Enforce proxy-level managed files on the file read/delete routes when + ``litellm.require_managed_files`` is enabled. + + Ownership is only recorded for LiteLLM managed files, so a raw provider file id sent to + retrieve/content/delete is forwarded to the provider under shared credentials without any + tenant check; knowing another tenant's provider file id would be enough to read or delete it. + + Raises: + HTTPException: 400 if ``file_id`` is not a LiteLLM managed file id. + """ + import litellm + from fastapi import HTTPException + + if litellm.require_managed_files is not True: + return + + if _is_base64_encoded_unified_file_id(file_id): + return + + raise HTTPException( + status_code=400, + detail=( + "Raw provider file ids cannot be used when require_managed_files is enabled in " + "litellm_settings. Use the LiteLLM managed file id returned when the file was created." + ), + ) + + def _extract_model_param(request: "Request", request_body: dict) -> str | None: """ Extract model parameter from request. diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 4e4718272bd..37f1ced6996 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, + validate_managed_file_id_requirement, validate_managed_files_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model @@ -612,6 +613,8 @@ async def get_file_content( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -908,6 +911,8 @@ async def get_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -1098,6 +1103,8 @@ async def delete_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ac01c6ae1d1..24b814bae1f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3051,3 +3051,119 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials( mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"] ) assert captured_kwargs.get("api_key") == "team-openai-key" + + +@pytest.mark.parametrize( + "http_method, url, patched_litellm_call", + [ + ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"), + ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"), + ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"), + ], +) +def test_require_managed_files_rejects_raw_provider_file_id( + mocker: MockerFixture, + monkeypatch, + llm_router: Router, + http_method: str, + url: str, + patched_litellm_call: str, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user" + ) + + try: + response = getattr(client, http_method)( + url, headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + monkeypatch.setattr("litellm.require_managed_files", False) + + assert response.status_code == 400, response.text + mock_call.assert_not_called() + + +def _unified_managed_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + +def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", True) + + validate_managed_file_id_requirement(file_id=_unified_managed_file_id()) + + +def test_managed_file_id_requirement_is_opt_in(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", False) + + validate_managed_file_id_requirement(file_id="file-victim-abc123") + + +def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", False) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_retrieve = mocker.patch( + "litellm.afile_retrieve", + new=mocker.AsyncMock( + return_value=OpenAIFileObject( + id="file-victim-abc123", + object="file", + bytes=3, + created_at=1234567890, + filename="test.txt", + purpose="user_data", + status="uploaded", + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user" + ) + + try: + response = client.get( + "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + mock_retrieve.assert_called_once() From a7250f4eeab560299215773b50ba32405f597a1c Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 17:11:50 +0000 Subject: [PATCH 05/74] fix(bedrock): normalize /v1/completions and /v1/responses batch records Bedrock managed-batch file upload read `messages` unconditionally, so a JSONL record shaped for /v1/completions (`prompt`) or /v1/responses (`input`) reached the per-provider transform with an empty message list. Anthropic and Nova rejected it at POST /v1/files, and the passthrough providers shipped an empty conversation to AWS. Classify each record by its OpenAI batch `url`, then normalize the non-embedding shapes to chat completions before the Bedrock transforms: `prompt` wraps into user messages the way litellm.text_completion does in real time, and `input` goes through the existing Responses-to-Chat bridge. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 22 +- litellm/llms/bedrock/files/transformation.py | 194 ++++++++-- litellm/types/llms/bedrock.py | 15 + ...ore_utils_prompt_templates_common_utils.py | 43 +++ .../test_bedrock_files_transformation.py | 347 ++++++++++++++++-- 5 files changed, 559 insertions(+), 62 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 639c93dfb80..777ba398d5a 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import io import json import mimetypes import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from os import PathLike from pathlib import Path from typing import ( @@ -1742,3 +1742,23 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: idx = end_idx return results + + +def text_completion_prompt_to_messages(prompt: str | Sequence[str]) -> tuple[AllMessageValues, ...]: + """ + Wrap an OpenAI ``/v1/completions`` ``prompt`` into Chat Completion messages. + + Mirrors what ``litellm.text_completion`` does on the real-time path: a + string becomes a single user message, and a list of strings becomes one + user message per element. Pre-tokenized prompts (``list[int]`` / + ``list[list[int]]``) are only meaningful for the OpenAI-family text + endpoints, so they are rejected here rather than silently forwarded, as is + an empty prompt, which every chat-shaped provider rejects downstream. + """ + if isinstance(prompt, str) and prompt: + return (ChatCompletionUserMessage(role="user", content=prompt),) + if isinstance(prompt, Sequence) and prompt and all(isinstance(entry, str) and entry for entry in prompt): + return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in prompt) + raise ValueError( + f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {type(prompt).__name__}." + ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 3656088cb9d..0aa832780e5 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -2,7 +2,9 @@ import base64 import json import os import time -from collections.abc import Mapping, MutableMapping +from collections.abc import Iterable, Mapping, MutableMapping +from functools import cache +from itertools import chain from types import MappingProxyType from typing import ( Any, @@ -12,7 +14,7 @@ from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -28,12 +30,16 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + text_completion_prompt_to_messages, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.types.llms.bedrock import BedrockBatchRecordKind from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -43,6 +49,8 @@ from litellm.types.llms.openai import ( OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, + ResponseInputParam, + ResponsesAPIOptionalRequestParams, ) from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider @@ -57,6 +65,26 @@ from ..common_utils import BedrockError S3_SIGNED_GET_HEADERS_PARAM = "_s3_signed_get_headers" +def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]: + return MappingProxyType(dict(items)) + + +# JSONL batch records are untyped json, so the `/v1/responses` fields are +# validated into their concrete Responses API types before being handed to the +# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't +# define, which is what the bridge would ignore anyway. Built on first use +# rather than at import: `ResponseInputParam` is a deep union and only batch +# files carrying `/v1/responses` records need it. +@cache +def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]: + return TypeAdapter(str | ResponseInputParam) + + +@cache +def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]: + return TypeAdapter(ResponsesAPIOptionalRequestParams) + + class _BedrockS3RequestParams(BaseModel): """Typed view of the credential/region params the S3 GetObject path reads.""" @@ -305,41 +333,55 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # example; add others here as they adopt the same schema. CONVERSE_INVOKE_PROVIDERS = ("nova",) - # OpenAI batch URL that signals an embedding request. Per OpenAI Batch API - # spec, every JSONL record carries a `url` field; we use it as the - # authoritative signal to route the line to the embedding code path - # instead of inferring from the presence of `input` vs `messages`. + # OpenAI batch URLs that select which request shape a JSONL line carries. + # Per the OpenAI Batch API spec every record carries a `url`, so we use it + # as the authoritative routing signal instead of inferring from the + # presence of `input` vs `prompt` vs `messages`. OPENAI_EMBEDDINGS_URL = "/v1/embeddings" + OPENAI_TEXT_COMPLETIONS_URL = "/v1/completions" + OPENAI_RESPONSES_URL = "/v1/responses" @staticmethod - def _is_embedding_record(openai_jsonl_record: dict[str, Any]) -> bool: + def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind: """ - Decide whether an OpenAI batch JSONL line is an embedding request. + Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries. - Precedence (strict - any explicit `url` short-circuits): - 1. `url == "/v1/embeddings"` -> embedding. Authoritative per the - OpenAI Batch API spec. - 2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT - embedding. We trust the caller's explicit signal even if the - body would otherwise suggest embedding; misrouting a chat - record into the embedding transformer would corrupt the - modelInput, while a chat-shaped body sent to the chat path - either succeeds or fails cleanly inside that transformer. - 3. `url` missing/empty -> fall back to body shape. Requires - `input` present AND `messages` absent so a malformed record - carrying both keys routes to the chat path (safer default: - Anthropic transforms ignore unknown top-level keys, whereas - the embedding transformer would silently drop the messages). + Precedence (strict - any recognized `url` short-circuits): + 1. A `url` matching a supported endpoint wins. Authoritative per the + OpenAI Batch API spec, which requires it on every record. + 2. Any other non-empty `url` -> chat. We trust the caller's explicit + signal rather than re-deriving it from the body, and an + unexpectedly-shaped body fails cleanly inside the chat + transformer instead of being silently misrouted. + 3. `url` missing/empty -> fall back to body shape. `messages` wins + over the other keys so a malformed record carrying several of + them keeps its conversation instead of having it dropped, and a + bare `input` stays an embedding for backwards compatibility + (that ambiguity with `/v1/responses` is only resolvable from + `url`). """ - url = openai_jsonl_record.get("url") - if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: - return True - if url: - return False - body = openai_jsonl_record.get("body", {}) - if not isinstance(body, dict): - return False - return "input" in body and "messages" not in body + match openai_jsonl_record.get("url"): + case BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: + return BedrockBatchRecordKind.EMBEDDING + case BedrockFilesConfig.OPENAI_TEXT_COMPLETIONS_URL: + return BedrockBatchRecordKind.TEXT_COMPLETION + case BedrockFilesConfig.OPENAI_RESPONSES_URL: + return BedrockBatchRecordKind.RESPONSES + case None | "": + pass + case _: + return BedrockBatchRecordKind.CHAT + + body = openai_jsonl_record.get("body") + if not isinstance(body, Mapping): + return BedrockBatchRecordKind.CHAT + if "messages" in body: + return BedrockBatchRecordKind.CHAT + if "prompt" in body: + return BedrockBatchRecordKind.TEXT_COMPLETION + if "input" in body: + return BedrockBatchRecordKind.EMBEDDING + return BedrockBatchRecordKind.CHAT # Identifier for the Bedrock Titan v2 InvokeModel body schema as stored # in `model_prices_and_context_window.json`. Centralized so future @@ -546,9 +588,83 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) return dict(titan_config._transform_request(input=input_text, inference_params=inference_params)) + @staticmethod + def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body. + + Bedrock batch `modelInput` is the model's InvokeModel/Converse body, and + no Bedrock batch model takes a bare `prompt`, so the wrapping that + `litellm.text_completion` does in real time has to happen here too. + """ + prompt = openai_request_body.get("prompt") + if prompt is None: + raise ValueError( + "Batch record for /v1/completions is missing required `prompt` field: " + f"model={openai_request_body.get('model', '')}" + ) + return _frozen_mapping( + chain( + ((key, value) for key, value in openai_request_body.items() if key != "prompt"), + (("messages", text_completion_prompt_to_messages(prompt)),), + ) + ) + + @staticmethod + def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body. + + Delegates to the same Responses-to-Chat bridge the real-time path uses + for providers without a native Responses API (which is every Bedrock + model), so `input`, `instructions`, `max_output_tokens` and the tool + params translate identically in batch and real time. The bridge always + emits a `tools` key; an empty one is dropped rather than shipped as an + empty array inside `modelInput`. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_input = openai_request_body.get("input") + if responses_input is None: + raise ValueError( + "Batch record for /v1/responses is missing required `input` field: " + f"model={openai_request_body.get('model', '')}" + ) + chat_body = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=openai_request_body.get("model", ""), + input=_responses_input_adapter().validate_python(responses_input), + responses_api_request=_responses_request_adapter().validate_python( + _frozen_mapping( + (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") + ) + ), + ) + return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) + + @staticmethod + def _transform_batch_body_to_chat_body( + openai_request_body: Mapping[str, Any], + record_kind: BedrockBatchRecordKind, + ) -> Mapping[str, Any]: + """ + Normalize a non-embedding batch body to the Chat Completions shape the + per-provider Bedrock transformations expect. + """ + match record_kind: + case BedrockBatchRecordKind.TEXT_COMPLETION: + return BedrockFilesConfig._transform_text_completion_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.RESPONSES: + return BedrockFilesConfig._transform_responses_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.CHAT: + return openai_request_body + case BedrockBatchRecordKind.EMBEDDING: + raise ValueError("Embedding batch records do not have a chat-completion equivalent") + def _map_openai_to_bedrock_params( self, - openai_request_body: dict[str, Any], + openai_request_body: Mapping[str, Any], provider: str | None = None, ) -> dict[str, Any]: """ @@ -659,14 +775,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): provider = self.get_bedrock_invoke_provider(model) # Route to the embedding transformer when the OpenAI batch line - # targets /v1/embeddings; otherwise fall back to the existing - # chat-completion path. We branch here (rather than inside + # targets /v1/embeddings; every other endpoint shape is normalized + # to chat completions first. We branch here (rather than inside # `_map_openai_to_bedrock_params`) so the chat helper keeps its # narrow contract and the embedding helper can evolve independently. - if self._is_embedding_record(_openai_jsonl_content): + record_kind = self._classify_batch_record(_openai_jsonl_content) + if record_kind is BedrockBatchRecordKind.EMBEDDING: model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body) else: - model_input = self._map_openai_to_bedrock_params(openai_request_body=openai_body, provider=provider) + model_input = self._map_openai_to_bedrock_params( + openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), + provider=provider, + ) # Create Bedrock batch record record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index d9f8229dbed..f7a4682cb03 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union from typing_extensions import TYPE_CHECKING, Required, TypedDict, override @@ -1100,3 +1101,17 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): # supported subset and strips the field entirely when nothing remains, so # other edit types (e.g. `clear_thinking_20251015`) never reach Bedrock. context_management: dict + + +class BedrockBatchRecordKind(Enum): + """ + Which OpenAI endpoint shape a line of a Bedrock managed-batch JSONL file + carries. Bedrock batch `modelInput` is always the model's InvokeModel / + Converse body, so every non-embedding shape is normalized to Chat + Completions before being handed to the per-provider transformation. + """ + + CHAT = "chat" + TEXT_COMPLETION = "text_completion" + RESPONSES = "responses" + EMBEDDING = "embedding" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1b1db634ed2..d10ccf77703 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -721,3 +721,46 @@ class TestUnpackLegacyDefs: out = unpack_legacy_defs(schema) assert "components" not in out assert out["properties"]["r0"]["properties"]["p0"] == {"type": "string"} + + +class TestTextCompletionPromptToMessages: + """`/v1/completions` prompt wrapping, shared by the real-time and batch paths.""" + + def test_string_prompt_becomes_single_user_message(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages("summarize this") == ( + {"role": "user", "content": "summarize this"}, + ) + + def test_list_of_strings_becomes_one_message_each(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages(["first", "second"]) == ( + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ) + + @pytest.mark.parametrize( + "prompt", + [ + [1, 2, 3], + [[1, 2], [3, 4]], + ["ok", 7], + [], + "", + None, + {"role": "user"}, + ], + ) + def test_unsupported_prompt_shapes_raise(self, prompt): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + with pytest.raises(ValueError, match="non-empty string or a non-empty list of strings"): + text_completion_prompt_to_messages(prompt) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..87b03b02e1a 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1072,21 +1072,34 @@ class TestBedrockFilesEmbeddingTransformation: is None ) - def test_is_embedding_record_helper(self): - """Helper detects embeddings via `url` first, then by body shape.""" + def test_classify_batch_record_helper(self): + """Helper classifies by `url` first, then by body shape.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind - assert BedrockFilesConfig._is_embedding_record( - {"url": "/v1/embeddings", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/embeddings", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.EMBEDDING ) # body-only fallback - assert BedrockFilesConfig._is_embedding_record({"body": {"input": "x"}}) - # chat shape - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/chat/completions", "body": {"messages": []}} + assert ( + BedrockFilesConfig._classify_batch_record({"body": {"input": "x"}}) + is BedrockBatchRecordKind.EMBEDDING + ) + # chat shape + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/chat/completions", "body": {"messages": []}} + ) + is BedrockBatchRecordKind.CHAT + ) + # ambiguous body without any recognized key is treated as chat + assert ( + BedrockFilesConfig._classify_batch_record({"body": {}}) + is BedrockBatchRecordKind.CHAT ) - # ambiguous body without `input` is treated as not-embedding - assert not BedrockFilesConfig._is_embedding_record({"body": {}}) def test_explicit_chat_url_with_input_body_short_circuits_to_chat(self): """Explicit url=/v1/chat/completions wins even if body looks like embedding. @@ -1097,15 +1110,20 @@ class TestBedrockFilesEmbeddingTransformation: """ from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + # Direct helper assertion - assert not BedrockFilesConfig._is_embedding_record( - { - "url": "/v1/chat/completions", - "body": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "input": "this would mis-route under the old precedence", - }, - } + assert ( + BedrockFilesConfig._classify_batch_record( + { + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "input": "this would mis-route under the old precedence", + }, + } + ) + is BedrockBatchRecordKind.CHAT ) # End-to-end: a record like this routes through the chat path. We @@ -1164,20 +1182,301 @@ class TestBedrockFilesEmbeddingTransformation: with pytest.raises(ValueError, match="must be a string"): BedrockFilesConfig._coerce_embedding_input_to_string({"unsupported": True}) - def test_other_non_embedding_urls_route_to_chat(self): - """Any non-/v1/embeddings url short-circuits to chat path.""" + def test_other_non_embedding_urls_do_not_route_to_embeddings(self): + """An `input` body only means "embedding" when the url says so.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind # /v1/completions (legacy completions endpoint) - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/completions", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/completions", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.TEXT_COMPLETION + ) + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/responses", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.RESPONSES ) # Arbitrary unknown url - caller's explicit signal still wins - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/responses", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/moderations", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.CHAT ) +class TestBedrockBatchNonChatEndpointRecords: + """`/v1/completions` and `/v1/responses` JSONL records (issue #35639). + + Bedrock batch `modelInput` is always the model's InvokeModel/Converse body, + so a record shaped for another OpenAI endpoint has to be normalized to chat + completions first. Before this normalization every record below either + raised `BadRequestError` at `POST /v1/files` (Anthropic, Nova) or silently + shipped an empty `messages` list to AWS (passthrough providers). + """ + + ANTHROPIC_MODEL = "bedrock/us.anthropic.claude-sonnet-4-6" + NOVA_MODEL = "bedrock/us.amazon.nova-pro-v1:0" + PASSTHROUGH_MODEL = "bedrock/openai.gpt-oss-120b-1:0" + + def _transform(self, record: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content([record]) + assert len(result) == 1 + assert result[0]["recordId"] == record["custom_id"] + return result[0]["modelInput"] + + def test_anthropic_text_completion_record_wraps_prompt(self): + model_input = self._transform( + { + "custom_id": "1", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": "Summarize the following call transcript", + "max_tokens": 64, + }, + } + ) + + assert model_input == { + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "Summarize the following call transcript"}], + } + ], + "max_tokens": 64, + "anthropic_version": "bedrock-2023-05-31", + } + + def test_anthropic_text_completion_record_keeps_every_prompt_in_a_list(self): + model_input = self._transform( + { + "custom_id": "2", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": ["first prompt", "second prompt"], + "max_tokens": 8, + }, + } + ) + + # Consecutive user messages are merged by the Anthropic transform, the + # same way they are on the real-time path. + assert model_input["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "first prompt"}, + {"type": "text", "text": "second prompt"}, + ], + } + ] + assert "prompt" not in model_input + + def test_anthropic_responses_record_wraps_string_input(self): + model_input = self._transform( + { + "custom_id": "3", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "input": "hi", + "max_output_tokens": 16, + }, + } + ) + + assert model_input == { + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + "max_tokens": 16, + "anthropic_version": "bedrock-2023-05-31", + } + assert "tools" not in model_input, "an empty tools array must not be shipped to Bedrock" + + def test_anthropic_responses_record_maps_instructions_and_input_items(self): + """The Responses-specific params go through the same bridge as real time.""" + model_input = self._transform( + { + "custom_id": "4", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": "what is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "and 3+3?"}, + ], + "max_output_tokens": 32, + "temperature": 0.2, + }, + } + ) + + assert model_input["system"] == [{"type": "text", "text": "be terse"}] + assert model_input["max_tokens"] == 32 + assert model_input["temperature"] == 0.2 + assert [message["role"] for message in model_input["messages"]] == [ + "user", + "assistant", + "user", + ] + assert model_input["messages"][-1]["content"] == [{"type": "text", "text": "and 3+3?"}] + assert "input" not in model_input + assert "max_output_tokens" not in model_input + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_nova_converse_record_wraps_prompt_and_input(self, body): + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "5", + "method": "POST", + "url": url, + "body": {"model": self.NOVA_MODEL, **body}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_passthrough_provider_record_no_longer_emits_empty_messages(self, body): + """The passthrough branch used to emit `{"messages": [], "prompt": ...}`. + + That shape is accepted by `POST /v1/files`, so the whole batch job was + submitted to AWS and only failed there. + """ + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "6", + "method": "POST", + "url": url, + "body": {"model": self.PASSTHROUGH_MODEL, **body}, + } + ) + + # Asserted on the serialized form, since the passthrough branch hands + # `messages` straight to S3 without a per-provider transform. + assert json.loads(json.dumps(model_input)) == {"messages": [{"role": "user", "content": "hi"}]} + + def test_mixed_endpoints_in_one_file_keep_their_own_shapes(self): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "chat", + "url": "/v1/chat/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4, + }, + }, + { + "custom_id": "text", + "url": "/v1/completions", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + }, + { + "custom_id": "responses", + "url": "/v1/responses", + "body": {"model": self.ANTHROPIC_MODEL, "input": "hi", "max_output_tokens": 4}, + }, + { + "custom_id": "embedding", + "url": "/v1/embeddings", + "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "hi"}, + }, + ] + ) + + assert [record["recordId"] for record in result] == [ + "chat", + "text", + "responses", + "embedding", + ] + for record in result[:3]: + assert record["modelInput"]["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]} + ] + assert result[3]["modelInput"] == {"inputText": "hi"} + + @pytest.mark.parametrize( + ("url", "expected_message"), + [ + ("/v1/completions", "missing required `prompt` field"), + ("/v1/responses", "missing required `input` field"), + ], + ) + def test_missing_required_field_raises_actionable_error(self, url, expected_message): + with pytest.raises(ValueError, match=expected_message): + self._transform( + { + "custom_id": "7", + "method": "POST", + "url": url, + "body": {"model": self.ANTHROPIC_MODEL, "max_tokens": 4}, + } + ) + + def test_prompt_body_without_url_is_still_wrapped(self): + """A record can omit `url`; the body shape then decides.""" + model_input = self._transform( + { + "custom_id": "8", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + + def test_messages_win_over_prompt_when_url_is_absent(self): + model_input = self._transform( + { + "custom_id": "9", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "from messages"}], + "prompt": "from prompt", + "max_tokens": 4, + }, + } + ) + + assert model_input["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "from messages"}]} + ] + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" From 6def61e672d95c4b145613009cd3064d0a133475 Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 17:22:46 +0000 Subject: [PATCH 06/74] fix(bedrock): keep /v1/responses batch metadata through the chat bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/files/transformation.py | 1 + .../files/test_bedrock_files_transformation.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 0aa832780e5..baa2630556a 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -640,6 +640,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") ) ), + metadata=openai_request_body.get("metadata"), ) return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 87b03b02e1a..09aa6b1cf2d 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1337,6 +1337,23 @@ class TestBedrockBatchNonChatEndpointRecords: assert "input" not in model_input assert "max_output_tokens" not in model_input + def test_responses_record_keeps_metadata(self): + """`metadata` reaches the bridge, which reads it as its own kwarg.""" + model_input = self._transform( + { + "custom_id": "4b", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.PASSTHROUGH_MODEL, + "input": "hi", + "metadata": {"tenant": "acct-1"}, + }, + } + ) + + assert model_input["metadata"] == {"tenant": "acct-1"} + @pytest.mark.parametrize( "body", [ From 7d00f9d019f84be709a7515094fed4ce7bbee900 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:46:00 -0700 Subject: [PATCH 07/74] fix(managed_files): return unified output file ids from GET /batches list_user_batches parsed each stored batch blob and returned it as-is, so any row whose blob still carried raw provider file ids (for example a batch that reached a terminal state through the cost poller, or rows written before output registration existed) leaked raw output_file_id and error_file_id values that clients cannot fetch through the proxy. The list path now runs each row through ensure_batch_response_managed_file_ids, which swaps in existing managed ids and registers missing ones under the batch owner's identity, matching what GET /batches/{id} already does --- .../proxy/hooks/managed_files.py | 12 ++ .../proxy/hooks/test_managed_files.py | 142 ++++++++++++++++++ 2 files changed, 154 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 0036603bcd1..07a1f959940 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, @@ -352,6 +353,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) batch_obj = LiteLLMBatch.model_validate(batch_data) batch_obj.id = batch.unified_object_id + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=batch, + unified_batch_id=_is_base64_encoded_unified_file_id( + batch.unified_object_id + ), + ) batch_objects.append(batch_obj) except Exception as e: diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 50af6465d06..fc10a1257e1 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1813,6 +1813,148 @@ def _create_unified_batch_id(model_id: str, batch_id: str) -> str: return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=") +def _decode_unified_id(b64_id: str) -> str: + return base64.urlsafe_b64decode(b64_id + "=" * (-len(b64_id) % 4)).decode() + + +def _terminal_batch_record( + unified_batch_uid: str, + raw_input_file_id: str, + raw_output_file_id: str, + raw_error_file_id: str, +): + record = MagicMock() + record.unified_object_id = unified_batch_uid + record.created_by = "owner-user" + record.team_id = "owner-team" + record.status = "cancelled" + record.file_object = json.dumps( + { + "id": "batch-raw-456", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "cancelled", + "created_at": 1234567890, + "input_file_id": raw_input_file_id, + "output_file_id": raw_output_file_id, + "error_file_id": raw_error_file_id, + } + ) + return record + + +@pytest.mark.asyncio +async def test_list_batches_registers_and_returns_unified_output_file_ids(): + """A stored batch blob with raw provider file IDs (e.g. persisted by the cost + poller for a cancelled batch) must be listed with unified managed IDs, and the + output/error files must be registered in the managed file table so GET + /files/{id}/content can route them.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_input_file_id = "file-list-in-1" + raw_output_file_id = "file-list-out-1" + raw_error_file_id = "file-list-err-1" + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" + ).decode() + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ + _terminal_batch_record( + unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id + ) + ] + + input_file_row = MagicMock() + input_file_row.unified_file_id = unified_input_file_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_input_file_id: + return input_file_row + return None + + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + listed = result["data"][0] + assert listed.id == unified_batch_uid + assert listed.input_file_id == unified_input_file_id + + decoded_output = _decode_unified_id(listed.output_file_id) + assert decoded_output.startswith("litellm_proxy") + assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output + assert "llm_output_file_model_id,model-123" in decoded_output + assert "target_model_names,gpt-5-batch" in decoded_output + + decoded_error = _decode_unified_id(listed.error_file_id) + assert f"llm_output_file_id,{raw_error_file_id}" in decoded_error + + upsert_calls = prisma_client.db.litellm_managedfiletable.upsert.await_args_list + stored_raw_ids = { + c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls + } + assert stored_raw_ids == {raw_output_file_id, raw_error_file_id} + for c in upsert_calls: + assert c.kwargs["data"]["create"]["created_by"] == "owner-user" + assert c.kwargs["data"]["create"]["team_id"] == "owner-team" + + +@pytest.mark.asyncio +async def test_list_batches_resolves_existing_managed_rows_without_minting(): + """When the raw provider file IDs already have managed file rows, listing must + swap in the existing unified IDs and must not upsert duplicate rows.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_output_file_id = "file-list-out-existing" + existing_unified_output_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-9;llm_output_file_id,{raw_output_file_id}".encode() + ).decode() + + record = _terminal_batch_record( + unified_batch_uid, "file-list-in-9", raw_output_file_id, "" + ) + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [record] + + existing_row = MagicMock() + existing_row.unified_file_id = existing_unified_output_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_output_file_id: + return existing_row + return None + + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + assert result["data"][0].output_file_id == existing_unified_output_id + prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth From 59041240f036fe80776b297b36757c48d85f7978 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:23:32 -0700 Subject: [PATCH 08/74] fix(managed_files): cap batch list page size at 100 and bulk-resolve raw file ids in one query --- .../proxy/hooks/managed_files.py | 104 +++++++++++++----- .../openai_files_endpoints/common_utils.py | 30 +++++ .../proxy/hooks/test_managed_files.py | 94 +++++++++++----- 3 files changed, 172 insertions(+), 56 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 07a1f959940..6fa6ef46ad4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,6 +3,7 @@ import base64 import json +from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast from uuid import NAMESPACE_URL, uuid5 @@ -31,10 +32,12 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + apply_unified_file_ids, ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, + map_raw_file_ids_to_unified, normalize_mime_type_for_provider, resolve_managed_output_file_model_name, ) @@ -62,6 +65,9 @@ if TYPE_CHECKING: if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prisma.models import ( + LiteLLM_ManagedObjectTable as PrismaManagedObjectRow, + ) from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient @@ -75,6 +81,20 @@ else: PrismaClient = Any +def _decode_json_blob(blob: object) -> object: + return json.loads(blob) if isinstance(blob, str) else blob + + +def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]: + try: + batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) + except Exception as e: + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {e}") + return None + batch_obj.id = row.unified_object_id + return batch_obj + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -329,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"Invalid 'after' cursor: no batch found with id '{after}'.", ) - page_size = limit or 20 + page_size: Final = min(limit or 20, 100) cursor_args: Dict[str, Any] = ( {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} ) @@ -343,36 +363,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size - batch_objects: List[LiteLLMBatch] = [] - for batch in batches[:page_size]: - try: - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) - batch_obj = LiteLLMBatch.model_validate(batch_data) - batch_obj.id = batch.unified_object_id - await ensure_batch_response_managed_file_ids( - response=batch_obj, - managed_files_obj=self, - prisma_client=self.prisma_client, - verbose_proxy_logger=verbose_logger, - user_api_key_dict=user_api_key_dict, - db_batch_object=batch, - unified_batch_id=_is_base64_encoded_unified_file_id( - batch.unified_object_id - ), - ) - batch_objects.append(batch_obj) + parsed_rows: Final = tuple( + (row, batch_obj) + for row in batches[:page_size] + if (batch_obj := _parse_managed_batch_row(row)) is not None + ) + unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( + raw_file_ids=frozenset( + file_id + for _, batch_obj in parsed_rows + for file_id in (batch_obj.input_file_id, batch_obj.output_file_id, batch_obj.error_file_id) + if file_id and not _is_base64_encoded_unified_file_id(file_id) + ), + prisma_client=self.prisma_client, + ) + resolved_batches: Final = [ + await self._resolve_listed_batch( + row=row, + batch_obj=batch_obj, + unified_id_by_raw_id=unified_id_by_raw_id, + user_api_key_dict=user_api_key_dict, + ) + for row, batch_obj in parsed_rows + ] + return build_list_page( + [batch_obj for batch_obj in resolved_batches if batch_obj is not None], + has_more=has_more, + ) - except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {batch.unified_object_id}: {e}" - ) - continue - - return build_list_page(batch_objects, has_more=has_more) + async def _resolve_listed_batch( + self, + row: "PrismaManagedObjectRow", + batch_obj: LiteLLMBatch, + unified_id_by_raw_id: Mapping[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> Optional[LiteLLMBatch]: + apply_unified_file_ids(batch_obj, unified_id_by_raw_id) + try: + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=row, + unified_batch_id=_is_base64_encoded_unified_file_id( + row.unified_object_id + ), + ) + except Exception as e: + verbose_logger.warning( + f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}" + ) + return None + return batch_obj async def get_user_created_file_ids( self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str] diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 080b8b80ae4..bf83a7cf25c 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,6 +1,7 @@ import base64 import mimetypes import re +from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional @@ -16,6 +17,7 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedObjectTable from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.utils import LiteLLMBatch @@ -1002,6 +1004,34 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: pass +async def map_raw_file_ids_to_unified( + raw_file_ids: frozenset[str], prisma_client: "PrismaClient | None" +) -> Mapping[str, str]: + if not raw_file_ids or not prisma_client: + return MappingProxyType({}) + managed_files: Final = await ManagedFileRepository(prisma_client).table.find_many( + where={"flat_model_file_ids": {"hasSome": sorted(raw_file_ids)}} # mutable-ok: prisma where is a plain dict + ) + return MappingProxyType( + { + raw_id: managed_file.unified_file_id + for managed_file in managed_files + for raw_id in managed_file.flat_model_file_ids + if raw_id in raw_file_ids + } + ) + + +def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mapping[str, str]) -> None: + for file_attr, raw_id in ( + ("input_file_id", getattr(response, "input_file_id", None)), + ("output_file_id", getattr(response, "output_file_id", None)), + ("error_file_id", getattr(response, "error_file_id", None)), + ): + if isinstance(raw_id, str) and raw_id in unified_id_by_raw_id: + setattr(response, file_attr, unified_id_by_raw_id[raw_id]) + + async def ensure_batch_response_managed_file_ids( response, managed_files_obj, diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index fc10a1257e1..e1e5cc6c532 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1869,15 +1869,12 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id + input_file_row.flat_model_file_ids = [raw_input_file_id] - def find_managed_file(where): - if where["flat_model_file_ids"]["has"] == raw_input_file_id: - return input_file_row - return None - - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[input_file_row] ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client @@ -1892,6 +1889,13 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): assert listed.id == unified_batch_uid assert listed.input_file_id == unified_input_file_id + bulk_lookup = prisma_client.db.litellm_managedfiletable.find_many.await_args + assert set(bulk_lookup.kwargs["where"]["flat_model_file_ids"]["hasSome"]) == { + raw_input_file_id, + raw_output_file_id, + raw_error_file_id, + } + decoded_output = _decode_unified_id(listed.output_file_id) assert decoded_output.startswith("litellm_proxy") assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output @@ -1914,33 +1918,43 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): @pytest.mark.asyncio async def test_list_batches_resolves_existing_managed_rows_without_minting(): """When the raw provider file IDs already have managed file rows, listing must - swap in the existing unified IDs and must not upsert duplicate rows.""" + swap in the existing unified IDs via one bulk lookup for the whole page, with + no per-row queries and no duplicate upserts.""" from litellm.proxy._types import UserAPIKeyAuth - unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") - raw_output_file_id = "file-list-out-existing" - existing_unified_output_id = base64.urlsafe_b64encode( - f"litellm_proxy:application/json;unified_id,u-9;llm_output_file_id,{raw_output_file_id}".encode() + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-9;target_model_names,gpt-5-batch" ).decode() + raw_output_file_ids = ["file-list-out-existing-1", "file-list-out-existing-2"] + existing_unified_output_ids = [ + base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-{i};llm_output_file_id,{raw_id}".encode() + ).decode() + for i, raw_id in enumerate(raw_output_file_ids) + ] - record = _terminal_batch_record( - unified_batch_uid, "file-list-in-9", raw_output_file_id, "" - ) + records = [ + _terminal_batch_record( + _create_unified_batch_id("model-123", f"batch-{i}"), + unified_input_file_id, + raw_id, + "", + ) + for i, raw_id in enumerate(raw_output_file_ids) + ] prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_many.return_value = [record] + prisma_client.db.litellm_managedobjecttable.find_many.return_value = records - existing_row = MagicMock() - existing_row.unified_file_id = existing_unified_output_id + existing_rows = [ + MagicMock(unified_file_id=unified_id, flat_model_file_ids=[raw_id]) + for raw_id, unified_id in zip(raw_output_file_ids, existing_unified_output_ids) + ] - def find_managed_file(where): - if where["flat_model_file_ids"]["has"] == raw_output_file_id: - return existing_row - return None - - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=existing_rows ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock() proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client @@ -1951,10 +1965,38 @@ async def test_list_batches_resolves_existing_managed_rows_without_minting(): limit=10, ) - assert result["data"][0].output_file_id == existing_unified_output_id + assert [b.output_file_id for b in result["data"]] == existing_unified_output_ids + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once() + prisma_client.db.litellm_managedfiletable.find_first.assert_not_awaited() prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() +@pytest.mark.asyncio +async def test_list_batches_caps_page_size_at_100(): + """The list page size must be capped at 100 rows (matching OpenAI's limit) + even when the caller asks for more, so one request cannot fan out into an + unbounded scan.""" + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=100000, + ) + + assert ( + prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"] + == 101 + ) + assert result["data"] == [] + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth From 845680ed1dc1e2f4b6c4493a00289e2f9422bbf0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:50:09 -0700 Subject: [PATCH 09/74] test(proxy): unit test batch file id mapping helpers directly --- .../test_common_utils.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py new file mode 100644 index 00000000000..4a021627c3e --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py @@ -0,0 +1,97 @@ +import os +import sys +from types import MappingProxyType +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.openai_files_endpoints.common_utils import ( + apply_unified_file_ids, + map_raw_file_ids_to_unified, +) +from litellm.types.utils import LiteLLMBatch + + +def _batch(input_file_id, output_file_id, error_file_id) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status="cancelled", + output_file_id=output_file_id, + error_file_id=error_file_id, + ) + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_empty_ids_skips_db(): + prisma_client = MagicMock() + + assert await map_raw_file_ids_to_unified(frozenset(), prisma_client) == {} + + prisma_client.db.litellm_managedfiletable.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_no_prisma_client_returns_empty(): + assert await map_raw_file_ids_to_unified(frozenset({"file-raw-1"}), None) == {} + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_bulk_queries_and_filters_to_requested_ids(): + row_a = MagicMock( + unified_file_id="unified-a", + flat_model_file_ids=["file-raw-a", "file-raw-other"], + ) + row_b = MagicMock(unified_file_id="unified-b", flat_model_file_ids=["file-raw-b"]) + prisma_client = MagicMock() + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row_a, row_b]) + + mapping = await map_raw_file_ids_to_unified( + frozenset({"file-raw-b", "file-raw-a", "file-raw-missing"}), prisma_client + ) + + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with( + where={"flat_model_file_ids": {"hasSome": ["file-raw-a", "file-raw-b", "file-raw-missing"]}} + ) + assert dict(mapping) == {"file-raw-a": "unified-a", "file-raw-b": "unified-b"} + + +def test_apply_unified_file_ids_swaps_only_mapped_ids(): + batch = _batch(input_file_id="file-raw-in", output_file_id="file-raw-out", error_file_id=None) + + apply_unified_file_ids(batch, MappingProxyType({"file-raw-out": "unified-out"})) + + assert batch.input_file_id == "file-raw-in" + assert batch.output_file_id == "unified-out" + assert batch.error_file_id is None + + +def test_apply_unified_file_ids_swaps_all_three_ids(): + batch = _batch( + input_file_id="file-raw-in", + output_file_id="file-raw-out", + error_file_id="file-raw-err", + ) + + apply_unified_file_ids( + batch, + MappingProxyType( + { + "file-raw-in": "unified-in", + "file-raw-out": "unified-out", + "file-raw-err": "unified-err", + } + ), + ) + + assert (batch.input_file_id, batch.output_file_id, batch.error_file_id) == ( + "unified-in", + "unified-out", + "unified-err", + ) From 6b2ac7cc5b79bc173e5aea73f963a8b34089f885 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:25:47 -0700 Subject: [PATCH 10/74] refactor(bedrock): freeze the SSE-KMS header and key-source collections The staging merge tightened the LIT002 ceiling, so the three mutable dict literals this branch added now breach it. Build the S3 request headers as MappingProxyType and resolve the encryption key from a tuple of sources. --- litellm/llms/bedrock/common_utils.py | 10 +++---- litellm/llms/bedrock/files/transformation.py | 28 +++++++++++--------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index c5ef78b44f3..d18cb7d8734 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1314,11 +1314,11 @@ def resolve_s3_encryption_key_id( Precedence: `s3_encryption_key_id` in litellm_params, then optional_params (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. """ - for source in (litellm_params, optional_params or {}): - value = source.get("s3_encryption_key_id") - if isinstance(value, str) and value: - return value - return get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + candidates: Final = tuple( + source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None + ) + explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None) + return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") class CommonBatchFilesUtils: diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index b4c4a270bd7..728fd02f001 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -788,20 +788,24 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # Prepare headers with required S3 headers (same as s3_v2.py) sse_headers: Final = ( - { - "x-amz-server-side-encryption": "aws:kms", - "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, - } + MappingProxyType( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + ) if s3_encryption_key_id - else {} + else MappingProxyType({}) + ) + request_headers: Final = MappingProxyType( + { + "Content-Type": "application/json", # JSONL files are JSON content + "x-amz-content-sha256": content_hash, # REQUIRED by S3 + "Content-Language": "en", + "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, + } ) - request_headers: Final = { - "Content-Type": "application/json", # JSONL files are JSON content - "x-amz-content-sha256": content_hash, # REQUIRED by S3 - "Content-Language": "en", - "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", - **sse_headers, - } # Use requests.Request to prepare the request (same pattern as s3_v2.py) req: Final = requests.Request("PUT", api_base, data=content, headers=request_headers) From 8466ed0920021b765c0ec6483b0e96d121f2373b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 21:05:51 +0000 Subject: [PATCH 11/74] fix(websearch_interception): bill and rate limit intercepted searches against the calling key An intercepted web search called litellm.asearch() with only the search tool's litellm_params, so the search request carried no owner. The proxy's spend hook skips any call with no key, user or team attached, so the search's provider cost never reached SpendLogs; it was missing from the Logs page and never counted against the caller's budget. The same path never ran the rate limiter either, so an intercepted search was free of the key's RPM/TPM limits. The search now carries the originating key's attribution metadata (key hash, alias, user, team, org, plus model_group set to the resolved search tool) and runs the caller's rate limit checks before hitting the provider, matching what a direct /v1/search request gets. SDK calls with no proxy auth context are unchanged. Resolves LIT-5033 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../websearch_interception/handler.py | 46 ++++++++- .../test_websearch_interception_handler.py | 97 +++++++++++++++++-- 2 files changed, 132 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 71388134e98..6be64c89828 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -40,7 +40,7 @@ from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import CallTypes, LlmProviders, StandardLoggingUserAPIKeyMetadata from litellm.utils import ProviderConfigManager if TYPE_CHECKING: @@ -1288,10 +1288,14 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None search_litellm_params: dict[str, Any] = {} + search_tool_name: str | None = None if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) search_provider = search_litellm_params.get("search_provider") + selected_tool_name = search_tool.get("search_tool_name") + if isinstance(selected_tool_name, str) and selected_tool_name: + search_tool_name = selected_tool_name # Fallback to perplexity if no router or no search tools configured if not search_provider: @@ -1304,10 +1308,22 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider ) + user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs) + search_metadata: Final = ( + None + if user_api_key_auth is None + else self._build_search_request_metadata( + user_api_key_auth=user_api_key_auth, + search_tool_name=search_tool_name, + ) + ) search_kwargs: Final = { - key: value - for key, value in search_litellm_params.items() - if key != "search_provider" and value is not None + **{ + key: value + for key, value in search_litellm_params.items() + if key != "search_provider" and value is not None + }, + **({} if search_metadata is None else {"litellm_metadata": search_metadata}), } result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) @@ -1366,6 +1382,28 @@ class WebSearchInterceptionLogger(CustomLogger): team_object=team_object, ) + @staticmethod + def _build_search_request_metadata( + user_api_key_auth: "UserAPIKeyAuth", + search_tool_name: str | None, + ) -> dict[str, object]: + """ + Spend-tracking metadata for the intercepted search, so its provider cost is logged + and billed against the key/user/team that made the originating LLM request instead + of being dropped by the proxy's spend hook for lack of an owner. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_metadata: StandardLoggingUserAPIKeyMetadata = ( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) + ) + return { + **user_api_key_metadata, + **({} if search_tool_name is None else {"model_group": search_tool_name}), + "user_api_key": user_api_key_auth.api_key, + "user_api_key_auth": user_api_key_auth, + } + @staticmethod def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": if not kwargs: diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index b6ff3b70a4d..f39f41a6d12 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -221,14 +221,97 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, ) - mock_asearch.assert_awaited_once_with( - query="what is litellm", - search_provider="tavily", - api_key="fake-ui-key", - api_base="https://api.tavily.com", - timeout=10.0, - max_retries=2, + forwarded_kwargs = mock_asearch.await_args.kwargs + assert forwarded_kwargs["query"] == "what is litellm" + assert forwarded_kwargs["search_provider"] == "tavily" + assert forwarded_kwargs["api_key"] == "fake-ui-key" + assert forwarded_kwargs["api_base"] == "https://api.tavily.com" + assert forwarded_kwargs["timeout"] == 10.0 + assert forwarded_kwargs["max_retries"] == 2 + + +@pytest.mark.asyncio +async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch): + """An intercepted search is billed and logged against the key that made the LLM request. + + Without the forwarded attribution metadata the proxy's spend hook skips the search + entirely, so its provider cost never reaches SpendLogs or any budget. + """ + import litellm + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _should_track_cost_callback + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + user_api_key_auth = UserAPIKeyAuth( + api_key="hashed-sk-1234", + key_alias="alice-key", + user_id="user-alice", + org_id="org-1", + ) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search( + "what is litellm", + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, + ) + + forwarded_metadata = mock_asearch.await_args.kwargs["litellm_metadata"] + assert forwarded_metadata["user_api_key"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_hash"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_alias"] == "alice-key" + assert forwarded_metadata["user_api_key_user_id"] == "user-alice" + assert forwarded_metadata["user_api_key_org_id"] == "org-1" + assert forwarded_metadata["model_group"] == "perplexity-sonar-pro" + assert ( + _should_track_cost_callback( + user_api_key=forwarded_metadata["user_api_key"], + user_id=forwarded_metadata["user_api_key_user_id"], + team_id=forwarded_metadata["user_api_key_team_id"], + end_user_id=None, + call_type="asearch", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_execute_search_without_proxy_auth_context_stays_sdk_only(monkeypatch): + """SDK callers have no key to attribute the search to, so no proxy metadata is invented.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", + ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm", kwargs={"litellm_params": {}}) + + assert "litellm_metadata" not in mock_asearch.await_args.kwargs @pytest.mark.asyncio From a2806d430f55c2d372dbd633687d84ad207b5131 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:58:30 -0700 Subject: [PATCH 12/74] test(router): drop the duplicate s3_encryption_key_id credential test test_get_deployment_credentials_with_provider_bedrock_batch_fields already covers s3_encryption_key_id on the base branch, and the new test passes with every production file in this branch reverted, so it guards nothing. --- tests/test_litellm/test_router.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1576fa3dbcd..4a3395a7d3f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3976,37 +3976,6 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" -def test_get_deployment_credentials_with_provider_includes_s3_encryption_key_id(): - """ - Regression: s3_encryption_key_id must survive the CredentialLiteLLMParams filter, - otherwise the Bedrock batch input-file upload loses the SSE-KMS key and S3 rejects - the PutObject on buckets whose policy requires aws:kms encryption. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "bedrock-batch", - "litellm_params": { - "model": "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", - "aws_region_name": "us-west-2", - "s3_bucket_name": "my-batch-bucket", - "s3_encryption_key_id": "arn:aws:kms:us-west-2:1234:key/abcd", - }, - } - ], - ) - - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch" - ) - - assert credentials is not None - assert ( - credentials["s3_encryption_key_id"] - == "arn:aws:kms:us-west-2:1234:key/abcd" - ) - - def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves From d59a492585c56fad079fd51e02e494d8402d2aa4 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 6 Aug 2026 20:52:43 -0700 Subject: [PATCH 13/74] fix(azure_sentinel): respect AZURE_AUTHORITY_HOST for the Entra token and audience (#36137) The Azure Sentinel logger hardcoded the commercial Entra authority and the commercial Azure Monitor audience, so Log Analytics ingestion could not work in Azure Government even when the ingestion endpoint pointed at a sovereign Data Collection Endpoint. Resolve the authority from AZURE_AUTHORITY_HOST and derive the matching Logs Ingestion audience from it. Moving only the token URL is not enough: sovereign Entra would then be asked for a token scoped to the commercial audience, which the sovereign endpoint rejects. --- .../azure_sentinel/azure_sentinel.py | 46 ++++++++- .../integrations/test_azure_sentinel.py | 93 +++++++++++++++++++ 2 files changed, 137 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index f23317ae9df..563f815b582 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -16,7 +16,10 @@ import asyncio import os import time import traceback +from collections.abc import Mapping +from types import MappingProxyType from typing import Final +from urllib.parse import urlparse from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -27,6 +30,16 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com" +DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default" + +MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( + { + "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE, + "login.microsoftonline.us": "https://monitor.azure.us/.default", + } +) + class AzureSentinelLogger(CustomBatchLogger): """ @@ -42,6 +55,7 @@ class AzureSentinelLogger(CustomBatchLogger): client_id: str | None = None, client_secret: str | None = None, audit_stream_name: str | None = None, + authority_host: str | None = None, **kwargs, ): """ @@ -62,6 +76,10 @@ class AzureSentinelLogger(CustomBatchLogger): If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var. audit_stream_name (str, optional): Stream name from DCR for audit logs. If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name. + authority_host (str, optional): Microsoft Entra authority host that issues the OAuth2 token, + e.g. "https://login.microsoftonline.us" for Azure Government. If not provided, will use + AZURE_AUTHORITY_HOST env var or default to the Azure Public Cloud authority. The Azure + Monitor audience is derived from it. """ self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) @@ -76,6 +94,9 @@ class AzureSentinelLogger(CustomBatchLogger): resolved_client_secret: Final = ( client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) + resolved_authority_host: Final = self._normalize_authority_host( + authority_host or os.getenv("AZURE_AUTHORITY_HOST") or DEFAULT_AZURE_AUTHORITY_HOST + ) if not resolved_dcr_immutable_id: raise ValueError( @@ -119,7 +140,8 @@ class AzureSentinelLogger(CustomBatchLogger): ) # OAuth2 scope for Azure Monitor - self.oauth_scope = "https://monitor.azure.com/.default" + self.authority_host = resolved_authority_host + self.oauth_scope = self._resolve_oauth_scope(authority_host=resolved_authority_host) self.oauth_token: str | None = None self.oauth_token_expires_at: float | None = None @@ -129,6 +151,26 @@ class AzureSentinelLogger(CustomBatchLogger): self.log_queue: list[StandardLoggingPayload] = [] self.audit_log_queue: list[StandardAuditLogPayload] = [] + @staticmethod + def _normalize_authority_host(authority_host: str) -> str: + """ + Normalize an authority host into an absolute URL with no trailing slash. + + Accepts the scheme-qualified form litellm documents ("https://login.microsoftonline.us") + and the bare-host form the azure-identity AzureAuthorityHosts constants use. + """ + stripped: Final = authority_host.strip().rstrip("/") + return stripped if "://" in stripped else f"https://{stripped}" + + @staticmethod + def _resolve_oauth_scope(authority_host: str) -> str: + """ + Map an authority host to the Azure Monitor Logs Ingestion audience for the same cloud, + falling back to the Azure Public Cloud audience for an unrecognized host. + """ + host: Final = urlparse(authority_host).hostname or "" + return MONITOR_SCOPE_BY_AUTHORITY_HOST.get(host, DEFAULT_AZURE_MONITOR_SCOPE) + @staticmethod def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str: return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01" @@ -150,7 +192,7 @@ class AzureSentinelLogger(CustomBatchLogger): assert self.client_id is not None, "client_id is required" assert self.client_secret is not None, "client_secret is required" - token_url: Final = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + token_url: Final = f"{self.authority_host}/{self.tenant_id}/oauth2/v2.0/token" token_data: Final = { "client_id": self.client_id, diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 55e462c82b8..56662eea633 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -296,3 +296,96 @@ async def test_azure_sentinel_audit_stream_name_from_env_var(monkeypatch): ) assert explicit_logger.audit_stream_name == "Custom-LiteLLM-Explicit" + + +def _build_logger(**overrides): + kwargs = { + "dcr_immutable_id": "dcr-test123456789", + "endpoint": "https://test-dce.eastus-1.ingest.monitor.azure.com", + "tenant_id": "test-tenant-id", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + **overrides, + } + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + return AzureSentinelLogger(**kwargs) + + +@pytest.fixture +def _no_authority_host_env(monkeypatch): + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + + +@pytest.mark.parametrize( + "authority_host, expected_authority, expected_scope", + [ + (None, "https://login.microsoftonline.com", "https://monitor.azure.com/.default"), + ("https://login.microsoftonline.us", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("https://login.microsoftonline.us/", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("login.microsoftonline.us", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("https://adfs.contoso.example", "https://adfs.contoso.example", "https://monitor.azure.com/.default"), + ], +) +def test_azure_sentinel_resolves_authority_host_and_audience_together( + _no_authority_host_env, authority_host, expected_authority, expected_scope +): + """Both the Entra authority and the Azure Monitor audience must follow the configured cloud. + + Moving only the authority leaves a sovereign deployment asking sovereign Entra for the + commercial audience, which the sovereign ingestion endpoint rejects. + """ + logger = _build_logger(**({} if authority_host is None else {"authority_host": authority_host})) + + assert logger.authority_host == expected_authority + assert logger.oauth_scope == expected_scope + + +def test_azure_sentinel_authority_host_from_env_var(_no_authority_host_env, monkeypatch): + """AZURE_AUTHORITY_HOST is the documented setting and the string callback constructs the logger + with no arguments, so the env var alone has to move both values.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +@pytest.mark.asyncio +async def test_azure_sentinel_token_request_uses_sovereign_authority_and_audience(_no_authority_host_env): + """The resolved values must reach the wire, not just the instance attributes.""" + logger = _build_logger(authority_host="https://login.microsoftonline.us") + logger.log_queue.append( + StandardLoggingPayload( + id="test_id", + call_type="completion", + model="gpt-3.5-turbo", + status="success", + messages=[{"role": "user", "content": "Hello"}], + response={"choices": [{"message": {"content": "Hi"}}]}, + ) + ) + + mock_token_response = MagicMock() + mock_token_response.status_code = 200 + mock_token_response.json = MagicMock(return_value={"access_token": "test-bearer-token", "expires_in": 3600}) + mock_token_response.text = "Success" + mock_api_response = MagicMock() + mock_api_response.status_code = 204 + mock_api_response.text = "Success" + + async def mock_post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return mock_token_response + return mock_api_response + + logger.async_httpx_client.post = AsyncMock(side_effect=mock_post) + + await logger.async_send_batch() + + token_calls = [ + call for call in logger.async_httpx_client.post.call_args_list if "oauth2/v2.0/token" in call.kwargs["url"] + ] + assert len(token_calls) == 1 + assert token_calls[0].kwargs["url"] == "https://login.microsoftonline.us/test-tenant-id/oauth2/v2.0/token" + assert token_calls[0].kwargs["data"]["scope"] == "https://monitor.azure.us/.default" From 166a97e3690bac17d5e7ac0bc0bd7dbedb4c5ca5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 21:17:56 -0700 Subject: [PATCH 14/74] fix(anthropic adapter): type _handle_choiceless_chunk param as ModelResponseStream --- .../experimental_pass_through/adapters/streaming_iterator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index aeb990c7bb4..1660f56378f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -348,7 +348,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return self._augment_message_delta_usage(merged_chunk) - def _handle_choiceless_chunk(self, chunk: Any) -> bool: + def _handle_choiceless_chunk(self, chunk: "ModelResponseStream") -> bool: """Consume an OpenAI-compatible chunk that carries no ``choices``. ``choices`` is legitimately empty on metadata-only chunks; the final @@ -361,7 +361,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): into the held stop-reason chunk); False when the chunk should be skipped entirely. """ - if self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None: + if self.holding_stop_reason_chunk is not None and _optional_attr(chunk, "usage") is not None: self.chunk_queue.append(self._merge_usage_into_held_stop_reason_chunk(chunk)) self.queued_usage_chunk = True self.holding_stop_reason_chunk = None From 8ad75e5ae9fd50e13fb750193b88bc027dc7ed0f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 22:35:56 -0700 Subject: [PATCH 15/74] test(bedrock): cover the batch record classifier fallbacks and pin metadata handling --- .../test_bedrock_files_transformation.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 67e34507b09..270add48e0e 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1173,6 +1173,36 @@ class TestBedrockFilesEmbeddingTransformation: is BedrockBatchRecordKind.CHAT ) + @pytest.mark.parametrize("body", ["not a mapping", ["messages"], 7, None], ids=["str", "list", "int", "missing"]) + def test_classify_batch_record_falls_back_to_chat_for_non_mapping_body(self, body): + """A malformed body must not crash the whole upload during classification. + + Chat is the only kind whose transformer tolerates an unexpected shape and + raises a readable error; routing a non-mapping body anywhere else would + blow up on attribute access before the caller sees which record is bad. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + + record = {"body": body} if body is not None else {} + assert BedrockFilesConfig._classify_batch_record(record) is BedrockBatchRecordKind.CHAT + + def test_embedding_kind_is_rejected_by_the_chat_normalizer(self): + """Embeddings have no chat equivalent, so the normalizer refuses them outright. + + The caller routes embeddings to the Titan transformer before ever getting + here; this guard is what keeps a future caller from quietly shipping an + embedding body through the chat path. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + + with pytest.raises(ValueError, match="do not have a chat-completion equivalent"): + BedrockFilesConfig._transform_batch_body_to_chat_body( + {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "hi"}, + BedrockBatchRecordKind.EMBEDDING, + ) + def test_explicit_chat_url_with_input_body_short_circuits_to_chat(self): """Explicit url=/v1/chat/completions wins even if body looks like embedding. @@ -1426,6 +1456,35 @@ class TestBedrockBatchNonChatEndpointRecords: assert model_input["metadata"] == {"tenant": "acct-1"} + @pytest.mark.parametrize("model_attr", ["ANTHROPIC_MODEL", "NOVA_MODEL"], ids=["anthropic", "nova"]) + def test_modelled_providers_do_not_smuggle_metadata_into_the_bedrock_body(self, model_attr): + """Providers with a real InvokeModel schema leave `metadata` out of `modelInput`. + + Batch `modelInput` has to match the model's own InvokeModel body, and + neither the Anthropic messages body nor the Nova body has a field for + arbitrary caller labels. Nova in particular answers `400 Malformed input + request` for any key it does not recognize, so translating `metadata` + into the Converse-level `requestMetadata` would fail the record rather + than preserve the labels. The passthrough providers keep it because + their body is the OpenAI request itself. + """ + model_input = self._transform( + { + "custom_id": "4c", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": getattr(self, model_attr), + "input": "hi", + "max_output_tokens": 8, + "metadata": {"tenant": "acct-1"}, + }, + } + ) + + assert "metadata" not in model_input + assert "requestMetadata" not in model_input + @pytest.mark.parametrize( "body", [ From c736a227d5627e950f0df84ca33b48a30d607dfc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 22:48:27 -0700 Subject: [PATCH 16/74] style(bedrock): mark the batch normalizer locals Final for the tightened lint ceiling --- litellm/llms/bedrock/files/transformation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index dc9f32acdcc..bd3570d50a3 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -595,7 +595,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): no Bedrock batch model takes a bare `prompt`, so the wrapping that `litellm.text_completion` does in real time has to happen here too. """ - prompt = openai_request_body.get("prompt") + prompt: Final = openai_request_body.get("prompt") if prompt is None: raise ValueError( "Batch record for /v1/completions is missing required `prompt` field: " @@ -624,13 +624,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): LiteLLMCompletionResponsesConfig, ) - responses_input = openai_request_body.get("input") + responses_input: Final = openai_request_body.get("input") if responses_input is None: raise ValueError( "Batch record for /v1/responses is missing required `input` field: " f"model={openai_request_body.get('model', '')}" ) - chat_body = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + chat_body: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=openai_request_body.get("model", ""), input=_responses_input_adapter().validate_python(responses_input), responses_api_request=_responses_request_adapter().validate_python( From a0e35990cf6eb9eac80a7ad36ee1853eff7440da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 23:15:48 -0700 Subject: [PATCH 17/74] test: rename openai files common utils test to a unique basename --- .../{test_common_utils.py => test_files_common_utils.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/test_litellm/proxy/openai_files_endpoint/{test_common_utils.py => test_files_common_utils.py} (100%) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py similarity index 100% rename from tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py rename to tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py From e1717c5e9c90c637a594b394d1ae558939631232 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 7 Aug 2026 01:02:52 -0700 Subject: [PATCH 18/74] fix(proxy): return the real status code when a credential update is rejected (#36166) * fix(proxy): return the real status code when a credential update is rejected update_credential ended its except clause with 'return handle_exception_on_proxy(e)'. Returning the exception makes it the response body, so FastAPI answers 200 and every rejection on this route reads as a successful write to any caller that checks the status; the admin dashboard's API client is one. Patching a name that does not exist answered 200 with the real 404 buried in the body. The sibling handlers in this file already raise. The route had no test coverage, which is why it survived. * Update tests/test_litellm/proxy/credential_endpoints/test_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../proxy/credential_endpoints/endpoints.py | 2 +- .../credential_endpoints/test_endpoints.py | 91 +++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/credential_endpoints/test_endpoints.py diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index a6141ac9217..3b3e9692eda 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -346,4 +346,4 @@ async def update_credential( return {"success": True, "message": "Credential updated successfully"} except Exception as e: - return handle_exception_on_proxy(e) + raise handle_exception_on_proxy(e) diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py new file mode 100644 index 00000000000..e2fa1de6962 --- /dev/null +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -0,0 +1,91 @@ +"""Tests for the credential management endpoints.""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.types.utils import CredentialItem + +client = TestClient(app) + + +def _as_admin(): + return UserAPIKeyAuth(api_key="test-key", user_role="proxy_admin") + + +def _patch_credential(name: str, body: dict): + missing = object() + previous_override = app.dependency_overrides.get(user_api_key_auth, missing) + app.dependency_overrides[user_api_key_auth] = _as_admin + try: + return client.patch( + f"/credentials/{name}", + json=body, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + if previous_override is missing: + app.dependency_overrides.pop(user_api_key_auth, None) + else: + app.dependency_overrides[user_api_key_auth] = previous_override + + +def test_update_credential_answers_404_when_the_credential_does_not_exist(): + """Regression: the handler used to ``return handle_exception_on_proxy(e)``, which makes + the exception the response body and lets FastAPI answer 200, so a write the handler + rejected read as a success to every caller that checks the status. The dashboard's API + client branches on the status, so it reported a failed edit as applied.""" + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository: + repository.return_value.find_by_name = AsyncMock(return_value=None) + + response = _patch_credential( + "definitely-not-there", + {"credential_name": "definitely-not-there", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) + + assert response.status_code == 404, f"rejected write answered {response.status_code}: {response.text}" + assert "error" in response.json() + + +def test_update_credential_answers_500_when_the_database_is_not_connected(): + """The other rejection this handler raises must carry its own status too.""" + with patch("litellm.proxy.proxy_server.prisma_client", None): + response = _patch_credential( + "any-name", + {"credential_name": "any-name", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) + + assert response.status_code == 500, f"rejected write answered {response.status_code}: {response.text}" + + +def test_update_credential_still_answers_200_on_a_successful_write(): + """The fix must not turn a legitimate update into an error; the dashboard and the + Playwright credentials spec both assert the success path.""" + stored = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository: + repository.return_value.find_by_name = AsyncMock(return_value=stored) + repository.return_value.update_by_name = AsyncMock(return_value=None) + + response = _patch_credential( + "existing", + {"credential_name": "existing", "credential_values": {"api_key": "sk-new"}, "credential_info": {}}, + ) + + assert response.status_code == 200, response.text + assert response.json()["success"] is True From d332accabc829126cd5aa8b3b4889ae61f938686 Mon Sep 17 00:00:00 2001 From: Aayush Gid <93825737+aayush598@users.noreply.github.com> Date: Fri, 7 Aug 2026 20:46:41 +0530 Subject: [PATCH 19/74] fix(proxy): improve Headroom 404 compression error diagnostics (#35952) --- .../guardrail_hooks/headroom/headroom.py | 25 +++++++- .../guardrail_hooks/test_headroom.py | 58 +++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 61220819d48..8bfd5cca58a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -148,6 +148,27 @@ def _restore_protected_messages( ] +def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, object]: + """Build error details for failed /v1/compress responses. + + Adds troubleshooting hints for known deployment-related errors while + preserving the upstream status code and response body. + """ + if status_code == 404: + return { + "status_code": status_code, + "body": body, + "hint": ( + "The Headroom compression endpoint returned HTTP 404. " + "Verify that the configured Headroom endpoint is correct and that " + "the compression endpoint is available. If you are using a " + "self-hosted deployment, some deployments require enabling remote " + "compression (for example, HEADROOM_COMPRESS_ALLOW_REMOTE=1)." + ), + } + return {"status_code": status_code, "body": body} + + def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: hashes: Final[list[str]] = [] for msg in messages: @@ -417,7 +438,7 @@ class HeadroomGuardrail(CustomGuardrail): self._handle_compress_failure( messages, "Headroom compression service returned an error", - {"status_code": e.response.status_code, "body": e.response.text}, + _build_compress_failure_detail(e.response.status_code, e.response.text), ), False, {}, @@ -449,7 +470,7 @@ class HeadroomGuardrail(CustomGuardrail): self._handle_compress_failure( messages, "Headroom compression service returned an error", - {"status_code": response.status_code, "body": response.text}, + _build_compress_failure_detail(response.status_code, response.text), ), False, {}, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 00ab39357b4..7a2772ce78c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1040,6 +1040,64 @@ async def test_apply_guardrail_http_status_error_raises(): assert exc_info.value.status_code == 502 +@pytest.mark.asyncio +async def test_apply_guardrail_404_error_includes_troubleshooting_hint(): + """404 responses include a troubleshooting hint for self-hosted Headroom deployments.""" + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(404, "Not Found"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["status_code"] == 404 + assert exc_info.value.detail["body"] == "Not Found" + assert "hint" in exc_info.value.detail + assert "HEADROOM_COMPRESS_ALLOW_REMOTE=1" in exc_info.value.detail["hint"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_404_error_omits_troubleshooting_hint(): + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(500, "headroom internal error"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["status_code"] == 500 + assert exc_info.value.detail["body"] == "headroom internal error" + assert "hint" not in exc_info.value.detail + + @pytest.mark.asyncio async def test_apply_guardrail_http_status_error_fail_open_forwards_uncompressed(): guardrail = _make_guardrail(unreachable_fallback="fail_open") From 83ab6e08dae6f7b8d1e0bcaa3abd4840805a071e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 08:19:00 -0700 Subject: [PATCH 20/74] fix(proxy): invalidate cached project object on project update and delete (#36028) * fix(proxy): invalidate cached project object on /project/update and /project/delete The auth path reads projects cache-first via get_project_object with a 60s TTL and no freshness check, but no project write endpoint ever evicted the project_id:{id} cache entry. A project cached before /project/update added a model allowlist kept an empty models list in cache, so _run_project_checks skipped can_project_access_model and project-bound keys could call team models outside the project allowlist until the TTL expired. The same staleness applied to blocked status and budget fields, and /project/delete left the deleted project enforceable from cache. Evict the cache entry after the DB write in update_project and delete_project via a shared delete_cached_project_object helper, with the cache key derivation shared with get_project_object. * fix(proxy): broadcast project cache invalidation to all workers and make eviction best-effort Single-worker eviction leaves every other worker serving its in-memory copy of the mutated project until the 60s TTL expires, so a project allowlist change was still bypassable on multi-worker deployments. Add a coordination Redis pub/sub channel (litellm_proxy.auth_cache_invalidation): project eviction publishes the cache key and a per-worker subscriber deletes the local in-memory entry, with the next auth read refetching from the DB. Subscriber starts on any deployment with a coordination Redis and falls back to the TTL when none is configured. Also wrap the eviction in a best-effort catch: the DB write has already committed when eviction runs, so a cache backend error must not turn a successful update into a 500 or abort the remaining ids in /project/delete. * fix(lint): sort auth cache invalidation import and suppress best-effort shutdown catch The strict-budget gate flagged the new import block as un-sorted (I001) and the broad except in stop_auth_cache_invalidation_subscriber (BLE001); the catch is intentional since a failing stop must not break proxy shutdown, so it carries a named suppression instead of counting against the budget. --- .../management_endpoints/project_endpoints.py | 14 +- litellm/proxy/auth/auth_checks.py | 32 +++- .../auth_cache_invalidation_pubsub.py | 153 +++++++++++++++ litellm/proxy/proxy_server.py | 35 ++++ .../test_project_endpoints_prisma.py | 175 ++++++++++++++++++ .../proxy/auth/test_auth_checks.py | 57 ++++++ .../test_auth_cache_invalidation_pubsub.py | 161 ++++++++++++++++ 7 files changed, 625 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py create mode 100644 tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 1f693526d1f..66fac8d76ee 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import delete_cached_project_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field from litellm.proxy.management_helpers.utils import ( @@ -514,6 +515,7 @@ async def update_project( litellm_proxy_admin_name, premium_user, prisma_client, + user_api_key_cache, ) try: @@ -672,6 +674,11 @@ async def update_project( include={"litellm_budget_table": True, "object_permission": True}, ) + await delete_cached_project_object( + project_id=data.project_id, + user_api_key_cache=user_api_key_cache, + ) + return updated_project except Exception as e: verbose_proxy_logger.exception( @@ -710,7 +717,7 @@ async def delete_project( }' ``` """ - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache try: if not premium_user: @@ -773,6 +780,11 @@ async def delete_project( prisma_models.LiteLLM_ProjectTable | None ) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id}) + await delete_cached_project_object( + project_id=project_id, + user_api_key_cache=user_api_key_cache, + ) + deleted_projects.append(deleted_project) return deleted_projects diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3fba464dd23..01034d0cf58 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4262,6 +4262,10 @@ async def _project_soft_budget_check( ) +def _project_cache_key(project_id: str) -> str: + return f"project_id:{project_id}" + + async def get_project_object( project_id: str, prisma_client: PrismaClient | None, @@ -4279,7 +4283,7 @@ async def get_project_object( return None # Check cache first - cache_key: Final = f"project_id:{project_id}" + cache_key: Final = _project_cache_key(project_id) deserialized_project: Final = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, @@ -4310,6 +4314,32 @@ async def get_project_object( return project_obj +async def delete_cached_project_object( + project_id: str, + user_api_key_cache: UserApiKeyCache, +) -> None: + """ + Every endpoint that mutates litellm_projecttable must call this: get_project_object + serves auth cache-first with no freshness check, so without invalidation a stale + project (e.g. a pre-update empty model allowlist) keeps being enforced until the + TTL expires (LIT-3803). Best-effort on both steps: the DB write has already + committed, so a cache backend error must not fail the endpoint; the stale entry + then expires via TTL. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + + cache_key: Final = _project_cache_key(project_id) + try: + await user_api_key_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to evict cached project entry %s; a stale project may be served until its TTL expires: %s", + cache_key, + e, + ) + await publish_auth_cache_invalidation(cache_key=cache_key) + + async def _organization_max_budget_check( valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py new file mode 100644 index 00000000000..7fc8da42a3d --- /dev/null +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -0,0 +1,153 @@ +import asyncio +import json +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_sync_pubsub import ( + _ConfigSyncPubSub, + _pubsub_capable_client, + coordination_redis_cache, +) + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" +_POLL_TIMEOUT_SECONDS: Final = 1.0 +_BACKOFF_INITIAL_SECONDS: Final = 5.0 +_BACKOFF_MAX_SECONDS: Final = 60.0 + + +def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: + if redis_cache.namespace is None: + return AUTH_CACHE_INVALIDATION_CHANNEL + return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}" + + +@dataclass(frozen=True, slots=True) +class _CacheInvalidationMessage: + cache_key: str + + +def _cache_invalidation_message_json(cache_key: str) -> str: + return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) + + +def _cache_key_from_message_data(data: object) -> str | None: + if isinstance(data, bytes): + data = data.decode("utf-8", errors="replace") + if not isinstance(data, str): + return None + try: + parsed: Final = json.loads(data) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + cache_key: Final = parsed.get("cache_key") + return cache_key if isinstance(cache_key, str) else None + + +async def publish_auth_cache_invalidation(cache_key: str) -> None: + """ + Best-effort broadcast so every worker drops its local in-memory copy of a + mutated management object; without this, only the handling worker and Redis + are evicted and other workers keep serving the stale object until its TTL. + """ + redis_cache: Final = coordination_redis_cache() + if redis_cache is None: + return + try: + client: Final = _pubsub_capable_client(redis_cache) + if client is None: + verbose_proxy_logger.debug( + "auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support", + cache_key, + ) + return + await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors + verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) + + +class AuthCacheInvalidationSubscriber: + __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + + def __init__( + self, + redis_cache: "RedisCache", + user_api_key_cache: "UserApiKeyCache", + ) -> None: + self._redis_cache = redis_cache + self._user_api_key_cache = user_api_key_cache + self._task: asyncio.Task[None] | None = None + + def start(self) -> None: + if self._task is not None: + return + self._task = asyncio.create_task(self._run()) + + async def stop(self) -> None: + task: Final = self._task + if task is None: + return + self._task = None + _ = task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def _run(self) -> None: + backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects + while True: + try: + client = _pubsub_capable_client(self._redis_cache) # rebind-ok: re-resolved on every reconnect + if client is None: + verbose_proxy_logger.warning( + "auth cache invalidation subscriber disabled: cluster redis client has no pub/sub support; " + "cross-worker eviction falls back to the local cache TTL" + ) + return + pubsub = client.pubsub() # rebind-ok: fresh pubsub per reconnect + try: + await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache)) + backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: reset after successful subscribe + await self._consume(pubsub) + finally: + await self._close_pubsub(pubsub) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect + verbose_proxy_logger.warning( + "auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs", + e, + backoff_seconds, + ) + await asyncio.sleep(backoff_seconds) + backoff_seconds = min(backoff_seconds * 2, _BACKOFF_MAX_SECONDS) # rebind-ok: backoff accumulator + + async def _consume(self, pubsub: _ConfigSyncPubSub) -> None: + while True: + message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + if message is None: + continue + self._apply_message(message) + + def _apply_message(self, message: object) -> None: + data: Final = message.get("data") if isinstance(message, dict) else None + cache_key: Final = _cache_key_from_message_data(data) + if cache_key is None: + return + in_memory_cache: Final = self._user_api_key_cache.in_memory_cache + if in_memory_cache is not None: + in_memory_cache.delete_cache(cache_key) + + @staticmethod + async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: + try: + await pubsub.aclose() + except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2e24a2d4f3c..14c6f8b8779 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -297,6 +297,9 @@ from litellm.proxy.common_request_processing import ( _should_return_raw_model_name, create_response, ) +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, +) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers @@ -1193,6 +1196,8 @@ async def proxy_startup_event(app: FastAPI): await proxy_config.stop_config_sync_subscriber() + await proxy_config.stop_auth_cache_invalidation_subscriber() + await proxy_shutdown_event() @@ -3904,6 +3909,7 @@ class ProxyConfig: self._last_hashicorp_vault_config: dict[str, Any] | None = None self.worker_registry: list[WorkerRegistryEntry] = [] self.config_sync_subscriber: ConfigSyncSubscriber | None = None + self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None from litellm.litellm_core_utils.get_model_cost_map import ( get_model_cost_map_loaded_at, ) @@ -6391,6 +6397,30 @@ class ProxyConfig: except Exception as e: verbose_proxy_logger.error("Error stopping config sync subscriber: %s", e) + def start_auth_cache_invalidation_subscriber( + self, + redis_cache: RedisCache | None, + user_api_key_cache: UserApiKeyCache, + ) -> None: + if redis_cache is None or self.auth_cache_invalidation_subscriber is not None: + return + subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=redis_cache, + user_api_key_cache=user_api_key_cache, + ) + self.auth_cache_invalidation_subscriber = subscriber + subscriber.start() + + async def stop_auth_cache_invalidation_subscriber(self) -> None: + subscriber: Final = self.auth_cache_invalidation_subscriber + if subscriber is None: + return + self.auth_cache_invalidation_subscriber = None + try: + await subscriber.stop() + except Exception as e: # noqa: BLE001 # best-effort: a failing stop must not break proxy shutdown + verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e) + async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): """ Use this to read non-llm objects from the db and initialize them @@ -8330,6 +8360,11 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + proxy_config.start_auth_cache_invalidation_subscriber( + redis_cache=redis_usage_cache, + user_api_key_cache=user_api_key_cache, + ) + if store_model_in_db is True: ### GET STORED CREDENTIALS ### scheduler.add_job( diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c55b66b402b..c29b4c68bb0 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -864,3 +864,178 @@ def test_litellm_project_table_has_timestamp_fields(): fields = LiteLLM_ProjectTable.model_fields assert "created_at" in fields, "LiteLLM_ProjectTable must have created_at field" assert "updated_at" in fields, "LiteLLM_ProjectTable must have updated_at field" + + +@pytest.mark.asyncio +async def test_update_project_invalidates_cached_project_object(monkeypatch): + """ + LIT-3803 regression: auth reads projects cache-first with no freshness check, + so /project/update must evict the cached project. Before the fix, a project + cached with models=[] kept bypassing the new allowlist until the TTL expired. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth.auth_checks import get_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + cache = UserApiKeyCache() + + stale_row = MagicMock() + stale_row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": []} + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=stale_row) + + seeded = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert seeded is not None and seeded.models == [] + + updated_models = ["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"] + existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None) + updated_row = MagicMock() + updated_row.model_dump = lambda: { + "project_id": project_id, + "team_id": None, + "models": updated_models, + } + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + + await update_project( + data=UpdateProjectRequest(project_id=project_id, models=updated_models), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=updated_row) + refreshed = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert refreshed is not None + assert refreshed.models == updated_models + + +@pytest.mark.asyncio +async def test_delete_project_invalidates_cached_project_object(monkeypatch): + """ + LIT-3803 regression: /project/delete must evict the cached project so auth + stops enforcing (or trusting) a project that no longer exists. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth.auth_checks import get_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + cache = UserApiKeyCache() + + row = MagicMock() + row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": ["gpt-5.5"]} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=row) + + seeded = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert seeded is not None + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_projecttable.delete = AsyncMock(return_value=row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + + await delete_project( + data=DeleteProjectRequest(project_ids=[project_id]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + assert ( + await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + is None + ) + + +@pytest.mark.asyncio +async def test_update_project_succeeds_when_cache_eviction_fails(monkeypatch): + """ + The DB write has already committed when eviction runs, so a cache backend + error must not turn a successful update into a 500; the stale entry is + bounded by the TTL instead. + """ + from unittest.mock import AsyncMock, MagicMock + + project_id = f"project-{uuid.uuid4()}" + failing_cache = MagicMock() + failing_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down")) + + existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None) + updated_row = MagicMock() + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", failing_cache) + + response = await update_project( + data=UpdateProjectRequest(project_id=project_id, models=["gpt-oss-120b"]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + assert response is updated_row + failing_cache.async_delete_cache.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_project_eviction_publishes_cross_worker_invalidation(monkeypatch): + """ + Single-worker eviction only fixes the handling worker; the broadcast is what + lets every other worker drop its in-memory copy instead of serving the stale + project until the TTL expires. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import delete_cached_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new=AsyncMock(), + ) as mock_publish: + await delete_cached_project_object( + project_id=project_id, user_api_key_cache=UserApiKeyCache() + ) + + mock_publish.assert_awaited_once_with(cache_key=f"project_id:{project_id}") diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a5211ba83e7..757991b8ff3 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5464,6 +5464,63 @@ async def test_get_project_object_db_fetch_returns_cached_obj(): assert result.project_alias == "proj" +@pytest.mark.asyncio +async def test_project_allowlist_enforced_when_key_models_empty(): + """ + LIT-3803: a project-bound key with models=[] has no key-level restriction, + but the project allowlist must still 403 team models outside it. + """ + from litellm.proxy._types import ( + LiteLLM_ProjectTableCachedObj, + ProxyErrorTypes, + ProxyException, + ) + from litellm.proxy.auth.auth_checks import _run_project_checks, can_key_call_model + + valid_token = UserAPIKeyAuth( + api_key="hashed-key", + project_id="p-1", + team_id="t-1", + models=[], + ) + project = LiteLLM_ProjectTableCachedObj( + project_id="p-1", + team_id="t-1", + models=["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"], + ) + + assert ( + await can_key_call_model( + model="gemini-2.5-flash", + llm_model_list=None, + valid_token=valid_token, + llm_router=None, + ) + is True + ) + + await _run_project_checks( + project_object=project, + _model="gemini-2.5-flash-image", + llm_router=None, + skip_budget_checks=True, + valid_token=valid_token, + proxy_logging_obj=MagicMock(), + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_project_checks( + project_object=project, + _model="gemini-2.5-flash", + llm_router=None, + skip_budget_checks=True, + valid_token=valid_token, + proxy_logging_obj=MagicMock(), + ) + assert exc_info.value.type == ProxyErrorTypes.project_model_access_denied + assert exc_info.value.code == "403" + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py new file mode 100644 index 00000000000..468e8aabae8 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -0,0 +1,161 @@ +import asyncio +import json +from typing import Iterable, List, Optional, Tuple +from unittest.mock import patch + +import pytest +from redis.asyncio import Redis + +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AUTH_CACHE_INVALIDATION_CHANNEL, + AuthCacheInvalidationSubscriber, + publish_auth_cache_invalidation, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +class _RecordingRedisClient(Redis): + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _FailingPublishRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + raise ConnectionError("redis down") + + +class _QueuePubSub: + def __init__(self, initial_messages: Iterable[object] = ()) -> None: + self.queue: "asyncio.Queue[object]" = asyncio.Queue() + for message in initial_messages: + self.queue.put_nowait(message) + self.subscribed_channels: List[str] = [] + self.closed = False + + async def subscribe(self, *channels: str) -> None: + self.subscribed_channels.extend(channels) + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: + try: + return await asyncio.wait_for(self.queue.get(), timeout) + except asyncio.TimeoutError: + return None + + async def aclose(self) -> None: + self.closed = True + + +class _ScriptedPubSubRedisClient(Redis): + def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: + self._scripted_pubsubs = iter(pubsubs) + + def pubsub(self) -> _QueuePubSub: + return next(self._scripted_pubsubs) + + +class _FakeRedisCache: + def __init__(self, client: object, namespace: Optional[str] = None) -> None: + self._client = client + self.namespace = namespace + + def init_async_client(self) -> object: + return self._client + + +def _invalidation_message(cache_key: str) -> dict: + return {"type": "message", "data": json.dumps({"cache_key": cache_key}).encode()} + + +@pytest.mark.asyncio +async def test_publish_sends_cache_key_json_on_channel() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + assert client.published == [(AUTH_CACHE_INVALIDATION_CHANNEL, json.dumps({"cache_key": "project_id:p-1"}))] + + +@pytest.mark.asyncio +async def test_publish_uses_namespaced_channel() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client, namespace="ns1"), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + assert client.published[0][0] == f"ns1:{AUTH_CACHE_INVALIDATION_CHANNEL}" + + +@pytest.mark.asyncio +async def test_publish_noops_without_coordination_redis() -> None: + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=None, + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + +@pytest.mark.asyncio +async def test_publish_swallows_redis_errors() -> None: + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=_FailingPublishRedisClient()), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + +@pytest.mark.asyncio +async def test_subscriber_deletes_local_cache_entry_on_message() -> None: + """ + The cross-worker half of LIT-3803: a worker that did not handle the project + mutation must drop its in-memory copy when the invalidation broadcast lands, + instead of serving the stale object until the TTL expires. + """ + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + for _ in range(200): + if cache.in_memory_cache.get_cache("project_id:p-1") is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert cache.in_memory_cache.get_cache("project_id:p-1") is None + assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] + + +@pytest.mark.asyncio +async def test_subscriber_ignores_malformed_messages() -> None: + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + user_api_key_cache=cache, + ) + subscriber._apply_message({"type": "message", "data": b"not json"}) + subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) + subscriber._apply_message("raw string") + subscriber._apply_message(None) + + assert cache.in_memory_cache.get_cache("project_id:p-1") is not None From 527dc0a8bbaa582b8770979a02e2b06e421505e7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 08:40:13 -0700 Subject: [PATCH 21/74] feat(proxy): add apply_user_budget_to_team_keys opt-in (#36102) * feat(proxy): add apply_user_budget_to_team_keys opt-in PR #32005 made a user's personal max_budget apply to their team-scoped keys too, and PR #35271 reverted the whole thing (behavior plus the skip_user_budget_on_team_key opt-out) because that flipped the default for everyone. This brings the behavior back the other way round: default is unchanged, and general_settings.apply_user_budget_to_team_keys opts a deployment into charging the key owner's personal budget on team keys. The flag reaches all three personal-budget gates so an opted-in deployment enforces consistently: the read-time check in common_checks, the optimistic reservation counter in _get_budget_counters, and the _PROXY_MaxBudgetLimiter pre-call hook. It is also in the /config/list allowed args and, unlike the reverted flag, in the _update_general_settings propagation allowlist, so the Admin UI General Settings toggle actually takes effect at runtime; an explicit YAML value still wins over the DB value on reload. get_config_list's allowed_args moves to a module-level frozen mapping of field name to type string, dropping 18 LIT002 violations and rebuilding one less dict per request. * style(proxy): drop explanatory comments from the budget flag paths --- litellm/proxy/_types.py | 10 +++ litellm/proxy/auth/auth_checks.py | 39 +++++----- litellm/proxy/auth/user_api_key_auth.py | 1 + litellm/proxy/hooks/max_budget_limiter.py | 9 ++- litellm/proxy/proxy_server.py | 60 +++++++++------ .../spend_tracking/budget_reservation.py | 6 +- .../proxy/auth/test_auth_checks.py | 74 +++++++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 40 ++++++++++ .../proxy/hooks/test_max_budget_limiter.py | 29 ++++++++ .../proxy/test_budget_reservation.py | 41 ++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 65 ++++++++++++++++ type-discipline-budget.json | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 ++ 13 files changed, 334 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b1134650f2..5b7dc3a7c73 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2488,6 +2488,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "is active as a reminder that hard enforcement is relaxed." ), ) + apply_user_budget_to_team_keys: bool | None = Field( + None, + description=( + "If True, a user's personal max_budget is enforced on every request they " + "make, including requests made with a team-scoped key. Defaults to False, " + "where a team-scoped key is governed only by the team and team-member " + "budgets and the key owner's personal max_budget does not apply " + "(see GitHub issue #12905)." + ), + ) user_url_validation: bool | None = Field( None, description=( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 01034d0cf58..3da899a5610 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -648,28 +648,29 @@ async def common_checks( ) async def _user_max_budget_check() -> None: - # 4.1 personal budget, if personal key - if ( - (team_object is None or team_object.team_id is None) - and user_object is not None - and user_object.max_budget is not None - ): - from litellm.proxy.proxy_server import get_current_spend + # 4.1 personal budget + if user_object is None or user_object.max_budget is None: + return + is_team_key: Final = team_object is not None and team_object.team_id is not None + if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True: + return - user_budget: Final = user_object.max_budget - user_spend: Final = await get_current_spend( - counter_key=f"spend:user:{user_object.user_id}", - fallback_spend=user_object.spend or 0.0, + from litellm.proxy.proxy_server import get_current_spend + + user_budget: Final = user_object.max_budget + user_spend: Final = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, + max_budget=user_budget, + ) + if math.isfinite(user_budget) and user_spend >= user_budget: + raise litellm.BudgetExceededError( + current_cost=user_spend, max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", + entity_type=Litellm_EntityType.USER.value, + entity_id=user_object.user_id, ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - entity_type=Litellm_EntityType.USER.value, - entity_id=user_object.user_id, - ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9248576b599..9dc450befc2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2470,6 +2470,7 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, ) diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index b8a69705b99..eaf37b0bcf1 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -32,9 +32,12 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): if max_budget is None or user_id is None: return - # Personal budget applies only to non-team requests, matching - # the explicit team-key exemption in common_checks section 4.1. - if user_api_key_dict.team_id is not None: + from litellm.proxy.proxy_server import general_settings + + if ( + user_api_key_dict.team_id is not None + and general_settings.get("apply_user_budget_to_team_keys") is not True + ): return # The reservation path admits at the strict-`<` boundary and diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 14c6f8b8779..84f8685f5dc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17,7 +17,7 @@ import traceback import warnings from collections.abc import AsyncGenerator, Callable, Mapping from datetime import datetime, timedelta, timezone -from types import UnionType +from types import MappingProxyType, UnionType from typing import ( TYPE_CHECKING, Any, @@ -6096,6 +6096,15 @@ class ProxyConfig: else: general_settings["disable_auto_add_proxy_admin_to_teams"] = value if value is None else bool(value) + if "apply_user_budget_to_team_keys" in _general_settings and ( + "apply_user_budget_to_team_keys" not in self._yaml_general_settings_keys + ): + db_value: Final = _general_settings["apply_user_budget_to_team_keys"] + if isinstance(db_value, str): + general_settings["apply_user_budget_to_team_keys"] = db_value.lower() == "true" + else: + general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] @@ -14982,6 +14991,29 @@ Keep it more precise, to prevent overwrite other values unintentially _PLUGIN_KEY_REDACTED: Final = "***" +_GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingProxyType( + { + "max_parallel_requests": "Integer", + "global_max_parallel_requests": "Integer", + "max_request_size_mb": "Integer", + "max_response_size_mb": "Integer", + "proxy_config_reload_interval_seconds": "Integer", + "pass_through_endpoints": "PydanticModel", + "store_model_in_db": "Boolean", + "store_prompts_in_spend_logs": "Boolean", + "maximum_spend_logs_retention_period": "String", + "mcp_internal_ip_ranges": "List", + "mcp_trusted_proxy_ranges": "List", + "mcp_xff_num_trusted_hops": "Integer", + "always_include_stream_usage": "Boolean", + "forward_client_headers_to_llm_api": "Boolean", + "mcp_required_fields": "List", + "cancel_on_disconnect": "Boolean", + "disable_auto_add_proxy_admin_to_teams": "Boolean", + "apply_user_budget_to_team_keys": "Boolean", + } +) + def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object: """Restore real plugin_key values the client never sees. @@ -15480,25 +15512,7 @@ async def get_config_list( else: db_general_settings_dict = {} - allowed_args: Final = { - "max_parallel_requests": {"type": "Integer"}, - "global_max_parallel_requests": {"type": "Integer"}, - "max_request_size_mb": {"type": "Integer"}, - "max_response_size_mb": {"type": "Integer"}, - "proxy_config_reload_interval_seconds": {"type": "Integer"}, - "pass_through_endpoints": {"type": "PydanticModel"}, - "store_model_in_db": {"type": "Boolean"}, - "store_prompts_in_spend_logs": {"type": "Boolean"}, - "maximum_spend_logs_retention_period": {"type": "String"}, - "mcp_internal_ip_ranges": {"type": "List"}, - "mcp_trusted_proxy_ranges": {"type": "List"}, - "mcp_xff_num_trusted_hops": {"type": "Integer"}, - "always_include_stream_usage": {"type": "Boolean"}, - "forward_client_headers_to_llm_api": {"type": "Boolean"}, - "mcp_required_fields": {"type": "List"}, - "cancel_on_disconnect": {"type": "Boolean"}, - "disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"}, - } + allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES return_val: Final = [] @@ -15506,7 +15520,7 @@ async def get_config_list( if field_name in allowed_args: ## HANDLE TYPED DICT - typed_dict_type = allowed_args[field_name]["type"] + typed_dict_type = allowed_args[field_name] if typed_dict_type == "PydanticModel": if field_name == "pass_through_endpoints": @@ -15548,7 +15562,7 @@ async def get_config_list( _response_obj = ConfigList( field_name=field_name, - field_type=allowed_args[field_name]["type"], + field_type=allowed_args[field_name], field_description=field_info.description or "", field_value=_redact_general_setting_value( field_name, @@ -15576,7 +15590,7 @@ async def get_config_list( _response_obj = ConfigList( field_name=field_name, - field_type=allowed_args[field_name]["type"], + field_type=allowed_args[field_name], field_description=field_info.description or "", field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin), stored_in_db=_stored_in_db, diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 907c4ac7344..58a85171cc7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -156,6 +156,7 @@ async def reserve_budget_for_request( proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: Any | None = None, + apply_user_budget_to_team_keys: bool = False, fail_closed_budget_enforcement: bool = False, ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): @@ -175,6 +176,7 @@ async def reserve_budget_for_request( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + apply_user_budget_to_team_keys=apply_user_budget_to_team_keys, ) if not counters: return None @@ -332,6 +334,7 @@ async def _get_budget_counters( proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: Any | None = None, + apply_user_budget_to_team_keys: bool = False, ) -> list[_BudgetCounter]: counters: Final[list[_BudgetCounter]] = [] @@ -380,8 +383,9 @@ async def _get_budget_counters( ) ) + is_team_key: Final = team_object is not None and team_object.team_id is not None if ( - (team_object is None or team_object.team_id is None) + (not is_team_key or apply_user_budget_to_team_keys) and user_object is not None and user_object.user_id is not None and user_object.max_budget is not None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 757991b8ff3..298f8a31b64 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5105,6 +5105,80 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): assert result is True +@pytest.mark.asyncio +async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag_enabled(): + """general_settings.apply_user_budget_to_team_keys opts a deployment into + charging the key owner's personal budget on team-scoped keys too. + + Same fixture as the default-off test above, so a regression that ignores the + flag lets this call through instead of raising. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=1000.0) + token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + async def _no_membership(*args, **kwargs): + return None + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={"apply_user_budget_to_team_keys": True}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + assert "ExceededBudget: User=u1" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_common_checks_personal_user_budget_still_enforced_on_personal_key_with_flag_enabled(): + """The flag only widens enforcement to team keys; personal keys keep blocking.""" + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + token = UserAPIKeyAuth(token="k1", user_id="u1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ): + with pytest.raises(litellm.BudgetExceededError): + await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={"apply_user_budget_to_team_keys": True}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + @pytest.mark.parametrize( "scope, route, expect_blocked", [ diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bc5bb877bd0..6fc44ef7519 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -218,6 +218,46 @@ async def test_fail_closed_budget_enforcement_reaches_reservation( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings,expected_flag", + [ + ({"apply_user_budget_to_team_keys": True}, True), + ({"apply_user_budget_to_team_keys": False}, False), + ({}, False), + ], +) +async def test_apply_user_budget_to_team_keys_reaches_reservation( + general_settings, expected_flag +): + """The opt-in lives in general_settings but is consumed inside + _get_budget_counters, so it has to be threaded through reserve_budget_for_request + or the reservation path keeps exempting team keys while the read path enforces.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=None), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings=general_settings, + ) + + assert ( + mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag + ) + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py index 0074d7062b8..71671966d1a 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -185,6 +185,35 @@ async def test_team_keys_skip_personal_budget(): mock_get_spend.assert_not_awaited() +@pytest.mark.asyncio +async def test_team_keys_enforce_personal_budget_when_flag_enabled(): + """This hook is the third personal-budget gate alongside common_checks and the + reservation path, so apply_user_budget_to_team_keys has to reach it too or an + opted-in deployment enforces in two places out of three.""" + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_max_budget=10.0, + team_id="team-1", + ) + + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"apply_user_budget_to_team_keys": True}, + ), patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=999.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + + @pytest.mark.asyncio async def test_no_max_budget_passes(): handler = _PROXY_MaxBudgetLimiter() diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1a584423fac..34adb4d2091 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -650,6 +650,47 @@ async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_should_reserve_user_budget_counter_for_team_key_when_flag_enabled(spend_counter_state): + """apply_user_budget_to_team_keys must widen the reservation path too. + + Read-time enforcement alone leaks budget under concurrency, so the opt-in has + to reserve against the personal counter as well or a burst of team-key + requests slips past the owner's max_budget. + """ + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-user-on-team-flagged", + spend=0.0, + user_id="user-on-team-flagged", + team_id="team-no-budget", + ) + team_object = LiteLLM_TeamTable(team_id="team-no-budget", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="user-on-team-flagged", spend=0.0, max_budget=5.0) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + apply_user_budget_to_team_keys=True, + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-flagged") == pytest.approx(0.3) + + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state): counter_cache, key_cache = spend_counter_state diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index efd2ccb3e53..580a58885d9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7043,6 +7043,39 @@ async def test_update_general_settings_store_model_in_db_false(): assert ps.general_settings["store_model_in_db"] is False +@pytest.mark.asyncio +async def test_update_general_settings_propagates_apply_user_budget_to_team_keys(): + """The Admin UI toggle writes to the DB config, so the flag has to be in the + runtime propagation allowlist. The reverted skip_user_budget_on_team_key was + exposed in /config/list but never propagated, so its toggle did nothing.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": "true"}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["apply_user_budget_to_team_keys"] is True + + +@pytest.mark.asyncio +async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins(): + """A DB value must not silently override an explicit YAML setting on reload.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"apply_user_budget_to_team_keys"} + + with patch("litellm.proxy.proxy_server.general_settings", {"apply_user_budget_to_team_keys": True}): + await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": False}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["apply_user_budget_to_team_keys"] is True + + @pytest.mark.asyncio @pytest.mark.parametrize( "db_value,expected", @@ -9536,6 +9569,38 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): app.dependency_overrides.clear() +def test_get_config_list_includes_apply_user_budget_to_team_keys(monkeypatch): + """Related to #12905: the opt-in must be discoverable via /config/list so it + renders as a Boolean toggle on the Admin UI General Settings table. This needs + both the ConfigGeneralSettings field and the allowed_args entry.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert "apply_user_budget_to_team_keys" in fields + assert fields["apply_user_budget_to_team_keys"]["field_type"] == "Boolean" + finally: + app.dependency_overrides.clear() + + def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..2921592acbd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23256 }, "LIT002": { - "limit": 27213 + "limit": 27195 }, "LIT003": { "limit": 269 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 752572f9863..7675c0506c0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23389,6 +23389,11 @@ export interface components { * @description Proxy API Endpoints you want users to be able to access */ allowed_routes?: unknown[] | null; + /** + * Apply User Budget To Team Keys + * @description If True, a user's personal max_budget is enforced on every request they make, including requests made with a team-scoped key. Defaults to False, where a team-scoped key is governed only by the team and team-member budgets and the key owner's personal max_budget does not apply (see GitHub issue #12905). + */ + apply_user_budget_to_team_keys?: boolean | null; /** * Background Health Checks * @description run health checks in background From c9292d3af2fb9a8aab2111a26082e5852b61dc75 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 09:57:30 -0700 Subject: [PATCH 22/74] fix(proxy): stop alerting on health probes that lose the planned engine-restart race (#36141) --- litellm/proxy/db/prisma_client.py | 32 ++ litellm/proxy/utils.py | 118 ++++- .../db/test_prisma_planned_engine_restart.py | 444 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 8 + 4 files changed, 593 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 6863687081c..5f86490a474 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -486,6 +486,38 @@ class PrismaWrapper: os.environ[self._db_url_env_var] = _db_url return _db_url + @property + def engine_generation(self) -> int: + """How many query-engine replacements have completed on this wrapper. + + Bumped under `_reconnection_lock` only after a replacement engine has + connected, so a change across an await proves a *successful* planned + replacement happened in between — a replacement that failed (a real + outage) leaves it untouched. + """ + return self._engine_generation + + async def _reconnection_settled(self) -> None: + async with self._reconnection_lock: + pass + + async def wait_for_planned_engine_replacement(self, timeout_seconds: float) -> None: + """Wait, bounded, for an in-flight planned engine replacement to finish. + + Both replacement paths (`recreate_prisma_client` and + `_safe_refresh_token`) hold `_reconnection_lock` across their whole + kill/connect window, so re-acquiring it means the replacement has + settled one way or the other. Gives up silently on timeout: a caller + that stopped waiting must treat the replacement as not completed and + consult `engine_generation` rather than assume success. + """ + if timeout_seconds <= 0 or not self._reconnection_lock.locked(): + return + try: + await asyncio.wait_for(self._reconnection_settled(), timeout=timeout_seconds) + except asyncio.TimeoutError: + return + async def recreate_prisma_client( self, new_db_url: str, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cdd41d2ed42..605455a2f73 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3000,6 +3000,13 @@ class PrismaClient: ] = [] # mutable-ok: drained queue, mirrors tool_usage_transactions _autorouter_turn_transactions_lock = asyncio.Lock() + # How long a health probe failure waits for an in-flight planned engine + # replacement to settle before deciding whether to report itself. Generous + # against a replacement that takes well under a second, and far short of the + # reconnect budget an outage-hung `connect()` runs under, so a real outage + # is never waited out. + PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS: ClassVar[float] = 5.0 + def __init__( self, database_url: str, @@ -4970,6 +4977,101 @@ class PrismaClient: else: verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e) + def _probe_target_wrapper(self) -> PrismaWrapper: + """The Prisma wrapper a `SELECT 1` health probe actually reaches. + + `health_check()` issues `query_raw`, which `RoutingPrismaWrapper` sends + to the reader unless the reader is degraded. The writer's engine state + therefore says nothing about a probe that failed against the reader, so + the gate has to follow the same routing rule the probe did. + """ + if isinstance(self.db, RoutingPrismaWrapper): + return self.db.writer if self.db.reader_unavailable else self.db.reader + return self.db + + async def _run_health_probe(self, wrapper: PrismaWrapper) -> object: + """Issue the `SELECT 1` a health check is made of, against `wrapper`. + + Takes the wrapper rather than re-reading `self.db`, because routing is + re-resolved on every attribute access: a reader that recovers between + the caller picking its target and the query going out would send the + probe to a different engine than the one whose generation the caller is + about to check, and attribute the failure to the wrong replacement. + """ + sql_query: Final = "SELECT 1" + response: Final = await wrapper.query_raw(sql_query) + return response + + async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: + try: + await self._run_health_probe(wrapper) + except Exception as probe_error: # noqa: BLE001 # any failure means the database is not answering + verbose_proxy_logger.debug("Prisma health_check() confirmation probe failed: %s", probe_error) + return False + return True + + async def _planned_engine_replacement_absorbed( + self, + e: Exception, + wrapper: PrismaWrapper, + generation_before: int, + ) -> bool: + """True iff `e` is a connection-class probe failure that a completed + planned query-engine replacement explains. + + Planned replacements (RDS IAM token refresh, guarded reconnect) kill the + running query engine and spawn a new one. A `SELECT 1` probe that races + that sub-second window fails with a transport error against the engine's + local HTTP port even though nothing is wrong with the database, and + reporting it drives a false-positive `db_exceptions` alert on every + replacement. + + Two things must both hold, because neither is sufficient alone. The + engine generation must have moved, which says a replacement completed + rather than merely being attempted: reconnect attempts during a real + outage hold the same lock for tens of seconds, so gating on an in-flight + replacement would swallow most of an outage's alerts. And a fresh probe + must succeed, because `Prisma.connect()` polls the query engine's own + `/status` endpoint rather than round-tripping to the database, so a + future engine that binds before it validates its connection pool would + let the generation advance with the database still unreachable. + + Waiting for an in-flight replacement to settle is what makes the + generation check meaningful, since the generation has not moved yet at + the instant the probe fails. The wait is generous against a replacement + that takes well under a second and short enough that an outage-hung + reconnect is not waited out; a replacement that has not settled by then + reports rather than stays silent. + """ + if not PrismaDBExceptionHandler.is_database_connection_error(e): + return False + await wrapper.wait_for_planned_engine_replacement(self.PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS) + if wrapper.engine_generation == generation_before: + return False + return await self._probe_answers_now(wrapper) + + async def _report_health_check_failure( + self, + e: Exception, + duration: float, + traceback_str: str, + wrapper: PrismaWrapper, + generation_before: int, + ) -> None: + if await self._planned_engine_replacement_absorbed(e, wrapper, generation_before): + verbose_proxy_logger.info( + "Prisma health_check() connection error raced a planned query-engine replacement; " + "not reporting it as a DB exception: %s", + e, + ) + return + await self.proxy_logging_obj.failure_handler( + original_exception=e, + duration=duration, + call_type="health_check", + traceback_str=traceback_str, + ) + @backoff.on_exception( backoff.expo, Exception, @@ -4982,13 +5084,10 @@ class PrismaClient: Health check endpoint for the prisma client """ start_time: Final = time.time() + probe_wrapper: Final = self._probe_target_wrapper() + generation_before: Final = probe_wrapper.engine_generation try: - sql_query: Final = "SELECT 1" - - # Execute the raw query - # The asterisk before `user_id_list` unpacks the list into separate arguments - response: Final = await self.db.query_raw(sql_query) - return response + return await self._run_health_probe(probe_wrapper) except Exception as e: import traceback @@ -4998,11 +5097,12 @@ class PrismaClient: end_time: Final = time.time() _duration: Final = end_time - start_time asyncio.create_task( - self.proxy_logging_obj.failure_handler( - original_exception=e, + self._report_health_check_failure( + e=e, duration=_duration, - call_type="health_check", traceback_str=error_traceback, + wrapper=probe_wrapper, + generation_before=generation_before, ) ) raise e diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py index 18a238b7545..c8e0338eeaa 100644 --- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py +++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py @@ -6,13 +6,24 @@ subprocess), and the engine-death watcher / in-flight transport-error retries must not treat that planned restart as a crash and recreate the client a second time. +A planned restart is also invisible to the database itself, so a ``SELECT 1`` +health probe that races the kill/connect window fails with a transport error +against the engine's local HTTP port. That failure must not be reported as a +``db_exceptions`` DB failure, or every IAM refresh cycle raises a false alarm. + Symbols pinned here: - ``PrismaWrapper._expected_engine_deaths`` - ``PrismaWrapper._engine_generation`` + - ``PrismaWrapper.engine_generation`` + - ``PrismaWrapper.wait_for_planned_engine_replacement`` - ``PrismaWrapper.on_engine_replaced`` - ``PrismaWrapper.recreate_prisma_client`` (expected_generation guard) - ``PrismaWrapper._safe_refresh_token`` (refresh coalescing) - ``RoutingPrismaWrapper.recreate_prisma_client`` (guard forwarding) + - ``PrismaClient.health_check`` (planned-replacement alert suppression) + - ``PrismaClient._probe_target_wrapper`` + - ``PrismaClient._probe_answers_now`` + - ``PrismaClient._planned_engine_replacement_absorbed`` """ import asyncio @@ -21,22 +32,29 @@ import signal import sys import urllib.parse from datetime import datetime, timedelta +from typing import Any, List from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest from prisma import Prisma as GeneratedPrisma +from prisma.engine.errors import EngineConnectionError sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.utils import PrismaClient @pytest.fixture(autouse=True) def mock_prisma_binary(): """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" mock_module = MagicMock() + # Production code isinstance-checks against this, which a bare MagicMock + # attribute cannot satisfy. + mock_module.engine.errors.EngineConnectionError = EngineConnectionError with patch.dict(sys.modules, {"prisma": mock_module}): yield mock_module @@ -49,6 +67,87 @@ def _make_wrapper(engine_pid: int = 111, iam: bool = False) -> PrismaWrapper: return PrismaWrapper(original_prisma=prisma, iam_token_db_auth=iam) +def _make_prisma_client(db: Any) -> PrismaClient: + """A ``PrismaClient`` whose ``db`` is a real wrapper and whose alerting + hook is observable.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + client = PrismaClient( + database_url="postgresql://user:pass@localhost:5432/db", + proxy_logging_obj=proxy_logging_obj, + ) + client.db = db + client._db_watchdog_reconnect_timeout_seconds = 5.0 + return client + + +_real_asyncio_sleep = asyncio.sleep + + +async def _yield_to_loop(times: int = 10) -> None: + """Let already-scheduled tasks make progress. + + Bound to the real ``asyncio.sleep`` at import time: the tests below patch + ``asyncio.sleep`` to skip the SIGTERM/SIGKILL grace, and an ``AsyncMock`` + stand-in never yields to the event loop, which would silently leave every + background task un-started and the assertions vacuous. + """ + for _ in range(times): + await _real_asyncio_sleep(0) + + +async def _await_health_check_reports() -> None: + """Await the fire-and-forget reporting tasks ``health_check()`` scheduled. + + Selected by coroutine qualname rather than by draining every pending task, + so an unrelated background task can never make these assertions pass by + accident. + """ + reports = [ + task + for task in asyncio.all_tasks() + if getattr(task.get_coro(), "__qualname__", "") + == "PrismaClient._report_health_check_failure" + ] + if reports: + await asyncio.gather(*reports, return_exceptions=True) + + +def _fails_then_answers(error: Exception, failures: int = 3) -> Any: + """Raise ``error`` for the first ``failures`` probes, then answer. + + ``health_check`` retries up to three times, so this exhausts the retries and + still lets the confirmation probe that decides suppression succeed. Without + that, a test would report for the wrong reason: the confirmation probe would + fail too, masking whether the error type was classified at all. + """ + seen: List[int] = [] + + async def _query_raw(_sql: str) -> Any: + seen.append(1) + if len(seen) <= failures: + raise error + return [{"?column?": 1}] + + return _query_raw + + +def _blocking_replacement(gate: asyncio.Event, fail: bool = False) -> MagicMock: + """A replacement Prisma whose ``connect()`` parks until ``gate`` is set. + + Holds ``_reconnection_lock`` open for as long as the test needs, which is + how a health probe is made to fail *while* a planned replacement is in + flight rather than after it. + """ + + async def _connect(*_: Any, **__: Any) -> None: + await gate.wait() + if fail: + raise ConnectionRefusedError("database is down") + + return MagicMock(connect=AsyncMock(side_effect=_connect)) + + def _token_db_url(created: datetime, expires_in: int = 900) -> str: """Build a DATABASE_URL whose password is a parseable RDS IAM token.""" token = ( @@ -603,3 +702,348 @@ async def test_recreate_caps_expected_engine_deaths_set(mock_prisma_binary): await wrapper.recreate_prisma_client("postgresql://new") assert wrapper._expected_engine_deaths == {111} + + +@pytest.mark.asyncio +async def test_wait_for_planned_engine_replacement_returns_once_recreate_settles( + mock_prisma_binary, +): + wrapper = _make_wrapper(engine_pid=111) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + waiter = asyncio.create_task(wrapper.wait_for_planned_engine_replacement(5.0)) + await _yield_to_loop() + blocked_while_in_flight = not waiter.done() + + gate.set() + await recreate + await waiter + + assert { + "blocked_while_in_flight": blocked_while_in_flight, + "generation": wrapper.engine_generation, + } == {"blocked_while_in_flight": True, "generation": 1} + + +@pytest.mark.asyncio +async def test_wait_for_planned_engine_replacement_gives_up_at_timeout( + mock_prisma_binary, +): + """A replacement that never settles must not stall the caller forever; the + caller then sees an unchanged generation and reports the failure.""" + wrapper = _make_wrapper(engine_pid=111) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + await asyncio.wait_for( + wrapper.wait_for_planned_engine_replacement(0.05), timeout=5.0 + ) + gave_up_with_replacement_still_in_flight = not recreate.done() + + gate.set() + await recreate + + assert gave_up_with_replacement_still_in_flight is True + + +@pytest.mark.asyncio +async def test_health_check_does_not_alert_when_probe_races_a_completed_replacement( + mock_prisma_binary, +): + """The reported bug: an IAM-refresh engine recreate makes a concurrent + readiness probe fail transiently, and that failure was alerting as a DB + exception on every refresh cycle. + + The reporting task is drained while the replacement is still in flight, + which is when it runs in production; a decision taken at that instant sees + an engine generation that has not moved yet. + """ + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + probe_result = await client.health_check() + + drain = asyncio.create_task(_await_health_check_reports()) + await _yield_to_loop() + alerts_while_replacement_in_flight = ( + client.proxy_logging_obj.failure_handler.await_count + ) + + gate.set() + await recreate + await drain + + assert { + "probe_result": probe_result, + "probe_attempts": wrapper.query_raw.await_count, + "alerts_while_in_flight": alerts_while_replacement_in_flight, + "alerts": client.proxy_logging_obj.failure_handler.await_count, + } == { + "probe_result": [{"?column?": 1}], + "probe_attempts": 3, + "alerts_while_in_flight": 0, + "alerts": 0, + } + + +@pytest.mark.asyncio +async def test_health_check_alerts_when_a_completed_replacement_still_cannot_reach_the_database( + mock_prisma_binary, +): + """``Prisma.connect()`` polls the query engine's own ``/status`` endpoint + rather than round-tripping to the database, so a replacement can complete + against a database that is still unreachable. The engine generation alone + must not be enough to stay silent.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + gate.set() + await recreate + await _await_health_check_reports() + + assert { + "replacement_completed": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"replacement_completed": 1, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_check_alerts_when_the_replacement_never_completes( + mock_prisma_binary, +): + """A real outage also has a replacement in flight, but it fails, so the + engine generation never advances and the probe failure must still alert.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate, fail=True) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + gate.set() + with pytest.raises(ConnectionRefusedError): + await recreate + await _await_health_check_reports() + + call_types: List[str] = [ + c.kwargs["call_type"] + for c in client.proxy_logging_obj.failure_handler.await_args_list + ] + assert { + "generation": wrapper.engine_generation, + "alerted": len(call_types) > 0, + "call_types": set(call_types), + } == {"generation": 0, "alerted": True, "call_types": {"health_check"}} + + +@pytest.mark.asyncio +async def test_health_check_alerts_for_non_connection_errors_during_a_replacement( + mock_prisma_binary, +): + """Suppression is scoped to transport failures. A query the database itself + rejected is a real defect and must alert even mid-replacement, and even + though the database is plainly reachable a moment later.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock(side_effect=_fails_then_answers(ValueError("malformed SELECT"))) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(ValueError): + await client.health_check() + + gate.set() + await recreate + await _await_health_check_reports() + + assert { + "generation": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"generation": 1, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_check_alerts_for_a_transient_failure_with_no_engine_replacement( + mock_prisma_binary, +): + """Suppression is scoped to failures a planned replacement explains. A + transport blip that self-heals with no engine replacement at all still + alerts, so the gate cannot be widened into silencing every failure whose + database happens to answer a moment later.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + + probe_result = await client.health_check() + await _await_health_check_reports() + + assert { + "probe_result": probe_result, + "generation": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"probe_result": [{"?column?": 1}], "generation": 0, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_probe_stays_on_its_target_when_reader_availability_flips( + mock_prisma_binary, monkeypatch +): + """Routing is re-resolved on every attribute access, so a reader that + recovers mid-call would otherwise send the probe to a different engine than + the one whose generation is being checked, and blame the wrong replacement. + The probe follows the wrapper it was handed.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader") + writer = _make_wrapper(engine_pid=111) + reader = _make_wrapper(engine_pid=222) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client = _make_prisma_client(routing) + writer.query_raw = AsyncMock(return_value=[{"writer": 1}]) + reader.query_raw = AsyncMock(return_value=[{"reader": 1}]) + + routing._reader_unavailable = False + target = client._probe_target_wrapper() + routing._reader_unavailable = True + result = await client._run_health_probe(target) + + assert { + "target_is_reader": target is reader, + "result": result, + "reader_probes": reader.query_raw.await_count, + "writer_probes": writer.query_raw.await_count, + } == { + "target_is_reader": True, + "result": [{"reader": 1}], + "reader_probes": 1, + "writer_probes": 0, + } + + +@pytest.mark.asyncio +async def test_health_check_consults_the_reader_wrapper_under_read_replica_routing( + mock_prisma_binary, monkeypatch +): + """``query_raw`` is routed to the reader, so a reader-side planned + replacement is the one that explains a probe failure.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader") + writer = _make_wrapper(engine_pid=111) + reader = _make_wrapper(engine_pid=222) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client = _make_prisma_client(routing) + reader.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + reader.recreate_prisma_client("postgresql://new-reader") + ) + await _yield_to_loop() + assert reader._reconnection_lock.locked() is True + + probe_result = await client.health_check() + + drain = asyncio.create_task(_await_health_check_reports()) + await _yield_to_loop() + alerts_while_replacement_in_flight = ( + client.proxy_logging_obj.failure_handler.await_count + ) + + gate.set() + await recreate + await drain + + assert { + "probe_target_is_the_reader": client._probe_target_wrapper() is reader, + "writer_generation": writer.engine_generation, + "reader_generation": reader.engine_generation, + "probe_result": probe_result, + "alerts_while_in_flight": alerts_while_replacement_in_flight, + "alerts": client.proxy_logging_obj.failure_handler.await_count, + } == { + "probe_target_is_the_reader": True, + "writer_generation": 0, + "reader_generation": 1, + "probe_result": [{"?column?": 1}], + "alerts_while_in_flight": 0, + "alerts": 0, + } diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index abd6220144b..a8e81e92ebd 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1096,6 +1096,7 @@ async def test_prisma_health_check_failure_names_itself_at_operator_visible_leve function and reads as "the check never ran", and reporting it only at debug level hides a database fault behind a flag nobody enables in production.""" import logging + from functools import partial from unittest.mock import AsyncMock from litellm.proxy.utils import PrismaClient @@ -1103,6 +1104,9 @@ async def test_prisma_health_check_failure_names_itself_at_operator_visible_leve client = MagicMock() client.db.query_raw = AsyncMock(side_effect=Exception("connection refused")) client.proxy_logging_obj.failure_handler = AsyncMock() + client._probe_target_wrapper = MagicMock(return_value=client.db) + client._run_health_probe = partial(PrismaClient._run_health_probe, client) + client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): with pytest.raises(Exception, match="connection refused"): @@ -1142,6 +1146,7 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): text can carry a full connection string, so the credential has to be gone from the emitted record.""" import logging + from functools import partial from unittest.mock import AsyncMock from litellm.proxy.utils import PrismaClient @@ -1151,6 +1156,9 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): side_effect=Exception("could not connect to postgresql://admin:hunter2@db.internal:5432/litellm") ) client.proxy_logging_obj.failure_handler = AsyncMock() + client._probe_target_wrapper = MagicMock(return_value=client.db) + client._run_health_probe = partial(PrismaClient._run_health_probe, client) + client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): with pytest.raises(Exception): From 6ba744b340fd0901f9bfab428a846bc1677ef3c4 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 09:57:37 -0700 Subject: [PATCH 23/74] test(docker): gate the componentized gateway and backend images on an arbitrary-uid offline boot (#36136) --- .github/ci-coverage-allowlist.yml | 7 - .github/workflows/image-scan.yml | 65 ++++++ .../test_component_image_serves_offline.py | 187 ++++++++++++++++++ 3 files changed, 252 insertions(+), 7 deletions(-) create mode 100644 tests/proxy_migration_tests/test_component_image_serves_offline.py diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 140f1155e3a..1423228e725 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -136,13 +136,6 @@ test_paths: - tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py dockerfiles: - - reason: >- - The componentized images the microservices chart deploys are built by no job; wiring both into - the scan workflow costs a full image build each and is deferred to a change that prices the - whole set - paths: - - backend/Dockerfile - - gateway/Dockerfile - reason: >- The dashboard container is a static Next.js export served by nginx, and the dashboard build and lint workflows already exercise that output, so building the image adds no signal about it diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index e0c0bfcedae..8faf3ef6229 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -12,6 +12,11 @@ on: - docker/Dockerfile.non_root - migrations/Dockerfile - migrations/run.py + - gateway/Dockerfile + - gateway/main.py + - backend/Dockerfile + - backend/main.py + - docker/component_entrypoint.sh - litellm-proxy-extras/** - tests/proxy_migration_tests/** - uv.lock @@ -147,3 +152,63 @@ jobs: run: | python -m pip install "pytest==9.0.3" python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + + gateway-image: + name: gateway-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build gateway image + run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the gateway serves offline as a non-root uid + env: + LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }} + LITELLM_COMPONENT_PORT: "4000" + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v + + backend-image: + name: backend-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build backend image + run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the backend serves offline as a non-root uid + env: + LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }} + LITELLM_COMPONENT_PORT: "4001" + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v diff --git a/tests/proxy_migration_tests/test_component_image_serves_offline.py b/tests/proxy_migration_tests/test_component_image_serves_offline.py new file mode 100644 index 00000000000..528c2960072 --- /dev/null +++ b/tests/proxy_migration_tests/test_component_image_serves_offline.py @@ -0,0 +1,187 @@ +"""Image-level regression net for the prisma bake in the componentized images. + +The gateway and backend serve requests; they never shell out to the Prisma CLI +(``PrismaManager.setup_database`` is reachable only from ``proxy_cli.py``, which +uvicorn'ing ``gateway.main:app`` bypasses). What they do need is the generated +client's baked query engine, and prisma-python resolves those baked paths +eagerly, with an existence check that propagates EACCES rather than skipping the +candidate. An engine baked under a build-time ``HOME`` is therefore unreadable to +any other runtime uid, and the process dies during startup before +``PRISMA_QUERY_ENGINE_BINARY`` is ever consulted. + +That is what an OpenShift ``restricted-v2`` namespace produces: the image +``USER`` is ignored and an arbitrary uid in GID 0 is assigned instead. The +symptom is not a degraded proxy, it is a proxy that does not serve at all. + +Booting the image the way that deployment does, and requiring it to answer a +request with a live database connection, is what catches the whole class: +a boot as the default uid, or one that reaches the internet, passes even when +the bake is unusable everywhere it actually ships. + +Gated on LITELLM_IMAGE (the tag of the image to exercise) so it is skipped in +the normal unit-test run and exercised only where an image has been built (the +image-scan workflow). Requires a working docker CLI. +""" + +import json +import shutil +import subprocess +import time +import uuid + +import os +import pytest + +IMAGE = os.getenv("LITELLM_IMAGE") +POSTGRES_IMAGE = os.getenv("LITELLM_TEST_POSTGRES_IMAGE", "postgres:16-alpine") +CURL_IMAGE = os.getenv("LITELLM_TEST_CURL_IMAGE", "curlimages/curl:8.11.1") +COMPONENT_PORT = os.getenv("LITELLM_COMPONENT_PORT", "4000") +NON_ROOT_UID = "12345:0" +STARTUP_TIMEOUT_SECONDS = int(os.getenv("LITELLM_COMPONENT_STARTUP_TIMEOUT", "180")) + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def _docker(*args: str, check: bool = True) -> subprocess.CompletedProcess: + return subprocess.run( + ["docker", *args], capture_output=True, text=True, check=check + ) + + +@pytest.fixture() +def offline_stack(): + """A component container and a fresh Postgres on a network with no egress. + + NON_ROOT_UID is an arbitrary uid in GID 0, the shape OpenShift restricted-v2 + assigns. Postgres and curl are pulled while egress still exists, because the + ``--internal`` network below has none: that is what makes a prisma engine + download (binaries.prisma.sh / npm) fail rather than mask a bake that is not + self-contained. + + The container runs with DISABLE_SCHEMA_UPDATE, since applying the schema is + the migration job's responsibility in this topology and needs the Prisma CLI + these images deliberately omit, and with LITELLM_LOCAL_MODEL_COST_MAP, or the + proxy spends the whole startup budget timing out on a cost-map fetch over the + network it does not have. + + Yields (network_name, component_container). Both are torn down afterwards. + """ + run_id = f"componentserve-{uuid.uuid4().hex[:8]}" + network = f"{run_id}-net" + pg = f"{run_id}-pg" + component = f"{run_id}-app" + + _docker("pull", "--quiet", POSTGRES_IMAGE) + _docker("pull", "--quiet", CURL_IMAGE) + _docker("network", "create", "--internal", network) + try: + _docker( + "run", "-d", "--name", pg, "--network", network, + "-e", "POSTGRES_PASSWORD=pw", "-e", "POSTGRES_DB=litellm", + POSTGRES_IMAGE, + ) + _wait_until_postgres_ready(pg) + assert IMAGE is not None + _docker( + "run", "-d", "--name", component, "--network", network, + "--user", NON_ROOT_UID, + "-e", f"DATABASE_URL=postgresql://postgres:pw@{pg}:5432/litellm", + "-e", "LITELLM_MASTER_KEY=sk-component-serve-test", + "-e", "DISABLE_SCHEMA_UPDATE=true", + "-e", "LITELLM_LOCAL_MODEL_COST_MAP=True", + IMAGE, + ) + yield network, component + finally: + _docker("logs", component, check=False) + _docker("rm", "-f", component, check=False) + _docker("rm", "-f", pg, check=False) + _docker("network", "rm", network, check=False) + + +def _wait_until_postgres_ready(pg: str, attempts: int = 60) -> None: + for _ in range(attempts): + running = _docker( + "ps", "--filter", f"name={pg}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout + if pg not in running: + logs = _docker("logs", pg, check=False) + pytest.fail(f"postgres container is not running:\n{logs.stdout}\n{logs.stderr}") + ready = _docker( + "exec", pg, "pg_isready", "-U", "postgres", "-d", "litellm", check=False + ) + if ready.returncode == 0: + return + time.sleep(1) + pytest.fail(f"postgres never became ready after {attempts}s") + + +def _container_logs(container: str) -> str: + logs = _docker("logs", container, check=False) + return f"stdout:\n{logs.stdout}\nstderr:\n{logs.stderr}" + + +def _is_running(container: str) -> bool: + return bool( + _docker( + "ps", "--filter", f"name={container}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout.strip() + ) + + +def _readiness(network: str, component: str) -> subprocess.CompletedProcess: + """Ask the component for its readiness, from a peer on the same egress-less network.""" + return _docker( + "run", "--rm", "--network", network, CURL_IMAGE, + "--silent", "--max-time", "10", + f"http://{component}:{COMPONENT_PORT}/health/readiness", + check=False, + ) + + +def test_component_serves_offline_as_non_root_uid(offline_stack): + """The component answers a request with a live DB connection, offline, as an arbitrary uid. + + On the pre-fix image this never gets a response: the engine baked under + /home/nonroot (mode 0700, owned by uid 65532) raises + ``PermissionError: .../query-engine-linux-...`` out of pathlib and uvicorn + reports ``Application startup failed. Exiting.``. A bake at the fixed, + world-readable /opt/prisma is what lets any uid start the client. + + `db: connected` is the load-bearing part of the assertion: it means the + query engine binary was found, executed, and reached Postgres. A liveness + probe alone would pass on an image whose engine never resolved. + """ + network, component = offline_stack + + deadline = time.time() + STARTUP_TIMEOUT_SECONDS + probe = None + while time.time() < deadline: + if not _is_running(component): + pytest.fail( + f"the component exited during startup as uid {NON_ROOT_UID} with no egress. " + "The prisma bake is not readable to a uid other than the one that built it, " + "so the proxy does not serve at all.\n" + f"{_container_logs(component)}" + ) + probe = _readiness(network, component) + if probe.returncode == 0 and probe.stdout.strip(): + break + time.sleep(2) + + assert probe is not None and probe.returncode == 0 and probe.stdout.strip(), ( + f"/health/readiness never answered within {STARTUP_TIMEOUT_SECONDS}s as uid " + f"{NON_ROOT_UID} with no egress.\n{_container_logs(component)}" + ) + + payload = json.loads(probe.stdout) + assert payload.get("db") == "connected", ( + f"the component answered but its database is {payload.get('db')!r}, so the baked " + f"query engine did not resolve as uid {NON_ROOT_UID}.\nresponse: {probe.stdout}\n" + f"{_container_logs(component)}" + ) From ae1d1cb05ebc921591c9717acbb5f007f7571805 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 11:05:59 -0700 Subject: [PATCH 24/74] fix(http): stop pooled clients persisting cookies on the aiohttp jar too (#36149) #35978 stopped the pooled A2A client replaying one upstream's Set-Cookie to another by installing a blocking policy on that client's httpx cookie jar. That covers only one of the two jars on the request path. AiohttpTransport is the default transport unless it is explicitly disabled, and the aiohttp ClientSession behind it keeps its own cookie jar which no httpx-level assertion can observe, so the leak is still live on the default path: a live proxy on that commit still delivers agent-alpha's session cookie to agent-beta's card fetch and JSON-RPC call. The reason it looked fixed is that aiohttp's default CookieJar is built with unsafe=False and refuses to store cookies for IP hosts, so a proof addressed to 127.0.0.1 comes back clean whether or not that jar is blocked. Cookie persistence is now blocked where the clients are built rather than at one call site: blocked_cookie_jar() gives every httpx client, async and sync, a jar whose DefaultCookiePolicy(allowed_domains=()) rejects every domain in both directions, and both ClientSession constructions litellm owns, the transport's session factory and the proxy's shared startup session, get a DummyCookieJar. LiteLLM reads a response cookie nowhere, and an explicitly supplied Cookie header still goes out, so passthrough forwarding and an agent's extra_headers are unaffected. The A2A-scoped policy #35978 added is removed, since it is now dead. The two suites that drive the aiohttp session factory synchronously mock ClientSession because a real one needs a running event loop; DummyCookieJar has the same requirement, so they mock it for the same reason. --- litellm/a2a_protocol/main.py | 4 - litellm/llms/custom_httpx/http_handler.py | 15 +++- litellm/proxy/proxy_server.py | 4 +- tests/test_litellm/a2a_protocol/test_main.py | 47 ++++++------ .../test_aiohttp_cleanup_closed.py | 4 +- .../custom_httpx/test_aiohttp_so_keepalive.py | 6 +- .../llms/custom_httpx/test_http_handler.py | 74 +++++++++++++++++++ 7 files changed, 116 insertions(+), 38 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index fe0e6837586..322393cd9c4 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -13,7 +13,6 @@ import asyncio import datetime import uuid from collections.abc import AsyncIterator, Coroutine -from http.cookiejar import DefaultCookiePolicy from typing import TYPE_CHECKING, Any, Final, Optional, cast import litellm @@ -80,8 +79,6 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError # Use our custom resolver instead of the default A2A SDK resolver A2ACardResolver: Final = LiteLLMA2ACardResolver -_BLOCK_ALL_COOKIES: Final = DefaultCookiePolicy(allowed_domains=()) - def _set_usage_on_logging_obj( kwargs: dict[str, Any], @@ -770,7 +767,6 @@ async def create_a2a_client( params={"timeout": timeout}, ) httpx_client: Final = _async_handler.client - httpx_client.cookies.jar.set_policy(_BLOCK_ALL_COOKIES) if extra_headers: verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys())) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index d586156b625..9ada3674d33 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,11 +8,12 @@ import sys import threading import time from collections.abc import Callable, Mapping +from http.cookiejar import CookieJar, DefaultCookiePolicy from typing import TYPE_CHECKING, Any, Final, Optional import certifi import httpx -from aiohttp import ClientSession, TCPConnector +from aiohttp import ClientSession, DummyCookieJar, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport from httpx._types import RequestFiles @@ -144,6 +145,15 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool: return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER +def blocked_cookie_jar() -> CookieJar: + """A jar that stores no response cookie and sends none, for httpx clients. + + LiteLLM's outbound clients are pooled and shared by every caller, so a cookie one + upstream sets would be replayed to every other upstream on a matching domain. + """ + return CookieJar(policy=DefaultCookiePolicy(allowed_domains=())) + + _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS: Final = 5.0 _STREAMING_ERROR_BODY_READ_EXECUTOR: Final = concurrent.futures.ThreadPoolExecutor( max_workers=50, @@ -587,6 +597,7 @@ class AsyncHTTPHandler: verify=ssl_config, cert=cert, headers=default_headers, + cookies=blocked_cookie_jar(), follow_redirects=True, ) @@ -1063,6 +1074,7 @@ class AsyncHTTPHandler: def session_factory() -> ClientSession: return ClientSession( connector=TCPConnector(**transport_connector_kwargs), + cookie_jar=DummyCookieJar(), trust_env=trust_env, ) @@ -1132,6 +1144,7 @@ class HTTPHandler: verify=ssl_config, cert=cert, headers=default_headers, + cookies=blocked_cookie_jar(), follow_redirects=True, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 84f8685f5dc..d016eb57dd4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -871,7 +871,7 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: - from aiohttp import ClientSession, TCPConnector + from aiohttp import ClientSession, DummyCookieJar, TCPConnector from litellm.llms.custom_httpx.http_handler import ( _build_aiohttp_keepalive_socket_factory, @@ -892,7 +892,7 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs["socket_factory"] = socket_factory connector: Final = TCPConnector(**connector_kwargs) - session: Final = ClientSession(connector=connector) + session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar()) verbose_proxy_logger.info( "SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)", diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 59cb8c2c438..08f6b9f25bb 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -174,21 +174,15 @@ _RPC_REPLY = { _AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"} _AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"} -_UPSTREAM_SESSION_COOKIE = "a2a_session=only-agent-a-may-hold-this; Path=/" class _RequestRecorder: - """Records the headers httpx put on the wire, per outbound request. + """Records the headers httpx put on the wire, per outbound request.""" - ``cookie_from_tenant`` makes that tenant's agent answer with a Set-Cookie, standing in - for an upstream that issues a session cookie. - """ - - def __init__(self, cookie_from_tenant: str | None = None): + def __init__(self): self.card_requests = [] self.rpc_requests = [] self.client = None - self.cookie_from_tenant = cookie_from_tenant def __call__(self, request: httpx.Request) -> httpx.Response: headers = {k.lower(): v for k, v in request.headers.items()} @@ -196,8 +190,6 @@ class _RequestRecorder: self.card_requests.append(headers) return httpx.Response(200, json=_AGENT_CARD) self.rpc_requests.append(headers) - if self.cookie_from_tenant is not None and headers.get("x-tenant") == self.cookie_from_tenant: - return httpx.Response(200, json=_RPC_REPLY, headers={"set-cookie": _UPSTREAM_SESSION_COOKIE}) return httpx.Response(200, json=_RPC_REPLY) @@ -205,15 +197,14 @@ def _a2a_client_cache_key(timeout: float) -> str: return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider -async def _seed_shared_a2a_client(cookie_from_tenant: str | None = None) -> _RequestRecorder: +async def _seed_shared_a2a_client() -> _RequestRecorder: """Put the one A2A client the cache will hand out behind a mock transport. Seeding has to happen on the test's own event loop, because the client cache keys on it. The injected client is a real httpx.AsyncClient, so the merge of per-request - headers over client defaults, and httpx's own cookie handling, which is what these - tests are about, stay real. + headers over client defaults, which is what these tests are about, stays real. """ - recorder = _RequestRecorder(cookie_from_tenant=cookie_from_tenant) + recorder = _RequestRecorder() handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT) owned_client = handler.client handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder)) @@ -333,17 +324,21 @@ async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cach @pytest.mark.asyncio -async def test_one_agents_session_cookie_never_reaches_another_agent(isolated_client_cache): - """One pooled client is also one httpx cookie jar. httpx stores every Set-Cookie on the - client and replays it on any later request to a matching domain, so an agent's session - cookie would ride along on a different agent's call to the same host.""" - recorder = await _seed_shared_a2a_client(cookie_from_tenant="tenant-a") +async def test_the_pooled_a2a_client_arrives_with_cookie_persistence_disabled(isolated_client_cache): + """create_a2a_client takes its client from the shared builder rather than building one, + and the builder is what refuses to persist cookies. This pins the join between those + two facts, so the A2A path cannot quietly start acquiring a client that keeps a jar. - client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS) - await _send_message(client_a, _send_request("a")) - client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS) - await _send_message(client_b, _send_request("b")) + test_callers_with_different_headers_reuse_one_pooled_client pins the other half, that + create_a2a_client hands back exactly this cached client.""" + handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.A2AProvider, + params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT}, + ) + request = httpx.Request("GET", "https://agent-a.example.com/") + handler.client.cookies.extract_cookies( + httpx.Response(200, headers={"set-cookie": "SESSION=only-agent-a-may-hold-this"}, request=request) + ) - assert dict(recorder.client.cookies) == {}, "the shared client kept an agent's session cookie" - assert "cookie" not in recorder.card_requests[-1] - assert "cookie" not in recorder.rpc_requests[-1] + assert dict(handler.client.cookies) == {}, "the pooled A2A client kept an upstream's cookie" + await handler.close() diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py index f279acfd60c..82010e82cea 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -13,7 +13,7 @@ def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeyp ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( shared_session=None ) @@ -36,7 +36,7 @@ def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed( ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( shared_session=None ) diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py index 5a37e681c4a..0065bf8f4ef 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py @@ -34,7 +34,7 @@ def test_socket_factory_omitted_when_disabled(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 @@ -55,7 +55,7 @@ def test_socket_factory_attached_when_enabled(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 @@ -77,7 +77,7 @@ def test_socket_factory_skipped_on_old_aiohttp(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 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 86d097c6123..fa1c7308c6f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1171,3 +1171,77 @@ async def test_client_handed_out_by_async_cache_survives_eviction_and_collection assert not consumer_client.is_closed await consumer_client.aclose() + + +_SET_COOKIE = "SESSION=upstream-a-secret; Path=/" + + +def _cookie_recorder(): + """A transport that hands out a Set-Cookie once, and records what comes back.""" + seen = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("cookie")) + if request.url.path == "/set": + return httpx.Response(200, headers={"set-cookie": _SET_COOKIE}) + return httpx.Response(200) + + return handler, seen + + +@pytest.mark.asyncio +async def test_async_client_never_replays_one_upstreams_cookie_to_another(): + """LiteLLM's async clients are pooled and shared by every caller, so a cookie one + upstream sets would be attached to every later request on a matching domain, reaching + a different tenant's upstream. The client must persist no response cookie.""" + handler, seen = _cookie_recorder() + http_handler = AsyncHTTPHandler() + client = http_handler.client + client._transport = httpx.MockTransport(handler) + + await client.get("https://upstream-a.example.com/set") + await client.get("https://upstream-b.example.com/rpc") + await client.aclose() + + assert dict(client.cookies) == {}, "the shared client stored an upstream's cookie" + assert seen == [None, None] + + +def test_sync_client_never_replays_one_upstreams_cookie_to_another(): + """Same invariant on the sync client, which is pooled the same way.""" + handler, seen = _cookie_recorder() + http_handler = HTTPHandler() + client = http_handler.client + client._transport = httpx.MockTransport(handler) + + client.get("https://upstream-a.example.com/set") + client.get("https://upstream-b.example.com/rpc") + client.close() + + assert dict(client.cookies) == {} + assert seen == [None, None] + + +@pytest.mark.asyncio +async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another(): + """The httpx jar is not the only one. AiohttpTransport is litellm's default transport + and the aiohttp ClientSession keeps its own cookie jar, which httpx-level assertions + cannot see, so blocking only the httpx jar leaves the leak intact on the real path. + + aiohttp's default jar refuses cookies for IP hosts, so this drives a hostname. An + IP-addressed check passes whether or not the session jar is blocked.""" + from aiohttp import DummyCookieJar + from yarl import URL + + http_handler = AsyncHTTPHandler(timeout=61.0) + transport = http_handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport), "aiohttp is no longer the default transport" + + session = transport.client() if callable(transport.client) else transport.client + jar = session.cookie_jar + assert isinstance(jar, DummyCookieJar) + + jar.update_cookies({"SESSION": "upstream-a-secret"}, URL("https://upstream-a.example.com")) + assert len(jar) == 0 + assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {} + await session.close() From 330a09235d1a8ca5cbd35acae2adc0e8319f4cf8 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 7 Aug 2026 11:07:09 -0700 Subject: [PATCH 25/74] fix(router): bound fallback-walk work and error-log volume (#36148) --- litellm/constants.py | 1 + litellm/router.py | 14 +- litellm/router_utils/common_utils.py | 17 ++ .../router_utils/fallback_event_handlers.py | 76 ++++++- litellm/types/utils.py | 1 + .../test_fallback_event_handlers.py | 204 ++++++++++++++++++ .../test_router_utils_common_utils.py | 26 +++ tests/test_litellm/test_router.py | 135 ++++++++++++ 8 files changed, 469 insertions(+), 5 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 75f12190b3e..6f0e9e7afe2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_non DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) +ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) diff --git a/litellm/router.py b/litellm/router.py index 1edb80da7ce..1a76eb5d59b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -109,6 +109,7 @@ from litellm.router_utils.common_utils import ( filter_team_based_models, filter_web_search_deployments, resolve_model_group_alias, + truncate_fallback_error_detail, ) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( @@ -342,6 +343,12 @@ def _replay_live_router_model_cost() -> None: set_live_deployment_replay(_replay_live_router_model_cost) +# Kwargs that log_retry must not copy into a retry breadcrumb. The breadcrumbs reach spend +# logs and logging callbacks, and these carry either the request payload or router-internal +# walk state rather than anything that identifies the failed attempt. +RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(("messages", "original_function", "attempted_targets")) + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -6361,17 +6368,16 @@ class Router: return response except Exception as new_exception: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) - fallback_failure_exception_str = redact_string(str(new_exception)) + fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception))) cooldown_info: Final = await _async_get_cooldown_deployments_with_debug_info( litellm_router_instance=self, parent_otel_span=parent_otel_span, ) verbose_router_logger.error( "litellm.router.py::async_function_with_fallbacks() - " - "Error occurred while trying to do fallbacks - %s\n%s\n" + "Error occurred while trying to do fallbacks - %s\n" "Debug Information:\nCooldown Deployments=%s", fallback_failure_exception_str, - redact_string(traceback.format_exc()), cooldown_info, ) @@ -7162,7 +7168,7 @@ class Router: k, v, ) in kwargs.items(): # log everything in kwargs except the old previous_models value - prevent nesting - if k not in [_metadata_var, "messages", "original_function"]: + if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS: previous_model[k] = v elif k == _metadata_var and isinstance(v, dict): previous_model[_metadata_var] = {} diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 1c628ef9ee3..6fad2dd31e9 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject from litellm._logging import verbose_logger +from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError from litellm.types.router import CredentialLiteLLMParams @@ -43,6 +44,22 @@ def resolve_model_group_alias(model_group_alias: object, model: str) -> str | No return target +def truncate_fallback_error_detail(detail: str) -> str: + """ + Bound a fallback failure detail before it is logged or appended to an exception message. + + Each level of the fallback walk records the failure of the level below it, so an + untruncated detail carries every nested failure with it and grows superlinearly with + the number of attempted model groups. One deterministic pre-network failure walked + through a small fallback graph is enough to turn that into hundreds of megabytes of + output on the event-loop thread, which starves the process that produced it. + """ + if len(detail) <= ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: + return detail + dropped: Final = len(detail) - ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + return f"{detail[:ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS]}... [truncated {dropped} characters]" + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 1c6bb52ccb8..ef48ccc821e 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,3 +1,6 @@ +import hashlib +import json +from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, Final @@ -19,6 +22,52 @@ else: LitellmRouter = Any +def fallback_attempt_key(fallback_target: object) -> str | None: + """ + Identity of one fallback attempt, so the same attempt is never made twice per request. + + A bare model group name and a `{"model": name}` entry describe the same attempt. An + entry carrying anything else describes a different one and keeps its own identity: a + client-side fallback list overrides request params such as `messages`, and the router + re-targets the group that just failed by attaching `_target_order` or + `_excluded_deployment_ids` to select a different set of deployments inside it. The + payload is hashed rather than kept, so a large `messages` override does not make the + request hold a second copy of itself. + + Returns None for a shape with no usable identity, which is never skipped. + """ + if isinstance(fallback_target, str): + return fallback_target + if not isinstance(fallback_target, dict): + return None + model: Final = fallback_target.get("model") + if tuple(fallback_target) == ("model",) and isinstance(model, str): + return model + serialized: Final = json.dumps(fallback_target, sort_keys=True, default=str) + return hashlib.sha256(serialized.encode()).hexdigest() + + +@dataclass(slots=True) +class AttemptedFallbackTargets: + """ + The fallback attempts a single request has already made. + + One instance is created on the first fallback hop and shared by reference for the rest + of the walk, so an attempt made in one branch is not repeated in a sibling branch. + Without it the walk enumerates paths rather than attempts: a fallback graph containing + a cycle retries one deterministic failure once per path through the cycle, and a + client-side fallback list is re-walked at every level of the recursion. + """ + + keys: frozenset[str] = frozenset() + + def __contains__(self, key: str) -> bool: + return key in self.keys + + def record(self, key: str) -> None: + self.keys = self.keys | frozenset((key,)) + + def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: """ Handles wildcard routing scenario @@ -106,7 +155,14 @@ async def run_async_fallback( fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"] original_model_group: The original model group. example: "gpt-3.5-turbo" original_exception: The original exception. - **kwargs: Keyword arguments. + **kwargs: Keyword arguments. `attempted_targets` carries the fallback attempts + already made for this request, created on the first hop and shared by reference + for the rest of the walk. A target already in it is skipped, so neither a + fallback graph that loops back on itself nor a client-side fallback list + re-walked at each level can repeat an attempt that has already failed. Identity + comes from `fallback_attempt_key`, so an entry that overrides request params or + re-targets the failed group with a different deployment selection stays distinct + from a bare name. Returns: The response from the successful fallback model group. @@ -120,10 +176,27 @@ async def run_async_fallback( error_from_fallbacks = original_exception fallback_errors = (get_fallback_error_info(original_exception),) + # Read out of kwargs and narrowed here rather than declared as a parameter: every caller + # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter + # would carry an annotation that no call site can actually be checked against. + carried_targets: Final = kwargs.get("attempted_targets") + attempted: Final = ( + carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() + ) + attempted.record(original_model_group) for mg in fallback_model_group: if mg == original_model_group: continue + attempt_key = fallback_attempt_key(mg) + if attempt_key is not None: + if attempt_key in attempted: + verbose_router_logger.info( + "Skipping fallback to model_group = %s, already attempted for this request", + mask_sensitive_structure(mg), + ) + continue + attempted.record(attempt_key) try: # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) @@ -138,6 +211,7 @@ async def run_async_fallback( fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks + kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 329b10e72e7..35d4250782f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3479,6 +3479,7 @@ all_litellm_params = ( "user_continue_message", "fallback_depth", "max_fallbacks", + "attempted_targets", "max_budget", "budget_duration", "use_in_pass_through", diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 98a34de295c..e2348d28701 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -3,6 +3,8 @@ import json import pytest from litellm.router_utils.fallback_event_handlers import ( + AttemptedFallbackTargets, + fallback_attempt_key, get_fallback_model_group, run_async_fallback, ) @@ -142,6 +144,208 @@ async def test_run_async_fallback_skips_original_model_group(): assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 +class RecordingFailRouter: + def __init__(self): + self.attempted_models = [] + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.attempted_models.append(kwargs.get("model")) + raise RuntimeError("fallback model also failed") + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_model_group_already_attempted(): + """A fallback graph that loops back on itself must not re-attempt a model group that + already failed for this request. Every group in a cycle fails identically, so + revisiting one multiplies the work and the error output without any chance of + succeeding.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="original failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["already-attempted"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_attempts_a_repeated_target_once(): + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model", "fallback-model", "other-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=5, + fallback_depth=0, + ) + + assert router.attempted_models == ["fallback-model", "other-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call(): + """The nested call is where the next hop of the walk decides what to skip, so the + accumulated set has to reach it, carrying both the group that just failed and the + target being attempted.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"earlier-model"})), + ) + + assert router.received_kwargs["attempted_targets"].keys == frozenset( + {"earlier-model", "primary-model", "fallback-model"} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entry", + [ + {"model": "primary-model", "_target_order": 2}, + {"model": "primary-model", "_excluded_deployment_ids": ["dep-1"]}, + ], +) +async def test_run_async_fallback_still_retargets_the_same_group_via_dict_entry(entry): + """Order-based fallback and weighted intra-group failover both re-target the group that + just failed, selecting a different set of deployments inside it. Those entries are dicts + rather than plain names and must survive a guard that skips already-attempted names.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[entry], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"primary-model"})), + ) + + assert router.received_kwargs["model"] == "primary-model" + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_a_repeated_dict_target(): + """A client-side fallback list names its targets with dicts, and that list is re-walked + at every level of the recursion, so an entry that carries no request override has to be + recognised as the same attempt as the bare name.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="original failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "already-attempted"}], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_attempts_a_repeated_dict_target_once(): + router = RecordingFailRouter() + entry = {"model": "fallback-model", "messages": [{"role": "user", "content": "shorter"}]} + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[entry, entry, {"model": "other-model"}], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=5, + fallback_depth=0, + ) + + assert router.attempted_models == ["fallback-model", "other-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bare_name(): + """The documented use of the client-side form is to retry a group with different request + params, so an entry carrying an override must survive even when the bare name of that + same group has already been attempted.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[ + {"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]} + ], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == ["already-attempted"] + + +@pytest.mark.parametrize( + "target, expected", + [ + ("group-a", "group-a"), + ({"model": "group-a"}, "group-a"), + (None, None), + (["group-a"], None), + ], +) +def test_fallback_attempt_key_identity(target, expected): + """A bare name and a `{"model": name}` entry are the same attempt. A shape with no + usable identity returns None and is never skipped, so an unrecognised entry keeps + today's behaviour rather than being silently dropped.""" + assert fallback_attempt_key(target) == expected + + +def test_fallback_attempt_key_gives_a_param_only_entry_its_own_identity(): + """An entry with no `model` re-targets the group currently being attempted with + different request params, so it is a distinct attempt and still needs an identity.""" + key = fallback_attempt_key({"messages": [{"role": "user", "content": "shorter"}]}) + + assert key is not None + assert key != fallback_attempt_key({"messages": [{"role": "user", "content": "other"}]}) + + +def test_fallback_attempt_key_separates_overrides_from_the_bare_name(): + bare = fallback_attempt_key("group-a") + override = fallback_attempt_key({"model": "group-a", "messages": [{"role": "user", "content": "x"}]}) + other_override = fallback_attempt_key({"model": "group-a", "messages": [{"role": "user", "content": "y"}]}) + order_retarget = fallback_attempt_key({"model": "group-a", "_target_order": 2}) + + assert len({bare, override, other_override, order_retarget}) == 4 + + +def test_fallback_attempt_key_is_stable_across_key_order(): + assert fallback_attempt_key({"model": "group-a", "_target_order": 2}) == fallback_attempt_key( + {"_target_order": 2, "model": "group-a"} + ) + + def test_get_fallback_model_group_does_not_mutate_fallbacks(): """A string fallback must be resolved without mutating the caller's fallbacks list, which is the live router config shared across requests.""" diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 7d453c72652..0d063ad14f5 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -4,6 +4,7 @@ from unittest.mock import Mock import pytest from litellm import Router +from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.common_utils import ( _deployment_supports_web_search, @@ -11,6 +12,7 @@ from litellm.router_utils.common_utils import ( filter_team_based_models, filter_web_search_deployments, resolve_model_group_alias, + truncate_fallback_error_detail, ) @@ -558,3 +560,27 @@ class TestResolveModelGroupAlias: assert router._get_model_from_alias("group-a") == "group-b" assert router._get_model_from_alias("group-item") == "group-b" assert router._get_model_from_alias("group-b") is None + + +class TestTruncateFallbackErrorDetail: + def test_short_detail_is_returned_unchanged(self): + assert truncate_fallback_error_detail("boom") == "boom" + + def test_detail_at_the_limit_is_returned_unchanged(self): + detail = "x" * ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + assert truncate_fallback_error_detail(detail) == detail + + def test_long_detail_is_bounded_and_reports_what_was_dropped(self): + detail = "x" * (ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + 500) + + truncated = truncate_fallback_error_detail(detail) + + assert truncated.startswith("x" * ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS) + assert truncated.endswith("... [truncated 500 characters]") + assert len(truncated) < len(detail) + + def test_a_megabyte_of_detail_comes_back_small(self): + """The detail is what a fallback level records about the level below it, so it has + to stay small enough that a walk over many model groups cannot compound it into an + output volume that starves the process.""" + assert len(truncate_fallback_error_detail("x" * 1_000_000)) < 3_000 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4a3395a7d3f..da2d78edb73 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import logging import os import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +15,7 @@ sys.path.insert( import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -7415,3 +7417,136 @@ class TestAutoRouterMaxInputCharsWiring: router = self._router() assert self._registered_auto_router(router).max_input_chars == DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS + + +class _LogCapture(logging.Handler): + def __init__(self, level): + super().__init__(level=level) + self._level = level + self.messages = [] + + def emit(self, record): + if record.levelno == self._level: + self.messages.append(record.getMessage()) + + +class _FallbackAttemptRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.failed_targets = [] + + async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception): + self.failed_targets.append(kwargs.get("model")) + + +def _cyclic_fallback_router(num_retries=0): + groups = ["group-a", "group-b", "group-c", "group-d"] + return litellm.Router( + model_list=[ + { + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "mock_response": "litellm.InternalServerError", + }, + } + for group in groups + ], + fallbacks=[ + {"group-a": ["group-b", "group-c"]}, + {"group-b": ["group-a", "group-c"]}, + {"group-c": ["group-d"]}, + {"group-d": ["group-b", "group-a"]}, + ], + num_retries=num_retries, + ) + + +async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwargs): + router_logger = logging.getLogger("LiteLLM Router") + previous_level = router_logger.level + router_logger.setLevel(capture.level) + router_logger.addHandler(capture) + if recorder is not None: + litellm.callbacks.append(recorder) + try: + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs + ) + finally: + router_logger.removeHandler(capture) + router_logger.setLevel(previous_level) + if recorder is not None: + litellm.callbacks.remove(recorder) + + +@pytest.mark.asyncio +async def test_cyclic_fallback_graph_does_not_amplify_one_request(): + """A fallback graph whose entries loop back on each other is easy to build by accident, + and every group in the loop fails identically on a deterministic error, so the walk must + not revisit a group and must not re-emit a growing chained traceback at each level. Left + unbounded, one request blocks the event loop long enough for health probes to fail.""" + recorder = _FallbackAttemptRecorder() + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(_cyclic_fallback_router(), capture, recorder) + + assert sorted(set(recorder.failed_targets)) == ["group-b", "group-c", "group-d"] + assert len(recorder.failed_targets) == len(set(recorder.failed_targets)) + assert not any("Traceback (most recent call last)" in message for message in capture.messages) + assert sum(len(message) for message in capture.messages) < 5_000 + + +@pytest.mark.asyncio +async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): + """log_retry copies every kwarg into previous_models, which reaches spend logs and + logging callbacks. The set of already-attempted groups is router-internal walk state + with no diagnostic value there, and it is the one entry that is not a plain scalar. + A retry has to be configured for the walk state to reach log_retry at all.""" + router = _cyclic_fallback_router(num_retries=1) + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(router, capture) + + assert router.previous_models, "no retry breadcrumbs were recorded" + assert any( + "fallback_depth" in breadcrumb for breadcrumb in router.previous_models + ), "no breadcrumb carried router walk state, so this test cannot see the leak" + for breadcrumb in router.previous_models: + assert "attempted_targets" not in breadcrumb + + +@pytest.mark.asyncio +async def test_fallback_traceback_stays_available_at_debug_level(): + """Dropping the stack from the ERROR line is only safe because the fallback path still + emits it once per level at DEBUG, which is what an operator needs to diagnose why every + fallback failed. This pins that remaining debug traceback.""" + capture = _LogCapture(logging.DEBUG) + + await _drive_cyclic_fallback(_cyclic_fallback_router(), capture) + + assert any("Traceback (most recent call last)" in message for message in capture.messages) + + +@pytest.mark.asyncio +async def test_fallback_failure_detail_from_upstream_is_bounded(): + """The detail each level records about the level below it is attacker-influenced, since + it carries whatever the upstream error said. It has to be bounded on its own, so a walk + over several groups cannot compound one large message into the log or into the message + handed back to the caller.""" + huge_message = "z" * 50_000 + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback( + _cyclic_fallback_router(), + capture, + mock_response=litellm.InternalServerError( + message=huge_message, llm_provider="openai", model="group-a" + ), + ) + + assert capture.messages, "the fallback failure path did not log at ERROR" + assert huge_message not in "".join(capture.messages) + assert max(len(message) for message in capture.messages) < 5_000 From d30e933e33ff2d98076431c44b668e5fbefc6b65 Mon Sep 17 00:00:00 2001 From: "cursor[bot]" <206951365+cursor[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:12:08 -0700 Subject: [PATCH 26/74] ci: wire credential_endpoints tests into the proxy endpoints job (#36187) PR #36166 added tests/test_litellm/proxy/credential_endpoints/test_endpoints.py but no CI job invokes it, so the CI Coverage guard failed on litellm_internal_staging. Add the directory to the proxy-endpoints job's test-path list so pytest actually runs the new tests and the coverage assertion is satisfied. Co-authored-by: Cursor Agent Co-authored-by: Krrish Dholakia --- .github/workflows/test-unit-proxy-endpoints.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 20b89c72440..645996f779d 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -44,6 +44,7 @@ jobs: tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/a2a + tests/test_litellm/proxy/credential_endpoints tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/shutdown From 0cd58a04d75419f94ec372a2e207f1dc2ea2659b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 11:13:03 -0700 Subject: [PATCH 27/74] docs(keys): document /key/info fields and clarify budget_reset_at is the next reset (#36127) * docs(keys): document /key/info fields and clarify budget_reset_at is the next reset * docs(keys): drop wrong budget window start math, note reset boundary alignment --- .../key_management_endpoints.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a5078c50fc0..f2cb1124fa0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3528,11 +3528,34 @@ async def info_key_fn( ): """ Retrieve information about a key. + Parameters: - key: Optional[str] = Query parameter representing the key in the request - user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + - key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash. + Defaults to the key in the Authorization header. + Returns: - Dict containing the key and its associated information + - key: str - The key that was looked up, echoed back as it was passed in + - info: dict - The key's row, minus the hashed token + - key_alias: str | None - User-friendly key alias + - spend: float - Amount spent by the key. When budget_duration is set this covers only the + current budget window, not the key's lifetime + - max_budget: float | None - Max budget for the key, enforced against spend + - budget_duration: str | None - Budget reset period ("30d", "1h", etc.) + - budget_reset_at: datetime | None - When the current budget window ends and spend is next + reset to 0, not when it was last reset. Reset times snap to standard boundaries in the + configured timezone (30d and 1mo land on the 1st of the month, 7d on Monday, 1h on the + hour), so subtracting budget_duration from it does not give the window's start + - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + - model_max_budget_usage: dict | None - Current-window spend per model, present only when + the key has per-model budgets + - models: list - Model_name's the key is allowed to call + - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits + - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} + - blocked: bool | None - Whether the key is blocked + - expires: datetime | None - When the key stops authenticating requests + - last_active: datetime | None - When the key was last used + - object_permission: dict | None - Resolved vector store / MCP permissions when the key has + an object_permission_id Example Curl: ``` From bd289c151c12a60edd6b0bb46adb224f15104ce0 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 7 Aug 2026 11:14:20 -0700 Subject: [PATCH 28/74] fix(azure_sentinel): add AZURE_SENTINEL_AUTHORITY_HOST as a Sentinel scoped override (#36165) Making Sentinel follow AZURE_AUTHORITY_HOST is a breaking change for a deployment that sets that variable for Azure OpenAI or the azure_storage callback while keeping a commercial Sentinel workspace. That deployment had no opt-out, because the proxy constructs the logger with no arguments and the authority_host parameter is reachable only from the SDK. Resolve the authority from AZURE_SENTINEL_AUTHORITY_HOST before falling back to AZURE_AUTHORITY_HOST, matching how tenant id, client id and client secret already resolve in this constructor. --- .../azure_sentinel/azure_sentinel.py | 9 +++-- .../integrations/test_azure_sentinel.py | 36 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 563f815b582..24328549094 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -78,8 +78,8 @@ class AzureSentinelLogger(CustomBatchLogger): If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name. authority_host (str, optional): Microsoft Entra authority host that issues the OAuth2 token, e.g. "https://login.microsoftonline.us" for Azure Government. If not provided, will use - AZURE_AUTHORITY_HOST env var or default to the Azure Public Cloud authority. The Azure - Monitor audience is derived from it. + AZURE_SENTINEL_AUTHORITY_HOST or AZURE_AUTHORITY_HOST env vars, or default to the Azure + Public Cloud authority. The Azure Monitor audience is derived from it. """ self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) @@ -95,7 +95,10 @@ class AzureSentinelLogger(CustomBatchLogger): client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) resolved_authority_host: Final = self._normalize_authority_host( - authority_host or os.getenv("AZURE_AUTHORITY_HOST") or DEFAULT_AZURE_AUTHORITY_HOST + authority_host + or os.getenv("AZURE_SENTINEL_AUTHORITY_HOST") + or os.getenv("AZURE_AUTHORITY_HOST") + or DEFAULT_AZURE_AUTHORITY_HOST ) if not resolved_dcr_immutable_id: diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 56662eea633..f48f5cb1784 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -313,6 +313,7 @@ def _build_logger(**overrides): @pytest.fixture def _no_authority_host_env(monkeypatch): + monkeypatch.delenv("AZURE_SENTINEL_AUTHORITY_HOST", raising=False) monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) @@ -389,3 +390,38 @@ async def test_azure_sentinel_token_request_uses_sovereign_authority_and_audienc assert len(token_calls) == 1 assert token_calls[0].kwargs["url"] == "https://login.microsoftonline.us/test-tenant-id/oauth2/v2.0/token" assert token_calls[0].kwargs["data"]["scope"] == "https://monitor.azure.us/.default" + + +def test_azure_sentinel_authority_host_prefers_the_sentinel_scoped_env_var(_no_authority_host_env, monkeypatch): + """AZURE_AUTHORITY_HOST is shared with Azure OpenAI and the azure_storage callback, so a deployment + whose Sentinel workspace lives in a different cloud than the rest of its Azure resources needs a + Sentinel-scoped override. This mirrors how tenant, client id and secret already resolve.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com") + monkeypatch.setenv("AZURE_SENTINEL_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +def test_azure_sentinel_falls_back_to_the_shared_authority_host(_no_authority_host_env, monkeypatch): + """With no Sentinel-scoped override the shared variable still applies, which is the behavior + shipped in the original fix.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +def test_azure_sentinel_authority_host_argument_outranks_the_scoped_env_var(_no_authority_host_env, monkeypatch): + """An explicit constructor argument is the most specific source and has to win, otherwise a + deployment that exports the scoped variable silently overrides an SDK caller.""" + monkeypatch.setenv("AZURE_SENTINEL_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger(authority_host="https://login.microsoftonline.com") + + assert logger.authority_host == "https://login.microsoftonline.com" + assert logger.oauth_scope == "https://monitor.azure.com/.default" From bf2def7eb565c2f8cdaeb4e42df19508e54a2ff9 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:32:30 -0700 Subject: [PATCH 29/74] docs(pr-template): add a User Flow section with authoring instructions (#36162) Adds a User Flow section right below the TLDR so every PR describes the same end user doing the same task before and after the change, plus comment instructions and a worked example so contributors can write it without any local tooling. --- .github/pull_request_template.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 85291b49880..d34b0ee2e0f 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -13,6 +13,33 @@ How it solves it: - - ... +## User Flow + + + ## Relevant issues From 493eb9054e906628f2d4eba0a17453576a94e6c9 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:13:12 +0000 Subject: [PATCH 30/74] chore(ui): regenerate schema.d.ts for the /key/info docstring update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7675c0506c0..f1660e77ad9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -6904,11 +6904,34 @@ export interface paths { /** * Info Key Fn * @description Retrieve information about a key. + * * Parameters: - * key: Optional[str] = Query parameter representing the key in the request - * user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + * - key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash. + * Defaults to the key in the Authorization header. + * * Returns: - * Dict containing the key and its associated information + * - key: str - The key that was looked up, echoed back as it was passed in + * - info: dict - The key's row, minus the hashed token + * - key_alias: str | None - User-friendly key alias + * - spend: float - Amount spent by the key. When budget_duration is set this covers only the + * current budget window, not the key's lifetime + * - max_budget: float | None - Max budget for the key, enforced against spend + * - budget_duration: str | None - Budget reset period ("30d", "1h", etc.) + * - budget_reset_at: datetime | None - When the current budget window ends and spend is next + * reset to 0, not when it was last reset. Reset times snap to standard boundaries in the + * configured timezone (30d and 1mo land on the 1st of the month, 7d on Monday, 1h on the + * hour), so subtracting budget_duration from it does not give the window's start + * - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + * - model_max_budget_usage: dict | None - Current-window spend per model, present only when + * the key has per-model budgets + * - models: list - Model_name's the key is allowed to call + * - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits + * - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} + * - blocked: bool | None - Whether the key is blocked + * - expires: datetime | None - When the key stops authenticating requests + * - last_active: datetime | None - When the key was last used + * - object_permission: dict | None - Resolved vector store / MCP permissions when the key has + * an object_permission_id * * Example Curl: * ``` From eb3c8c168f533640bf16d70fdeea148e18f65de0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 12:19:49 -0700 Subject: [PATCH 31/74] fix(proxy): derive config agent ids from agent_name so grants survive secret rotation (#36020) * fix(proxy): derive config agent ids from agent_name so grants survive secret rotation Config-defined A2A agents were identified by a sha256 of the whole resolved config entry, secrets included, so rotating an os.environ secret re-minted the agent_id on restart and orphaned every object_permission.agents grant while grant-less keys kept access (LIT-5144). The id now hashes only agent_name, and the old full-entry hash is kept as a legacy alias: permission checks, GET /v1/agents filtering, spend and key attachment, and public_agent_groups all normalize legacy ids so pre-upgrade grants keep working * fix(proxy): persist stable agent ids into stored grants at startup The runtime alias only translates a legacy grant while the current config still hashes to it, so a secret rotation after upgrading would orphan the grant, and an orphaned grant intersecting a stable team grant collapses to an empty list that downstream reads as allow-all. Rewriting the stored ids once at boot removes both. This cannot be a SQL migration because only the running proxy can recompute the legacy hash from resolved config secrets * fix(proxy): make the grant id migration a compare-and-swap A grant edited between the migration's read and write kept the stale snapshot. The update now predicates on the agents array read at scan time via update_many, so a concurrently modified row is skipped and the runtime alias covers it until the next boot retries * fix(proxy): retry the grant id migration and stay within the LIT002 ceiling The one-shot startup task now retries up to three times with a short delay so a transient DB error at boot cannot leave a legacy grant unmigrated until an operator's next restart is the rotation itself. The new list constructions in the migration and the alias-expanded agent id lookups are tuples now, keeping the branch under the mutable-collection budget * fix(proxy): count compare-and-swap misses in the grant id migration migrate_legacy_grant_ids now returns rewritten and missed counts from the update_many results instead of reporting scanned rows as migrated, and the startup task retries while any rows remain unmigrated, not just on errors * fix(lint): clear basedpyright budget breaches in agent id aliasing --- .../proxy/agent_endpoints/agent_registry.py | 119 ++++++++++++-- .../auth/agent_permission_handler.py | 38 +++-- litellm/proxy/agent_endpoints/endpoints.py | 34 +++- litellm/proxy/proxy_server.py | 31 ++++ .../public_endpoints/public_endpoints.py | 2 +- litellm/repositories/table_repositories.py | 4 + ruff-strict-budget.json | 2 +- .../auth/test_agent_permission_handler.py | 80 +++++++++ .../agent_endpoints/test_agent_registry.py | 155 +++++++++++++++++- .../proxy/agent_endpoints/test_endpoints.py | 3 +- .../public_endpoints/test_public_endpoints.py | 2 + type-discipline-budget.json | 4 +- 12 files changed, 426 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 476bd725c73..742fdf35b1e 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,8 +1,10 @@ +import asyncio import hashlib import json from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Protocol, TypedDict +from types import MappingProxyType +from typing import Any, Final, NamedTuple, Protocol, TypedDict import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -10,7 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient -from litellm.repositories.table_repositories import AgentsRepository +from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -86,10 +88,32 @@ def agents_table(prisma_client: PrismaClient) -> AgentTableClient: return table +class ObjectPermissionGrantRecord(Protocol): + object_permission_id: str + agents: list[str] | None + + +class ObjectPermissionTableClient(Protocol): + async def find_many(self, where: Mapping[str, object]) -> Sequence[ObjectPermissionGrantRecord]: ... + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient: + table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table + return table + + +class GrantMigrationResult(NamedTuple): + rewritten: int + missed: int + + class AgentRegistry: def __init__(self): self.agent_list: list[AgentResponse] = [] self.config_agents: tuple[AgentConfig, ...] = () + self.config_agent_legacy_ids: Mapping[str, str] = MappingProxyType({}) def reset_agent_list(self): self.agent_list = [] @@ -100,23 +124,33 @@ class AgentRegistry: def deregister_agent(self, agent_name: str): self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name] - def get_agent_list(self, agent_names: Sequence[str] | None = None): + def get_agent_list(self, agent_names: Sequence[str] | None = None) -> tuple[AgentResponse, ...]: if agent_names is not None: - return [agent for agent in self.agent_list if agent.agent_name in agent_names] - return self.agent_list + return tuple(agent for agent in self.agent_list if agent.agent_name in agent_names) + return tuple(self.agent_list) - def get_public_agent_list(self) -> list[AgentResponse]: - public_agent_list: Final[list[AgentResponse]] = [] - if litellm.public_agent_groups is None: - return public_agent_list - for agent in self.agent_list: - if agent.agent_id in litellm.public_agent_groups: - public_agent_list.append(agent) - return public_agent_list + def get_public_agent_list(self) -> tuple[AgentResponse, ...]: + public_agent_groups: Final = litellm.public_agent_groups + if public_agent_groups is None: + return () + return tuple( + agent for agent in self.agent_list if not self.ids_for_agent(agent.agent_id).isdisjoint(public_agent_groups) + ) def _create_agent_id(self, agent_config: AgentConfig) -> str: + return hashlib.sha256(agent_config["agent_name"].encode()).hexdigest() + + def _create_legacy_agent_id(self, agent_config: AgentConfig) -> str: return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() + def ids_for_agent(self, agent_id: str) -> frozenset[str]: + return frozenset( + {agent_id, *(legacy for legacy, stable in self.config_agent_legacy_ids.items() if stable == agent_id)} + ) + + def stable_agent_id(self, agent_id: str) -> str: + return self.config_agent_legacy_ids.get(agent_id, agent_id) + def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): """ Register the agents declared in config.yaml and remember them for later rebuilds. @@ -131,12 +165,20 @@ class AgentRegistry: if agent_config is None: return - self.config_agents = tuple(agent_config) - for agent_config_item in agent_config: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") + self.config_agents = tuple(agent_config) + self.config_agent_legacy_ids = MappingProxyType( + { + self._create_legacy_agent_id(agent_config_item): self._create_agent_id(agent_config_item) + for agent_config_item in agent_config + if agent_config_item.get("agent_name") and agent_config_item.get("agent_card_params") + } + ) + + for agent_config_item in agent_config: agent_name = agent_config_item.get("agent_name") agent_card_params = agent_config_item.get("agent_card_params") if not all([agent_name, agent_card_params]): @@ -180,6 +222,45 @@ class AgentRegistry: self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents) return self.agent_list + async def migrate_legacy_grant_ids(self, table: ObjectPermissionTableClient) -> GrantMigrationResult: + """ + Rewrite object_permission.agents rows holding a legacy full-entry hash to the + stable name-derived id. + + Only the running proxy can do this: the legacy hash is computed from the + resolved config entry (secrets included), so no SQL migration can know it. + Persisting the stable id here is what keeps a grant alive across a later + secret rotation, which re-mints the legacy hash and would otherwise orphan + the stored value. Idempotent; runs of it after the first find no rows. + + Each write is a compare-and-swap against the agents array read above, so a + grant edited concurrently is left untouched; the runtime alias keeps covering + it and the next boot retries the rewrite. + """ + legacy_ids: Final = tuple(legacy for legacy, stable in self.config_agent_legacy_ids.items() if legacy != stable) + if not legacy_ids: + return GrantMigrationResult(rewritten=0, missed=0) + rows: Final = await table.find_many(where={"agents": {"has_some": legacy_ids}}) + updates: Final = tuple( + ( + row.object_permission_id, + tuple(row.agents or ()), + tuple(dict.fromkeys(self.stable_agent_id(agent_id) for agent_id in row.agents or ())), + ) + for row in rows + ) + counts: Final = await asyncio.gather( + *( + table.update_many( + where={"object_permission_id": object_permission_id, "agents": {"equals": snapshot_agents}}, + data={"agents": translated_agents}, + ) + for object_permission_id, snapshot_agents, translated_agents in updates + ) + ) + rewritten: Final = sum(counts) + return GrantMigrationResult(rewritten=rewritten, missed=len(updates) - rewritten) + ########################################################### ########### DB management helpers for agents ########### ############################################################ @@ -492,6 +573,14 @@ class AgentRegistry: if agent.agent_id == agent_id: return agent + translated_id: Final = self.config_agent_legacy_ids.get(agent_id) + if translated_id is None: + return None + + for agent in self.agent_list: + if agent.agent_id == translated_id: + return agent + return None except Exception as e: raise Exception(f"Error getting agent from DB: {e}") diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index af7e731b385..6baeab3dfff 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -42,24 +42,25 @@ class AgentRequestHandler: List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). """ try: - allowed_agents: list[str] = [] - allowed_agents_for_key: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - allowed_agents_for_team: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + raw_key_grants: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) + raw_team_grants: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + allowed_agents_for_key: Final = frozenset( + global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_key_grants + ) + allowed_agents_for_team: Final = frozenset( + global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_team_grants + ) # If team has agent restrictions, handle inheritance and intersection logic - if len(allowed_agents_for_team) > 0: - if len(allowed_agents_for_key) > 0: - # Key has its own agent permissions - use intersection with team permissions - for agent_id in allowed_agents_for_key: - if agent_id in allowed_agents_for_team: - allowed_agents.append(agent_id) - else: - # Key has no agent permissions - inherit from team - allowed_agents = allowed_agents_for_team - else: - allowed_agents = allowed_agents_for_key - - return list(set(allowed_agents)) + if allowed_agents_for_team and allowed_agents_for_key: + # Key has its own agent permissions - use intersection with team permissions + return sorted(allowed_agents_for_key & allowed_agents_for_team) + if allowed_agents_for_team: + # Key has no agent permissions - inherit from team + return sorted(allowed_agents_for_team) + return sorted(allowed_agents_for_key) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) return [] @@ -79,13 +80,16 @@ class AgentRequestHandler: Returns: bool: True if agent is allowed, False otherwise """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + allowed_agents: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth) # Empty list means no restrictions - allow all if len(allowed_agents) == 0: return True - return agent_id in allowed_agents + stable_id: Final = global_agent_registry.stable_agent_id(agent_id) + return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agents) @staticmethod def _get_key_object_permission( diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 1f9c6e1cc05..028885aa065 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -101,7 +101,11 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) foreign key. Mirrors how spend is joined into the agent response so the UI never has to cross-reference a full key dump client-side. Only non-secret fields are exposed (alias, masked key_name, hashed token).""" - agent_ids: Final = [agent.agent_id for agent in agents] + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent_ids: Final = tuple( + alias_id for agent in agents for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + ) if not agent_ids: return key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many( @@ -117,7 +121,12 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) ) ) for agent in agents: - agent.keys = keys_by_agent.get(agent.agent_id) + matched_keys = [ + key_summary + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + for key_summary in keys_by_agent.get(alias_id) or () + ] + agent.keys = matched_keys or None def _redact_sensitive_agent_fields( @@ -266,23 +275,32 @@ async def get_agents( from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - agent_ids: Final = [agent.agent_id for agent in returned_agents] + agent_ids: Final = tuple( + alias_id + for agent in returned_agents + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + ) if agent_ids: db_agents: Final = await agents_table(prisma_client).find_many( where={"agent_id": {"in": agent_ids}}, ) spend_map: Final = {a.agent_id: a.spend for a in db_agents} for agent in returned_agents: - if agent.agent_id in spend_map: - agent.spend = spend_map[agent.agent_id] + matched_spends = tuple( + spend_map[alias_id] + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + if alias_id in spend_map + ) + if matched_spends: + agent.spend = sum(matched_spends) await _attach_keys_to_agents(returned_agents, prisma_client) # add is_public field to each agent - we do it this way, to allow setting config agents as public for agent in returned_agents: if agent.litellm_params is None: agent.litellm_params = {} - agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and ( - agent.agent_id in litellm.public_agent_groups + agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and not ( + global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ) # Redact sensitive fields for non-admin users @@ -863,7 +881,7 @@ async def make_agent_public( if litellm.public_agent_groups is None: litellm.public_agent_groups = [] # handle duplicates - if agent.agent_id in litellm.public_agent_groups: + if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups): raise HTTPException( status_code=400, detail=f"Agent with name {agent.agent_name} already in public agent groups", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d016eb57dd4..90daaaeae6b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1016,6 +1016,37 @@ async def proxy_startup_event(app: FastAPI): asyncio.create_task(_run_pw_migration()) + async def _run_agent_grant_id_migration() -> None: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + object_permission_table, + ) + + for attempt in range(3): + try: + result = await global_agent_registry.migrate_legacy_grant_ids( + table=object_permission_table(prisma_client) + ) + if result.rewritten: + verbose_proxy_logger.info( + "Rewrote %s object_permission rows from legacy config agent ids", result.rewritten + ) + if result.missed == 0: + return + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 left %s rows unmigrated", + attempt + 1, + result.missed, + ) + except Exception as e: # noqa: BLE001 # startup task must survive any DB error and retry + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 failed: %s", attempt + 1, e + ) + if attempt < 2: + await asyncio.sleep(5) + + asyncio.create_task(_run_agent_grant_id_migration()) + ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before ## the coordination Redis is published to its consumers below. diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 78c3e9fd31b..79791607b2e 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -220,7 +220,7 @@ async def get_agents(request: Request): "url": get_custom_url(str(request.base_url), route=f"a2a/{agent.agent_id}"), } for agent in agents - if agent.agent_id in litellm.public_agent_groups + if not global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ] diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 66e0b6d59e7..be19f290ba6 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -42,6 +42,10 @@ class AgentsRepository(PrismaTableRepository): table_name = "litellm_agentstable" +class ObjectPermissionRepository(PrismaTableRepository): + table_name = "litellm_objectpermissiontable" + + class GuardrailsRepository(PrismaTableRepository): table_name = "litellm_guardrailstable" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 60356eda05b..afd8c0107ea 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -135,7 +135,7 @@ "limit": 27 }, "PERF401": { - "limit": 23 + "limit": 13 }, "PERF402": { "limit": 0 diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 6bdea9c2615..c57296dfc3f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -2,8 +2,11 @@ Unit tests for AgentRequestHandler - Agent permission management for keys and teams. """ +import hashlib +import json import os import sys +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -11,6 +14,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../..")) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) @@ -212,3 +216,79 @@ class TestAgentRequestHandler: user_api_key_auth=mock_user_auth ) assert sorted(result) == ["agent-from-ag", "native-agent-1"] + + async def test_is_agent_allowed_accepts_legacy_config_agent_id_grants(self): + """LIT-5144: object_permission grants stored under the pre-fix full-entry hash + must keep authorizing the agent after its id became name-based.""" + entry: Final = { + "agent_name": "granted-agent", + "agent_card_params": { + "name": "Granted Agent", + "url": "http://localhost", + "version": "1.0.0", + }, + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("granted-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ): + with patch.object(AgentRequestHandler, "get_allowed_agents") as mock_get_allowed: + for grant, expected in ( + ([legacy_id], True), + ([agent.agent_id], True), + (["unrelated-agent-id"], False), + ([], True), + ): + mock_get_allowed.return_value = grant + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent.agent_id, + user_api_key_auth=mock_user_auth, + ) + is expected + ), grant + + async def test_get_allowed_agents_intersects_legacy_team_grant_with_stable_key_grant(self): + """LIT-5144: a team grant stored under the pre-fix full-entry hash and a key grant + stored under the name-based id name the same agent; the intersection must resolve + to that agent instead of collapsing to the allow-all empty list.""" + entry: Final = { + "agent_name": "shared-agent", + "agent_card_params": { + "name": "Shared Agent", + "url": "http://localhost", + "version": "1.0.0", + }, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("shared-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ): + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + for key_grant, team_grant in ( + ([agent.agent_id], [legacy_id]), + ([legacy_id], [agent.agent_id]), + ([legacy_id], []), + ): + mock_key.return_value = key_grant + mock_team.return_value = team_grant + assert await AgentRequestHandler.get_allowed_agents(user_api_key_auth=mock_user_auth) == [ + agent.agent_id + ], (key_grant, team_grant) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index f506535b5e1..a87f5384c6f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -1,10 +1,14 @@ """Unit tests for AgentRegistry DB operations.""" +import hashlib +import json +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, GrantMigrationResult def _sample_agent_card_params() -> dict: @@ -153,7 +157,7 @@ def test_load_agents_from_db_and_config_skips_incomplete_config_entries(): registry.load_agents_from_db_and_config(db_agents=None) - assert registry.get_agent_list() == [] + assert registry.get_agent_list() == () @pytest.mark.parametrize( @@ -278,4 +282,149 @@ def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_agents assert registry.config_agents == () registry.load_agents_from_db_and_config(db_agents=None) - assert registry.get_agent_list() == [], "a removed config agent must not come back on the next rebuild" + assert registry.get_agent_list() == (), "a removed config agent must not come back on the next rebuild" + + +def test_config_agent_id_survives_static_header_secret_rotation(): + """LIT-5144: the id was a hash of the whole entry, so rotating a static_headers secret silently + re-identified the agent and orphaned every grant pointing at it.""" + base_entry: Final = { + "agent_name": "rotating-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry_v1: Final = AgentRegistry() + registry_v1.load_agents_from_config([base_entry]) + agent_v1: Final = registry_v1.get_agent_by_name("rotating-agent") + assert agent_v1 is not None + + registry_v2: Final = AgentRegistry() + registry_v2.load_agents_from_config([{**base_entry, "static_headers": {"x-upstream-token": "token-v2"}}]) + agent_v2: Final = registry_v2.get_agent_by_name("rotating-agent") + assert agent_v2 is not None + + assert agent_v1.agent_id == agent_v2.agent_id + + +def test_config_agent_ids_differ_when_only_the_agent_name_differs(): + """Two entries identical except for agent_name must not collapse onto one id.""" + registry: Final = AgentRegistry() + registry.load_agents_from_config( + [ + {"agent_name": "agent-a", "agent_card_params": _sample_agent_card_params()}, + {"agent_name": "agent-b", "agent_card_params": _sample_agent_card_params()}, + ] + ) + + ids: Final = {agent.agent_id for agent in registry.get_agent_list()} + assert len(ids) == 2 + + +def test_legacy_full_entry_hash_still_resolves_the_config_agent(): + """Grants and clients created before LIT-5144 hold the old full-entry hash; it must keep resolving.""" + entry: Final = { + "agent_name": "legacy-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("legacy-agent") + assert agent is not None + + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + assert registry.config_agent_legacy_ids[legacy_id] == agent.agent_id + assert legacy_id in registry.ids_for_agent(agent.agent_id) + assert agent.agent_id in registry.ids_for_agent(agent.agent_id) + + resolved: Final = registry.get_agent_by_id(legacy_id) + assert resolved is not None + assert resolved.agent_id == agent.agent_id + assert registry.get_agent_by_id("nonexistent-id") is None + + +def test_public_agent_groups_holding_the_legacy_id_still_mark_the_config_agent_public(monkeypatch): + """LIT-5144: config.yaml written before the fix stores the full-entry hash in + public_agent_groups; the agent must stay public after its id became name-based.""" + import litellm + + entry: Final = { + "agent_name": "public-agent", + "agent_card_params": _sample_agent_card_params(), + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("public-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + + monkeypatch.setattr(litellm, "public_agent_groups", [legacy_id]) + assert [a.agent_id for a in registry.get_public_agent_list()] == [agent.agent_id] + + monkeypatch.setattr(litellm, "public_agent_groups", ["unrelated-id"]) + assert registry.get_public_agent_list() == () + + monkeypatch.setattr(litellm, "public_agent_groups", None) + assert registry.get_public_agent_list() == () + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_persists_stable_ids_into_grant_rows(): + """LIT-5144: the startup migration rewrites stored legacy full-entry hashes to the stable + name id, so a later secret rotation (which re-mints the legacy hash) cannot orphan grants.""" + entry: Final = { + "agent_name": "migrated-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("migrated-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + + row: Final = SimpleNamespace(object_permission_id="op-1", agents=[legacy_id, "unrelated-id", agent.agent_id]) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[row]) + table.update_many = AsyncMock(return_value=1) + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=1, missed=0) + table.find_many.assert_awaited_once_with(where={"agents": {"has_some": (legacy_id,)}}) + table.update_many.assert_awaited_once_with( + where={"object_permission_id": "op-1", "agents": {"equals": (legacy_id, "unrelated-id", agent.agent_id)}}, + data={"agents": (agent.agent_id, "unrelated-id")}, + ) + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_reports_compare_and_swap_misses(): + """A concurrently edited row makes the CAS update affect zero rows; the result must + surface that as missed so the startup task knows to retry instead of reporting success.""" + entry: Final = { + "agent_name": "contended-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + + row: Final = SimpleNamespace(object_permission_id="op-1", agents=[legacy_id]) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[row]) + table.update_many = AsyncMock(return_value=0) + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=1) + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_no_ops_without_config_agents(): + """Without config agents there are no legacy hashes to translate, so the DB is never queried.""" + registry: Final = AgentRegistry() + table: Final = MagicMock() + table.find_many = AsyncMock() + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=0) + table.find_many.assert_not_awaited() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 3e097711ad7..6dbd7475be5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -308,7 +308,7 @@ async def test_attach_keys_to_agents_groups_by_agent_and_omits_secret(): # Query is scoped to the agents being returned, not the whole key table. where = mock_prisma.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] - assert where == {"agent_id": {"in": ["agent-1", "agent-2"]}} + assert where == {"agent_id": {"in": ("agent-1", "agent-2")}} # agent-1 gets both of its keys; agent-2 gets None. assert agent_without_keys.keys is None @@ -503,6 +503,7 @@ class TestAgentRBACProxyAdminViewOnly: ] self.mock_registry = MagicMock() self.mock_registry.get_agent_list = MagicMock(return_value=self.agents) + self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry) self.allowed_agents_spy = AsyncMock(return_value=["someone-elses-agent"]) diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 8b8b7871cce..88dc07e741b 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -582,6 +582,7 @@ def test_public_agent_hub_rewrites_upstream_url_to_proxy(): mock_registry = MagicMock() mock_registry.get_public_agent_list.return_value = [agent] + mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) with ( patch("litellm.public_agent_groups", ["agent-123"]), @@ -631,6 +632,7 @@ def test_public_agent_hub_serializes_http_security_scheme_without_bearer_format( mock_registry = MagicMock() mock_registry.get_public_agent_list.return_value = [agent] + mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) with ( patch("litellm.public_agent_groups", ["agent-123"]), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2921592acbd..b824f5bb550 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23256 + "limit": 23250 }, "LIT002": { "limit": 27195 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16783 + "limit": 16777 }, "LIT011": { "limit": 5602 From 1f61c4399790ce8eba7853a65807ce365674ab0e Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:23:57 +0000 Subject: [PATCH 32/74] build(deps): bump gitpython to 3.1.58, defer pypdf advisory Closes the six gitpython advisories flagged on litellm_internal_staging (GHSA-jm78-9fvv-mhgr and GHSA-wvpp-8hx9-p66j at 8.8, GHSA-hmq2-w58f-27jc at 8.2, GHSA-4gmw-gg2m-w46p at 8.1, GHSA-9rj7-rf2p-w77r at 7.5, GHSA-hh9p-6wh2-4mfc at 6.5). gitpython comes in transitively through mlflow-skinny, so this is a lock-only change re-derived with 'uv lock --upgrade-package gitpython'. The seventh finding, GHSA-fwg2-594c-jp42 on pypdf, cannot be fixed the same way today: pypdf 6.15.0 published 2026-08-06 and the repo pins exclude-newer to a 3 day window, so uv will not resolve it before 2026-08-09. Rather than widen that window, the advisory gets a short dated IgnoredVulns entry that expires 2026-08-12, which leaves a hard deadline to land the real bump. It is a local, user-interaction denial of service on crafted CID font widths at CVSS 4.8, so a couple of days of exposure in the lock is acceptable. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- osv-scanner.toml | 5 +++++ uv.lock | 8 ++++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/osv-scanner.toml b/osv-scanner.toml index 7ab450945f5..efcbe6c8c16 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -1,3 +1,8 @@ +[[IgnoredVulns]] +id = "GHSA-fwg2-594c-jp42" +ignoreUntil = 2026-08-12 +reason = "pypdf 6.15.0 (the fix) published 2026-08-06 and is still inside the P3D exclude-newer window, so uv cannot lock it yet; bump and drop this entry from 2026-08-09" + [[IgnoredVulns]] id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 diff --git a/uv.lock b/uv.lock index 5ac80a92568..779bf128293 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-04T01:43:29.894567Z" +exclude-newer = "2026-08-04T19:23:00.022310687Z" exclude-newer-span = "P3D" [manifest] @@ -2364,14 +2364,14 @@ wheels = [ [[package]] name = "gitpython" -version = "3.1.57" +version = "3.1.58" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "gitdb" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ba/0d/132ed135c871b6bf91adf16a0e43797cd535b81d4973b5d09291c54fc5ee/gitpython-3.1.57.tar.gz", hash = "sha256:c493ec57c0ef6b19743798b6a5af859c71814b524e7e6f97baa2f8e658961488", size = 225898, upload-time = "2026-07-26T07:33:26.351Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/d6/5f358ff283325580c2003a6d953aea18cfe10ae87b46f5ebc80fa3a386dc/gitpython-3.1.58.tar.gz", hash = "sha256:621416df10ef3fd0e19fabf9172ddeed0fa704d353d04f194eec56a625a95b22", size = 228498, upload-time = "2026-08-04T15:05:49.47Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/41/6e/2139de986d9c7c3ac86f1f8be43858ce90bdfe2f7175e6c80c650ba15242/gitpython-3.1.57-py3-none-any.whl", hash = "sha256:4ccf7d73c10f5c9e76043fbb2675ac5a1b3ff5b41e648f56bcbed5f63792ecaf", size = 217151, upload-time = "2026-07-26T07:33:24.838Z" }, + { url = "https://files.pythonhosted.org/packages/ec/0c/9d8752098bc442f0726e64aa6135940b3a96809915d1aa4206c1bb97881d/gitpython-3.1.58-py3-none-any.whl", hash = "sha256:d331e722577f0fd7fc1f857419b3ecc07af66282b933d2a4d95f84a042fdd50f", size = 220183, upload-time = "2026-08-04T15:05:48.025Z" }, ] [[package]] From ecd5ad49f709e549e2f28e9bc71d4f49e3efb1d0 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:24:46 +0000 Subject: [PATCH 33/74] ci: always run the UI API types sync check so it can be required Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/check-ui-api-types.yml | 42 +++++++++++++++++++----- 1 file changed, 34 insertions(+), 8 deletions(-) diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 597daebd720..587d6595bcc 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -2,17 +2,19 @@ name: Check UI API Types Sync on: pull_request: - paths: - - "litellm/proxy/**" - - "litellm/types/**" - - "ui/litellm-dashboard/src/lib/http/schema.d.ts" - - "ui/litellm-dashboard/scripts/gen-api-types.mjs" - - "ui/litellm-dashboard/package.json" - - "ui/litellm-dashboard/package-lock.json" - - ".github/workflows/check-ui-api-types.yml" + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" permissions: contents: read + pull-requests: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true jobs: check-sync: @@ -25,17 +27,35 @@ jobs: with: persist-credentials: false + - name: Detect changes that can affect the generated types + id: changes + env: + GH_TOKEN: ${{ github.token }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ github.event.pull_request.number }} + run: | + set -euo pipefail + files="$(gh api --paginate "repos/$REPO/pulls/$PR_NUMBER/files" --jq '.[].filename')" + if printf '%s\n' "$files" | grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)'; then + echo "relevant=true" >> "$GITHUB_OUTPUT" + else + echo "relevant=false" >> "$GITHUB_OUTPUT" + fi + - name: Set up Python + if: steps.changes.outputs.relevant == 'true' uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - name: Set up uv + if: steps.changes.outputs.relevant == 'true' uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - name: Cache uv dependencies + if: steps.changes.outputs.relevant == 'true' uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: | @@ -46,14 +66,17 @@ jobs: ${{ runner.os }}-uv- - 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 - name: Generate Prisma client + if: steps.changes.outputs.relevant == 'true' env: PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - name: Set up Node.js + if: steps.changes.outputs.relevant == 'true' uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0 with: node-version-file: ui/litellm-dashboard/.nvmrc @@ -61,16 +84,19 @@ jobs: cache-dependency-path: ui/litellm-dashboard/package-lock.json - name: Install dashboard dependencies + if: steps.changes.outputs.relevant == 'true' working-directory: ui/litellm-dashboard run: npm ci - name: Regenerate types from the live spec + if: steps.changes.outputs.relevant == 'true' working-directory: ui/litellm-dashboard env: LITELLM_PYTHON: "uv run --no-sync python" run: npm run gen:api - name: Fail if types are stale + if: steps.changes.outputs.relevant == 'true' run: | if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec." From bc71564d1d4132005f9c7fa62279b88b8dd43ca1 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:58:22 +0000 Subject: [PATCH 34/74] build(deps): defer the second pypdf advisory until the 6.15.0 bump GHSA-fp3f-mc75-235c was published at 19:29 UTC today, hours after #36212 landed, and covers a different pypdf issue (large memory use on big /ToUnicode streams) that lands on the same 6.15.0 fix. It reds osv-scan on every PR into staging again. pypdf 6.15.0 is still inside the repo's P3D exclude-newer window until 08-09, so uv cannot lock it yet and widening exclude-newer to pull it in early would weaken the freshness guard for every package. This gets the same 08-12 ignoreUntil as the first pypdf entry so both drop together in the bump. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- osv-scanner.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/osv-scanner.toml b/osv-scanner.toml index efcbe6c8c16..4ef612e3a70 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -3,6 +3,11 @@ id = "GHSA-fwg2-594c-jp42" ignoreUntil = 2026-08-12 reason = "pypdf 6.15.0 (the fix) published 2026-08-06 and is still inside the P3D exclude-newer window, so uv cannot lock it yet; bump and drop this entry from 2026-08-09" +[[IgnoredVulns]] +id = "GHSA-fp3f-mc75-235c" +ignoreUntil = 2026-08-12 +reason = "second pypdf advisory with the same 6.15.0 fix, published 2026-08-07 after the first; drop alongside GHSA-fwg2-594c-jp42 in the same bump" + [[IgnoredVulns]] id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 From 78addb230b3f8bdf7286ef3d7cd22fc19fad0bab Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 13:44:11 -0700 Subject: [PATCH 35/74] fix(proxy): deny agent access when key and team grants resolve to nothing (#36221) * fix(proxy): deny when agent grants resolve to nothing `get_allowed_agents` returned a plain list where the empty value meant both "this caller was never restricted" and "this caller's grants resolved to nothing". Downstream read either as allow-all, so a key restricted to one agent inside a team restricted to another reached every agent on the proxy, and an access group that resolved to no agents did the same. Replace it with `resolve_agent_access`, returning a tagged UnrestrictedAgentAccess | RestrictedAgentAccess. Only a caller with no grant anywhere is unrestricted; an empty restricted set denies. Access group lookup failures now propagate to the key/team resolvers so a DB error still fails open exactly as before, while a group that genuinely resolves to nothing denies. * style(proxy): drop redundant comments from the agent access match --- basedpyright-code-budget.json | 2 +- .../auth/agent_permission_handler.py | 241 ++++++++++-------- litellm/proxy/agent_endpoints/endpoints.py | 45 ++-- .../agent_endpoints/model_list_helpers.py | 76 +++--- ruff-strict-budget.json | 6 +- .../proxy/agent_endpoints/test_agent_rbac.py | 8 +- .../auth/test_agent_permission_handler.py | 209 +++++++++++---- .../proxy/agent_endpoints/test_endpoints.py | 16 +- .../test_model_list_helpers.py | 11 +- .../test_activity_tenant_scoping.py | 20 +- type-discipline-budget.json | 6 +- 11 files changed, 401 insertions(+), 239 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 632743236c4..c8765eb0bd0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1825 + "limit": 1824 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 6baeab3dfff..81586b4eef3 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -5,7 +5,8 @@ Handles agent permission checking for keys and teams using object_permission_id. Follows the same pattern as MCP permission handling. """ -from typing import Final +from dataclasses import dataclass +from typing import Final, TypeAlias from litellm._logging import verbose_logger from litellm.proxy._types import ( @@ -17,6 +18,27 @@ from litellm.proxy._types import ( from litellm.repositories.table_repositories import AgentsRepository +@dataclass(frozen=True, slots=True) +class UnrestrictedAgentAccess: + """No agent grant exists on the key or its team, so every agent is reachable.""" + + +@dataclass(frozen=True, slots=True) +class RestrictedAgentAccess: + """Only ``agent_ids`` are reachable. An empty set denies every agent.""" + + agent_ids: frozenset[str] + + +AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess + + +def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids) + + class AgentRequestHandler: """ Class to handle agent permission checking, including: @@ -32,38 +54,32 @@ class AgentRequestHandler: """ @staticmethod - async def get_allowed_agents( + async def resolve_agent_access( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ - Get list of allowed agent IDs for the given user/key based on permissions. + Resolve the agents the given user/key may reach. - Returns: - List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). + ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team + carries any grant. Grants that intersect to nothing stay restricted, so + narrowing a caller can never widen what it reaches. """ try: - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) + team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - raw_key_grants: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - raw_team_grants: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - allowed_agents_for_key: Final = frozenset( - global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_key_grants - ) - allowed_agents_for_team: Final = frozenset( - global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_team_grants - ) - - # If team has agent restrictions, handle inheritance and intersection logic - if allowed_agents_for_team and allowed_agents_for_key: - # Key has its own agent permissions - use intersection with team permissions - return sorted(allowed_agents_for_key & allowed_agents_for_team) - if allowed_agents_for_team: - # Key has no agent permissions - inherit from team - return sorted(allowed_agents_for_team) - return sorted(allowed_agents_for_key) + match (key_access, team_access): + case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()): + return UnrestrictedAgentAccess() + case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)): + return RestrictedAgentAccess(_to_stable_ids(team_ids)) + case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()): + return RestrictedAgentAccess(_to_stable_ids(key_ids)) + case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)): + return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids)) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod async def is_agent_allowed( @@ -82,14 +98,12 @@ class AgentRequestHandler: """ from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - allowed_agents: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth) - - # Empty list means no restrictions - allow all - if len(allowed_agents) == 0: - return True - - stable_id: Final = global_agent_registry.stable_agent_id(agent_id) - return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agents) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth): + case UnrestrictedAgentAccess(): + return True + case RestrictedAgentAccess(allowed_agent_ids): + stable_id: Final = global_agent_registry.stable_agent_id(agent_id) + return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids) @staticmethod def _get_key_object_permission( @@ -143,55 +157,58 @@ class AgentRequestHandler: @staticmethod async def _get_allowed_agents_for_key( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ Get allowed agents for a key. 1. First checks native key-level agent permissions (object_permission) 2. Also includes agents from key's access_group_ids (unified access groups) + A key that declares agents or access groups is restricted even when those + declarations resolve to nothing, so an emptied or deleted access group denies + rather than opening the key up. Lookup failures still propagate to the caller, + which keeps them fail-open. + Note: object_permission is already loaded by get_key_object() in main auth flow. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgentAccess() try: - all_agents: list[str] = [] - # 1. Get agents from object_permission (native permissions) key_object_permission: Final = AgentRequestHandler._get_key_object_permission(user_api_key_auth) - if key_object_permission is not None: - # Get direct agents - direct_agents: Final = key_object_permission.agents or [] - - # Get agents from access groups - access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups( - key_object_permission.agent_access_groups or [] - ) - - all_agents = direct_agents + access_group_agents - + direct_agents: Final = tuple( + key_object_permission.agents or () if key_object_permission is not None else () + ) + declared_access_groups: Final = tuple( + key_object_permission.agent_access_groups or () if key_object_permission is not None else () + ) # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) - key_access_group_ids: Final = user_api_key_auth.access_group_ids or [] - if key_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ()) - unified_agents: Final = await _get_agent_ids_from_access_groups( - access_group_ids=key_access_group_ids, - ) - all_agents.extend(unified_agents) + if not direct_agents and not declared_access_groups and not key_access_group_ids: + return UnrestrictedAgentAccess() - return list(set(all_agents)) + access_group_agents: Final = ( + tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple(await AgentRequestHandler._get_unified_access_group_agents(list(key_access_group_ids))) + if key_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: verbose_logger.warning("Failed to get allowed agents for key: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod async def _get_allowed_agents_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ Get allowed agents for a team. @@ -199,12 +216,13 @@ class AgentRequestHandler: 2. Also includes agents from team's access_group_ids (unified access groups) Fetches the team object once and reuses it for both permission sources. + Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgentAccess() if user_api_key_auth.team_id is None: - return [] + return UnrestrictedAgentAccess() try: from litellm.proxy.auth.auth_checks import get_team_object @@ -215,7 +233,7 @@ class AgentRequestHandler: ) if not prisma_client: - return [] + return UnrestrictedAgentAccess() # Fetch the team object once for both permission sources team_obj: Final = await get_team_object( @@ -227,42 +245,38 @@ class AgentRequestHandler: ) if team_obj is None: - return [] - - all_agents: list[str] = [] + return UnrestrictedAgentAccess() # 1. Get agents from object_permission (native permissions) object_permissions: Final = team_obj.object_permission - if object_permissions is not None: - # Get direct agents - direct_agents: Final = object_permissions.agents or [] - - # Get agents from access groups - access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups( - object_permissions.agent_access_groups or [] - ) - - all_agents = direct_agents + access_group_agents - + direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ()) + declared_access_groups: Final = tuple( + object_permissions.agent_access_groups or () if object_permissions is not None else () + ) # 2. Also include agents from team's access_group_ids (unified access groups) - team_access_group_ids: Final = team_obj.access_group_ids or [] - if team_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + team_access_group_ids: Final = tuple(team_obj.access_group_ids or ()) - unified_agents: Final = await _get_agent_ids_from_access_groups( - access_group_ids=team_access_group_ids, - ) - all_agents.extend(unified_agents) + if not direct_agents and not declared_access_groups and not team_access_group_ids: + return UnrestrictedAgentAccess() - return list(set(all_agents)) + access_group_agents: Final = ( + tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple(await AgentRequestHandler._get_unified_access_group_agents(list(team_access_group_ids))) + if team_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: verbose_logger.warning("Failed to get allowed agents for team: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]: @@ -281,18 +295,26 @@ class AgentRequestHandler: async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. + + Query failures propagate so the caller can tell "this group is empty" (deny) + apart from "the lookup failed" (fail-open). """ - agent_ids: Final[set[str]] = set() - if access_groups and prisma_client is not None: - try: - agents: Final = await AgentsRepository(prisma_client).table.find_many( - where={"agent_access_groups": {"hasSome": access_groups}} - ) - for agent in agents: - agent_ids.add(agent.agent_id) - except Exception as e: - verbose_logger.debug("Error getting agents from access groups: %s", e) - return agent_ids + if not access_groups or prisma_client is None: + return set() + + agents: Final = await AgentsRepository(prisma_client).table.find_many( + where={"agent_access_groups": {"hasSome": access_groups}} + ) + return {agent.agent_id for agent in agents} + + @staticmethod + async def _get_unified_access_group_agents(access_group_ids: list[str]) -> list[str]: + """ + Resolve unified access group ids to agent IDs. + """ + from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups + + return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids) @staticmethod async def _get_agents_from_access_groups( @@ -304,20 +326,17 @@ class AgentRequestHandler: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.proxy_server import prisma_client - try: - # Use the helper for config-loaded agents - agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( - global_agent_registry.agent_list, access_groups - ) + # Use the helper for config-loaded agents + config_agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( + global_agent_registry.agent_list, access_groups + ) - # Use the helper for DB agents - db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups(prisma_client, access_groups) - agent_ids.update(db_agent_ids) + # Use the helper for DB agents + db_agent_ids: Final = await AgentRequestHandler._get_db_agent_ids_for_access_groups( + prisma_client, access_groups + ) - return list(agent_ids) - except Exception as e: - verbose_logger.warning("Failed to get agents from access groups: %s", e) - return [] + return list(config_agent_ids | db_agent_ids) @staticmethod async def get_agent_access_groups( diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 028885aa065..5bc5b5566f8 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -248,6 +248,8 @@ async def get_agents( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) try: @@ -261,15 +263,14 @@ async def get_agents( returned_agents = global_agent_registry.get_agent_list() else: # Get allowed agents from object_permission (key/team level) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) + agent_access: Final = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict) + all_agents: Final = global_agent_registry.get_agent_list() - # If no restrictions (empty list), return all agents - if len(allowed_agent_ids) == 0: - returned_agents = global_agent_registry.get_agent_list() - else: - # Filter agents by allowed IDs - all_agents: Final = global_agent_registry.get_agent_list() - returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] + match agent_access: + case UnrestrictedAgentAccess(): + returned_agents = all_agents + case RestrictedAgentAccess(allowed_agent_ids): + returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] # Fetch current spend from DB for all returned agents from litellm.proxy.proxy_server import prisma_client @@ -1061,27 +1062,29 @@ async def get_agent_daily_activity( # intersect their explicit `agent_ids` filter with the same allowlist. from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view where_condition: Final[dict[str, object]] = {} if not _user_has_admin_view(user_api_key_dict): - permitted_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - # `get_allowed_agents` returns an empty list when the caller's key - # and team carry no agent restrictions. For activity scoping that's - # not "see everything" — fall back to the agents the caller - # created so they cannot enumerate other tenants' agents. + permitted_agent_ids: list[str] = [] + # An unrestricted caller is not "see everything" for activity scoping. Fall + # back to the agents the caller created so they cannot enumerate other + # tenants' agents. # Guard against `user_id is None`: a literal None in Prisma # `where={"created_by": None}` resolves to ``created_by IS NULL`` # and would expose every ownerless agent's rows. - if not permitted_agent_ids: - if user_api_key_dict.user_id is None: - permitted_agent_ids = [] - else: - owned_records: Final = await agents_table(prisma_client).find_many( - where={"created_by": user_api_key_dict.user_id} - ) - permitted_agent_ids = [a.agent_id for a in owned_records] + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + permitted_agent_ids = list(allowed_agent_ids) + case UnrestrictedAgentAccess(): + if user_api_key_dict.user_id is not None: + owned_records: Final = await agents_table(prisma_client).find_many( + where={"created_by": user_api_key_dict.user_id} + ) + permitted_agent_ids = [a.agent_id for a in owned_records] if agent_ids_list: permitted_agent_id_set: Final = set(permitted_agent_ids) diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 862b66d3860..4a88644bae5 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -4,8 +4,6 @@ Helper functions for appending A2A agents to model lists. Used by proxy model endpoints to make agents appear in UI alongside models. """ -from typing import Final - from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -27,20 +25,23 @@ async def append_agents_to_model_group( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, ) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - for agent_id in allowed_agent_ids: - agent = global_agent_registry.get_agent_by_id(agent_id) - if agent is not None: - model_groups.append( - ModelGroupInfoProxy( - model_group=f"a2a/{agent.agent_name}", - mode="chat", - providers=["a2a"], - ) - ) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + for agent_id in allowed_agent_ids: + agent = global_agent_registry.get_agent_by_id(agent_id) + if agent is not None: + model_groups.append( + ModelGroupInfoProxy( + model_group=f"a2a/{agent.agent_name}", + mode="chat", + providers=["a2a"], + ) + ) + case _: + pass except Exception as e: verbose_proxy_logger.debug("Error appending agents to model_group/info: %s", e) @@ -61,30 +62,33 @@ async def append_agents_to_model_info( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, ) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - for agent_id in allowed_agent_ids: - agent = global_agent_registry.get_agent_by_id(agent_id) - if agent is not None: - models.append( - { - "model_name": f"a2a/{agent.agent_name}", - "litellm_params": { - "model": f"a2a/{agent.agent_name}", - "custom_llm_provider": "a2a", - }, - "model_info": { - "id": agent.agent_id, - "mode": "chat", - "db_model": True, - "created_by": agent.created_by, - "created_at": agent.created_at, - "updated_at": agent.updated_at, - }, - } - ) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + for agent_id in allowed_agent_ids: + agent = global_agent_registry.get_agent_by_id(agent_id) + if agent is not None: + models.append( + { + "model_name": f"a2a/{agent.agent_name}", + "litellm_params": { + "model": f"a2a/{agent.agent_name}", + "custom_llm_provider": "a2a", + }, + "model_info": { + "id": agent.agent_id, + "mode": "chat", + "db_model": True, + "created_by": agent.created_by, + "created_at": agent.created_at, + "updated_at": agent.updated_at, + }, + } + ) + case _: + pass except Exception as e: verbose_proxy_logger.debug("Error appending agents to v2/model/info: %s", e) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index afd8c0107ea..8ff4bfb36c0 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 834 }, "ANN201": { - "limit": 2033 + "limit": 2032 }, "ANN202": { "limit": 865 @@ -60,7 +60,7 @@ "limit": 0 }, "BLE001": { - "limit": 2926 + "limit": 2924 }, "C401": { "limit": 8 @@ -135,7 +135,7 @@ "limit": 27 }, "PERF401": { - "limit": 13 + "limit": 12 }, "PERF402": { "limit": 0 diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py index 78cb1488533..842a635ecad 100644 --- a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -11,6 +11,10 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: @@ -64,8 +68,8 @@ async def test_get_agents_allowed_when_not_disabled(): MagicMock(get_agent_list=MagicMock(return_value=[])), ): with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ): result = await get_agents(request=request_mock, user_api_key_dict=user) assert result == [] diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index c57296dfc3f..066d33dc187 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -17,6 +17,8 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) @@ -26,11 +28,11 @@ class TestAgentRequestHandler: Test suite for AgentRequestHandler permission logic. """ - async def test_get_allowed_agents_intersection_logic(self): + async def test_resolve_agent_access_intersection_logic(self): """ Test key/team intersection: when both have restrictions, only common agents are allowed. When team has restrictions but key has none, key inherits from team. - When neither has restrictions, returns empty list (meaning allow all). + Only a caller with no grant anywhere is unrestricted. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -45,13 +47,17 @@ class TestAgentRequestHandler: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = ["agent1", "agent2", "agent3"] - mock_team.return_value = ["agent2", "agent4"] + mock_key.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2", "agent3"}) + ) + mock_team.return_value = RestrictedAgentAccess( + frozenset({"agent2", "agent4"}) + ) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent2"] + assert result == RestrictedAgentAccess(frozenset({"agent2"})) # Case 2: Team has agents, key has none - inherit from team with patch.object( @@ -60,41 +66,124 @@ class TestAgentRequestHandler: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = [] - mock_team.return_value = ["team_agent1", "team_agent2"] + mock_key.return_value = UnrestrictedAgentAccess() + mock_team.return_value = RestrictedAgentAccess( + frozenset({"team_agent1", "team_agent2"}) + ) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["team_agent1", "team_agent2"] + assert result == RestrictedAgentAccess( + frozenset({"team_agent1", "team_agent2"}) + ) - # Case 3: No restrictions - returns empty list (allow all) + # Case 3: Key has agents, team has none - key restrictions stand with patch.object( AgentRequestHandler, "_get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = [] - mock_team.return_value = [] + mock_key.return_value = RestrictedAgentAccess(frozenset({"key_agent1"})) + mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert result == [] + assert result == RestrictedAgentAccess(frozenset({"key_agent1"})) + + # Case 4: No grant anywhere - unrestricted (documented open-by-default) + with patch.object( + AgentRequestHandler, "_get_allowed_agents_for_key" + ) as mock_key: + with patch.object( + AgentRequestHandler, "_get_allowed_agents_for_team" + ) as mock_team: + mock_key.return_value = UnrestrictedAgentAccess() + mock_team.return_value = UnrestrictedAgentAccess() + + result = await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) + assert result == UnrestrictedAgentAccess() + + async def test_disjoint_key_and_team_grants_deny_every_agent(self): + """LIT-5143: a key restricted to one agent inside a team restricted to another + must reach nothing. The empty intersection used to read as "no restrictions", + so adding the team grant handed the key every agent on the proxy.""" + mock_user_auth: Final = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + mock_key.return_value = RestrictedAgentAccess(frozenset({"agent-alpha"})) + mock_team.return_value = RestrictedAgentAccess(frozenset({"agent-beta"})) + + assert await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset()) + + for agent_id in ("agent-alpha", "agent-beta", "agent-secret"): + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent_id, user_api_key_auth=mock_user_auth + ) + is False + ), agent_id + + async def test_empty_access_group_denies_every_agent(self): + """LIT-5143: a key restricted to an access group that resolves to no agents is + restricted to nothing, not unrestricted. A failed group lookup still fails open.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + mock_user_auth.object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="obj-1", + agents=[], + agent_access_groups=["group-with-no-agents"], + ) + + with patch.object( + AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.return_value = [] + + assert await AgentRequestHandler._get_allowed_agents_for_key( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset()) + + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id="agent-secret", user_api_key_auth=mock_user_auth + ) + is False + ) + + with patch.object( + AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.side_effect = Exception("DB Error") + + assert await AgentRequestHandler._get_allowed_agents_for_key( + user_api_key_auth=mock_user_auth + ) == UnrestrictedAgentAccess() async def test_is_agent_allowed_respects_permissions(self): """ - Test is_agent_allowed: returns True if agent in allowed list or if no restrictions. + Test is_agent_allowed: returns True if agent in allowed list or if unrestricted. Returns False if agent not in allowed list. """ mock_user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") # Agent in allowed list - should be allowed with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + mock_get_allowed.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2"}) + ) assert ( await AgentRequestHandler.is_agent_allowed( agent_id="agent1", user_api_key_auth=mock_user_auth @@ -104,9 +193,11 @@ class TestAgentRequestHandler: # Agent not in allowed list - should be denied with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + mock_get_allowed.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2"}) + ) assert ( await AgentRequestHandler.is_agent_allowed( agent_id="agent3", user_api_key_auth=mock_user_auth @@ -114,11 +205,23 @@ class TestAgentRequestHandler: is False ) - # Empty list means no restrictions - should allow any agent + # Restricted to nothing - should deny every agent with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = [] + mock_get_allowed.return_value = RestrictedAgentAccess(frozenset()) + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id="any_agent", user_api_key_auth=mock_user_auth + ) + is False + ) + + # Unrestricted - should allow any agent + with patch.object( + AgentRequestHandler, "resolve_agent_access" + ) as mock_get_allowed: + mock_get_allowed.return_value = UnrestrictedAgentAccess() assert ( await AgentRequestHandler.is_agent_allowed( agent_id="any_agent", user_api_key_auth=mock_user_auth @@ -130,17 +233,19 @@ class TestAgentRequestHandler: """ Test that when user_api_key_auth is None, all agents are allowed (no restrictions). """ - result = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=None) - assert result == [] + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=None) + assert result == UnrestrictedAgentAccess() is_allowed = await AgentRequestHandler.is_agent_allowed( agent_id="any_agent", user_api_key_auth=None ) assert is_allowed is True - async def test_get_allowed_agents_handles_errors_gracefully(self): + async def test_resolve_agent_access_handles_errors_gracefully(self): """ - Test that errors during permission lookup are handled gracefully (returns empty list). + Test that errors during permission lookup are handled gracefully. This stays + fail-open for now to preserve existing availability behavior; fail-closed is + tracked separately. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -156,12 +261,12 @@ class TestAgentRequestHandler: AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: mock_key.side_effect = Exception("DB Error") - mock_team.return_value = [] + mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert result == [] + assert result == UnrestrictedAgentAccess() async def test_get_allowed_agents_for_key_via_access_group_ids(self): """ @@ -185,7 +290,9 @@ class TestAgentRequestHandler: result = await AgentRequestHandler._get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent-from-ag-1", "agent-from-ag-2"] + assert result == RestrictedAgentAccess( + frozenset({"agent-from-ag-1", "agent-from-ag-2"}) + ) async def test_get_allowed_agents_for_key_combines_native_and_access_groups(self): """ @@ -215,7 +322,9 @@ class TestAgentRequestHandler: result = await AgentRequestHandler._get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent-from-ag", "native-agent-1"] + assert result == RestrictedAgentAccess( + frozenset({"agent-from-ag", "native-agent-1"}) + ) async def test_is_agent_allowed_accepts_legacy_config_agent_id_grants(self): """LIT-5144: object_permission grants stored under the pre-fix full-entry hash @@ -241,12 +350,13 @@ class TestAgentRequestHandler: "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry, ): - with patch.object(AgentRequestHandler, "get_allowed_agents") as mock_get_allowed: + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: for grant, expected in ( - ([legacy_id], True), - ([agent.agent_id], True), - (["unrelated-agent-id"], False), - ([], True), + (RestrictedAgentAccess(frozenset({legacy_id})), True), + (RestrictedAgentAccess(frozenset({agent.agent_id})), True), + (RestrictedAgentAccess(frozenset({"unrelated-agent-id"})), False), + (RestrictedAgentAccess(frozenset()), False), + (UnrestrictedAgentAccess(), True), ): mock_get_allowed.return_value = grant assert ( @@ -257,10 +367,10 @@ class TestAgentRequestHandler: is expected ), grant - async def test_get_allowed_agents_intersects_legacy_team_grant_with_stable_key_grant(self): + async def test_resolve_agent_access_intersects_legacy_team_grant_with_stable_key_grant(self): """LIT-5144: a team grant stored under the pre-fix full-entry hash and a key grant stored under the name-based id name the same agent; the intersection must resolve - to that agent instead of collapsing to the allow-all empty list.""" + to that agent instead of collapsing to an empty set.""" entry: Final = { "agent_name": "shared-agent", "agent_card_params": { @@ -283,12 +393,21 @@ class TestAgentRequestHandler: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: for key_grant, team_grant in ( - ([agent.agent_id], [legacy_id]), - ([legacy_id], [agent.agent_id]), - ([legacy_id], []), + ( + RestrictedAgentAccess(frozenset({agent.agent_id})), + RestrictedAgentAccess(frozenset({legacy_id})), + ), + ( + RestrictedAgentAccess(frozenset({legacy_id})), + RestrictedAgentAccess(frozenset({agent.agent_id})), + ), + ( + RestrictedAgentAccess(frozenset({legacy_id})), + UnrestrictedAgentAccess(), + ), ): mock_key.return_value = key_grant mock_team.return_value = team_grant - assert await AgentRequestHandler.get_allowed_agents(user_api_key_auth=mock_user_auth) == [ - agent.agent_id - ], (key_grant, team_grant) + assert await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset({agent.agent_id})), (key_grant, team_grant) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 6dbd7475be5..ea196bda529 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -6,6 +6,10 @@ from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) from litellm.proxy.agent_endpoints.endpoints import ( _attach_keys_to_agents, _check_agent_management_permission, @@ -506,9 +510,11 @@ class TestAgentRBACProxyAdminViewOnly: self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry) - self.allowed_agents_spy = AsyncMock(return_value=["someone-elses-agent"]) + self.allowed_agents_spy = AsyncMock( + return_value=RestrictedAgentAccess(frozenset({"someone-elses-agent"})) + ) monkeypatch.setattr( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", self.allowed_agents_spy, ) @@ -537,9 +543,9 @@ class TestAgentRBACProxyAdminViewOnly: self.allowed_agents_spy.assert_awaited_once() def test_should_still_redact_secrets_for_view_only_admin(self): - """An unrestricted viewer (empty allowlist means no restrictions) sees the - same agents as an admin but with keys stripped and litellm_params masked.""" - self.allowed_agents_spy.return_value = [] + """An unrestricted viewer sees the same agents as an admin but with keys + stripped and litellm_params masked.""" + self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py index 86eb7a83079..ccf5942c89d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -13,6 +13,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, +) from litellm.proxy.agent_endpoints.model_list_helpers import ( append_agents_to_model_group, append_agents_to_model_info, @@ -37,14 +40,14 @@ async def test_append_agents_to_model_group(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["test-agent-id"]) + mock_get_allowed_agents = AsyncMock(return_value=RestrictedAgentAccess(frozenset({"test-agent-id"}))) # Mock global_agent_registry mock_registry = Mock() mock_registry.get_agent_by_id = Mock(return_value=mock_agent) with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", mock_get_allowed_agents, ): with patch( @@ -80,14 +83,14 @@ async def test_append_agents_to_model_info(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["agent-123"]) + mock_get_allowed_agents = AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))) # Mock global_agent_registry mock_registry = Mock() mock_registry.get_agent_by_id = Mock(return_value=mock_agent) with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", mock_get_allowed_agents, ): with patch( diff --git a/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py b/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py index 0855c201945..61583d11dfa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py +++ b/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py @@ -13,6 +13,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) # --------------------------------------------------------------------------- @@ -227,8 +231,8 @@ async def test_agent_activity_non_admin_no_perms_falls_back_to_owned(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), # no explicit agent permissions + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), # no explicit agent permissions ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -275,8 +279,8 @@ async def test_agent_activity_non_admin_intersects_explicit_agent_ids(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=["agent-permitted"]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-permitted"}))), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -321,8 +325,8 @@ async def test_agent_activity_keyless_caller_does_not_query_created_by_null(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -366,8 +370,8 @@ async def test_agent_activity_non_admin_no_access_returns_empty_page(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", diff --git a/type-discipline-budget.json b/type-discipline-budget.json index b824f5bb550..991f8eaa934 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23250 + "limit": 23245 }, "LIT002": { - "limit": 27195 + "limit": 27179 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16777 + "limit": 16769 }, "LIT011": { "limit": 5602 From 16d650ca94e00a21ce0aad1fe292d174c2fbc493 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 21:05:59 +0000 Subject: [PATCH 36/74] test(proxy): compare empty agent list to the tuple get_agent_list returns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/proxy_server/test_proxy_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 30156849628..91a7e1bc2c2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2638,4 +2638,4 @@ async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembe assert clean_agent_registry.config_agents == () clean_agent_registry.load_agents_from_db_and_config(db_agents=None) - assert clean_agent_registry.get_agent_list() == [] + assert clean_agent_registry.get_agent_list() == () From 4306c4ad70e4f50f8415d5e47e13884e0152596a Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 21:07:27 +0000 Subject: [PATCH 37/74] ci: detect relevant changes from git instead of the paginated files API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/check-ui-api-types.yml | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 587d6595bcc..02543d67a82 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -10,7 +10,6 @@ on: permissions: contents: read - pull-requests: read concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} @@ -26,19 +25,22 @@ jobs: uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false + fetch-depth: 2 - name: Detect changes that can affect the generated types id: changes - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ github.event.pull_request.number }} run: | set -euo pipefail - files="$(gh api --paginate "repos/$REPO/pulls/$PR_NUMBER/files" --jq '.[].filename')" - if printf '%s\n' "$files" | grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)'; then + if ! base="$(git rev-parse --verify --quiet HEAD^2 >/dev/null && git rev-parse HEAD^1)"; then + echo "Not a pull request merge commit, running the full check." + echo "relevant=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + files="$(git diff --name-only "$base" HEAD)" + if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)' <<< "$files"; then echo "relevant=true" >> "$GITHUB_OUTPUT" else + echo "No proxy, types or generator changes in this pull request, nothing to verify." echo "relevant=false" >> "$GITHUB_OUTPUT" fi From 860e37597fd20e2fe62c952d674c8b5c0ebddae9 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:56:29 +0000 Subject: [PATCH 38/74] fix(a2a): align agent list annotation and test with the tuple return type PR #36020 changed AgentRegistry.get_agent_list to return tuple[AgentResponse, ...], and PR #35163 added a test asserting the result equals []. Both were green on their own branches and only collided once they were both on litellm_internal_staging, so proxy-server has been failing on every PR since with 'assert () == []'. Nothing user-facing was wrong: get_agents only iterates the result and rebuilds it with comprehensions, and FastAPI serializes a tuple to the same JSON array. The test expectation was simply stale, so it now compares against (). The get_agents local was still annotated list[AgentResponse] while two branches assign the registry tuple straight into it, so it widens to Sequence[AgentResponse]. That covers both the tuple and the list branches without pretending the value is mutable. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/endpoints.py | 2 +- tests/test_litellm/proxy/proxy_server/test_proxy_config.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 5bc5b5566f8..d348bc01153 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -253,7 +253,7 @@ async def get_agents( ) try: - returned_agents: list[AgentResponse] = [] + returned_agents: Sequence[AgentResponse] = () # Admin users get all agents if ( 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 30156849628..91a7e1bc2c2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2638,4 +2638,4 @@ async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembe assert clean_agent_registry.config_agents == () clean_agent_registry.load_agents_from_db_and_config(db_agents=None) - assert clean_agent_registry.get_agent_list() == [] + assert clean_agent_registry.get_agent_list() == () From 2f70f4edbaf7aa042bac1b097bae5199774ad56b Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 21:30:09 +0000 Subject: [PATCH 39/74] build(deps): bump nanoid to 3.3.17 in the dashboard lockfile GHSA-2v37-7h3g-55p8 (CVE-2026-67213) rates 8.2 against nanoid 3.3.16 and reds osv-scan on every PR into staging. Custom generators can loop indefinitely when size is zero, so a caller that passes through a zero size hangs the process. nanoid is transitive through the dashboard's toolchain and 3.3.17 is a patch release that published 08-03, so it already clears the .npmrc min-release-age=3 guard. The lock diff is the version, resolved url, and integrity hash for that one package. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/package-lock.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 34dea2a39f5..515a992bc85 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -10319,9 +10319,9 @@ "license": "MIT" }, "node_modules/nanoid": { - "version": "3.3.16", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.16.tgz", - "integrity": "sha512-bzlKTyNJ7+LdGIIwy8ijFpIqEQIvafahV7eYykJ8Cvh42EdJeODoJ6gUJXpQJvej1BddH8OqTXZNE/KfbWAu8Q==", + "version": "3.3.17", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.17.tgz", + "integrity": "sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==", "funding": [ { "type": "github", From 8db2fbaad0c064b61e8f33d0c21c06d36e63d05d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 15:53:59 -0700 Subject: [PATCH 40/74] feat(ui): show user email or alias in usage data export (#36232) * feat(ui): resolve user email/alias in usage export instead of raw user id * test(ui): cover email/alias resolution in usage export data builders * chore(ui): drop explanatory comment per repo comment policy --- .../EntityUsageExport/utils.test.ts | 172 ++++++++++++++++++ .../src/components/EntityUsageExport/utils.ts | 26 ++- 2 files changed, 189 insertions(+), 9 deletions(-) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index e91b7b73a1e..2f551176091 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -2197,4 +2197,176 @@ describe("EntityUsageExport utils", () => { }); }); }); + + describe("display name resolution from entity metadata", () => { + const entityMetrics = { + spend: 12.25, + api_requests: 40, + successful_requests: 39, + failed_requests: 1, + total_tokens: 900, + prompt_tokens: 500, + completion_tokens: 400, + cache_read_input_tokens: 20, + cache_creation_input_tokens: 10, + }; + + const makeSpendData = (entity: string, metadata?: Record): EntitySpendData => ({ + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + [entity]: { + metrics: entityMetrics, + metadata, + api_key_breakdown: { + key1: { + metrics: entityMetrics, + metadata: { key_alias: "prod-key" }, + }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }); + + it("should export the user email as the entity label and keep the raw user id in the id column", () => { + const result = generateDailyData( + makeSpendData("user-123", { user_email: "ada@example.com", user_alias: "Ada" }), + "User", + ); + + expect(result).toHaveLength(1); + expect(result[0]["User"]).toBe("ada@example.com"); + expect(result[0]["User ID"]).toBe("user-123"); + }); + + it("should fall back to the user alias when the user has no email", () => { + const nullEmail = generateDailyData( + makeSpendData("user-123", { user_email: null, user_alias: "Ada Lovelace" }), + "User", + ); + const missingEmail = generateDailyData(makeSpendData("user-123", { user_alias: "Ada Lovelace" }), "User"); + + expect(nullEmail[0]["User"]).toBe("Ada Lovelace"); + expect(missingEmail[0]["User"]).toBe("Ada Lovelace"); + }); + + it("should fall back to the raw entity key when the entity carries no metadata", () => { + const noMetadata = generateDailyData(makeSpendData("my-tag"), "Tag"); + const emptyMetadata = generateDailyData(makeSpendData("customer-9", {}), "Customer"); + const blankNames = generateDailyData(makeSpendData("user-123", { user_email: null, user_alias: null }), "User"); + + expect(noMetadata[0]["Tag"]).toBe("my-tag"); + expect(emptyMetadata[0]["Customer"]).toBe("customer-9"); + expect(blankNames[0]["User"]).toBe("user-123"); + }); + + it("should prefer the team alias map over any alias in entity metadata", () => { + const result = generateDailyData( + makeSpendData("team-1", { team_alias: "Stale Alias", user_email: "ada@example.com" }), + "Team", + mockTeamAliasMap, + ); + + expect(result[0]["Team"]).toBe("Team One"); + }); + + it("should use the team alias from entity metadata when the alias map has no entry for the team", () => { + const result = generateDailyData( + makeSpendData("team-9", { team_alias: "Team Nine", user_email: "ada@example.com" }), + "Team", + mockTeamAliasMap, + ); + + expect(result[0]["Team"]).toBe("Team Nine"); + }); + + it("should resolve metadata.alias to the user email in getEntityBreakdown", () => { + const withEmail = getEntityBreakdown( + makeSpendData("user-123", { user_email: "ada@example.com", user_alias: "Ada" }), + ); + const withoutEmail = getEntityBreakdown(makeSpendData("user-123", { user_alias: "Ada" })); + + expect(withEmail[0].metadata.alias).toBe("ada@example.com"); + expect(withEmail[0].metadata.id).toBe("user-123"); + expect(withoutEmail[0].metadata.alias).toBe("Ada"); + }); + + it("should resolve the user email on every key row of the keys scope", () => { + const spendData: EntitySpendData = { + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + "user-123": { + metrics: entityMetrics, + metadata: { user_email: "ada@example.com", user_alias: "Ada" }, + api_key_breakdown: { + key1: { metrics: entityMetrics, metadata: { key_alias: "prod-key" } }, + key2: { metrics: entityMetrics, metadata: { key_alias: "dev-key" } }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }; + + const result = generateDailyWithKeysData(spendData, "User"); + + expect(result).toHaveLength(2); + expect(result.map((r) => r["User"])).toEqual(["ada@example.com", "ada@example.com"]); + expect(result.map((r) => r["User ID"])).toEqual(["user-123", "user-123"]); + expect(result.find((r) => r["Key ID"] === "key1")?.["Key Alias"]).toBe("prod-key"); + expect(result.find((r) => r["Key ID"] === "key2")?.["Key Alias"]).toBe("dev-key"); + }); + + it("should resolve each entity's own email in the models scope", () => { + const spendData: EntitySpendData = { + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + "user-a": { + metrics: entityMetrics, + metadata: { user_email: "ada@example.com", user_alias: "Ada" }, + api_key_breakdown: { key1: { metrics: entityMetrics, metadata: {} } }, + }, + "user-b": { + metrics: entityMetrics, + metadata: { user_email: null, user_alias: "Grace" }, + api_key_breakdown: { key2: { metrics: entityMetrics, metadata: {} } }, + }, + }, + models: { + "claude-sonnet-4-5": { + metrics: entityMetrics, + api_key_breakdown: { + key1: { metrics: entityMetrics, metadata: {} }, + key2: { metrics: entityMetrics, metadata: {} }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }; + + const result = generateDailyWithModelsData(spendData, "User"); + + expect(result).toHaveLength(2); + expect(result.every((r) => r.Model === "claude-sonnet-4-5")).toBe(true); + expect(result.find((r) => r["User ID"] === "user-a")?.["User"]).toBe("ada@example.com"); + expect(result.find((r) => r["User ID"] === "user-b")?.["User"]).toBe("Grace"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 8ff389ded5e..53d100040d7 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -3,12 +3,18 @@ import type { DateRangePickerValue } from "@tremor/react"; import Papa from "papaparse"; import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; -// Resolve display name for an entity. For teams the teamAliasMap provides -// a human-readable alias; for every other entity type the entity key itself -// (tag name, org id, customer id, …) is already the correct label. -const resolveEntityDisplay = (entity: string, teamAliasMap: Record): { id: string; alias: string } => ({ +const resolveEntityDisplay = ( + entity: string, + teamAliasMap: Record, + entityMetadata?: Record, +): { id: string; alias: string } => ({ id: entity, - alias: teamAliasMap[entity] || entity, + alias: + teamAliasMap[entity] || + entityMetadata?.team_alias || + entityMetadata?.user_email || + entityMetadata?.user_alias || + entity, }); // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). @@ -68,7 +74,7 @@ export const getEntityBreakdown = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); if (!entitySpend[entity]) { entitySpend[entity] = { @@ -113,7 +119,7 @@ export const generateDailyData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); dailyBreakdown.push({ Date: day.date, @@ -164,7 +170,7 @@ export const generateDailyWithKeysData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap); + const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); const apiKeyBreakdown = data.api_key_breakdown || {}; // Iterate through each API key in the breakdown @@ -241,11 +247,13 @@ export const generateDailyWithModelsData = ( spendData.results.forEach((day) => { const dailyEntityModels: { [key: string]: { [key: string]: any } } = {}; + const dailyEntityMetadata: { [key: string]: Record | undefined } = {}; Object.entries(resolveEntities(day.breakdown)).forEach(([entity, entityData]: [string, any]) => { if (!dailyEntityModels[entity]) { dailyEntityModels[entity] = {}; } + dailyEntityMetadata[entity] = entityData.metadata; Object.entries(day.breakdown.models || {}).forEach(([model, modelData]: [string, any]) => { const entityApiKeys = entityData.api_key_breakdown || {}; @@ -282,7 +290,7 @@ export const generateDailyWithModelsData = ( }); Object.entries(dailyEntityModels).forEach(([entity, models]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, dailyEntityMetadata[entity]); Object.entries(models).forEach(([model, metrics]: [string, any]) => { dailyModelBreakdown.push({ From 3238ce840601c820f31b01f1731b704a6d3f8f2e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 7 Aug 2026 17:03:34 -0700 Subject: [PATCH 41/74] feat(auto-router): track turns per complexity tier (LIT-5302) (#36209) * feat(auto-router): track turns per complexity tier (LIT-5302) Stamps complexity tier at decision time (rollup never re-derives from routed model, since tier->model mapping is mutable config). Records per-tier turn counts in LiteLLM_AutoRouterSession.tier_turns (jsonb), rolls up per router in benchmarks SQL via jsonb_object_agg, returns on AutoRouterBenchmarkGroup for dashboard turns/share metrics. Addresses Greptile/Bugbot findings: - Missing _SessionAggRow.tier_turns field: added with field_validator to parse jsonb text cast and handle NULL. Would 500 every benchmarks read. - Missing ::text cast on tier parameter: Postgres fails type inference on parameterized CASE/IS NULL without explicit cast. Added to all usages. - Docstring false claim (only complexity routers produce tiers): quality router stamps numeric tier '1'/'2'/'3'. Per-type grouping in SQL prevents cross-contamination. Rewrote docstring to clarify isolation. - Comment convention violations: stripped per CLAUDE.md rule. - Test gaps: 8 unit tests for extraction/validation/aggregation, 7 behavior tests for SQL semantics against real Postgres. 12 mutations killed. Fixed fragile complexity_router test that broke on nested function calls. No API change; extends existing GET /auto_router/benchmarks response only. Co-Authored-By: Claude * fix(auto-router): address review findings on tier turns tracking - Guard router_type update so a mid-session reconfigure can't pool foreign tier names into tier_turns - Keep pinned turns attributed to the tier that actually serves them - Drop stray -- AlterTable comment from hand-written migration - Drop the now-unnecessary ::text/json.loads round-trip; prisma already returns tier_turns as a parsed dict Co-Authored-By: Claude * fix(auto-router): satisfy type-discipline lint gate - tier_turns fields: dict[str, int] -> Mapping[str, int] (LIT001, mutable collection in annotation); these are read-only after construction - _summed_agg_row: {} -> MappingProxyType({}) (LIT002, mutable dict literal) - default-fallback branch: replace the reassigned-without-Final fallback_tier with a Final default_model_first flag and a single ternary assignment (LIT010) Verified locally: type_discipline_gate.py, ruff_strict_gate.py, and type_check_gate.py all pass against the litellm_internal_staging merge-base; full test_complexity_router.py (374), auto_router management-endpoint tests (26), db-layer rollup tests (31), and the live-Postgres proxy_behavior rollup suite (17) all pass. Co-Authored-By: Claude --------- Co-authored-by: Claude --- .../migration.sql | 1 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/db/autorouter_session_rollup.py | 13 ++- .../auto_router_endpoints.py | 30 +++++- litellm/proxy/schema.prisma | 1 + .../complexity_router/complexity_router.py | 6 +- .../auto_router_endpoints.py | 11 ++ schema.prisma | 1 + .../spend/test_autorouter_session_rollup.py | 102 ++++++++++++++++++ .../db/test_autorouter_session_rollup.py | 22 +++- .../test_auto_router_endpoints.py | 39 +++++++ .../router_strategy/test_complexity_router.py | 56 ++++++++-- 12 files changed, 268 insertions(+), 15 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql new file mode 100644 index 00000000000..81b1cbc7ec3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index da1652cdb61..b1f074c26b7 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -49,6 +49,7 @@ class AutoRouterTurnTransaction: cache_hit: bool cache_ttl_seconds: int | None cache_touched: bool + tier: str | None = None class TurnCacheFacts(NamedTuple): @@ -152,11 +153,13 @@ def build_autorouter_turn_transaction( return None usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) + tier_raw: Final = routing_decision.get("tier") return AutoRouterTurnTransaction( api_key=api_key, session_id=_bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), + tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None, model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), @@ -184,6 +187,8 @@ _COVERED: Final = _p("covered") _CACHE_HIT: Final = _p("cache_hit") _CACHE_TTL: Final = _p("cache_ttl_seconds") _TOUCHED: Final = _p("cache_touched") +_TIER: Final = f"{_p('tier')}::text" +_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)" _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at" _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}" @@ -201,7 +206,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( last_model, models, turns, unordered_turns, covered_turns, cache_hits, same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -211,7 +216,8 @@ VALUES ( 0, 0, 0, 0, (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), - {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8 + {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, + {_TIER_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, @@ -242,6 +248,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END) )), last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END), + tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")} + THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1) + ELSE t.tier_turns END), first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at), last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at) """ diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 141094f4d4c..d4221845b0c 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -4,8 +4,9 @@ AUTO ROUTER MANAGEMENT ENDPOINTS POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final from pydantic import BaseModel, TypeAdapter @@ -260,6 +261,7 @@ async def preview_auto_router_routing( class _SessionAggRow(BaseModel): router_name: str router_type: str + tier_turns: Mapping[str, int] sessions: int turns: int unordered_turns: int @@ -284,6 +286,23 @@ class _SessionAggRow(BaseModel): _SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow]) _BENCHMARKS_SQL: Final = """ +WITH windowed AS ( + SELECT * FROM "LiteLLM_AutoRouterSession" + WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +), +tier_maps AS ( + SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns + FROM ( + SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns + FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv + GROUP BY router_name, router_type, kv.key + ) per_tier + GROUP BY router_name, router_type +) +SELECT + agg.*, + COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns +FROM ( SELECT router_name, router_type, @@ -306,10 +325,11 @@ SELECT COALESCE(SUM(spend), 0)::float8 AS spend, COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds -FROM "LiteLLM_AutoRouterSession" -WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +FROM windowed GROUP BY router_name, router_type -ORDER BY SUM(spend) DESC +) agg +LEFT JOIN tier_maps USING (router_name, router_type) +ORDER BY agg.spend DESC """ @@ -366,6 +386,7 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: return _SessionAggRow( router_name="", router_type="", + tier_turns=MappingProxyType({}), sessions=sum(row.sessions for row in rows), turns=sum(row.turns for row in rows), unordered_turns=sum(row.unordered_turns for row in rows), @@ -443,6 +464,7 @@ async def get_auto_router_benchmarks( AutoRouterBenchmarkGroup( router_name=row.router_name, router_type=row.router_type, + tier_turns=row.tier_turns, **_benchmark_totals(row).model_dump(), ) for row in rows diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a69509fc37a..f6ced0bb9d2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1730,6 +1730,7 @@ class ComplexityRouter(CustomLogger): routing_decision=self._build_routing_decision( routed_model=routed_model, cause=cause, + tier=self._tier_for_model(routed_model), escalation_keyword=pin_escalation_keyword, escalated=escalated, conversation_continuing=conversation_continuing, @@ -1797,7 +1798,8 @@ class ComplexityRouter(CustomLogger): if user_message is None: verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model") - if not self.config.plugins and self.config.default_model: + default_model_first: Final = not self.config.plugins and self.config.default_model + if default_model_first: # No plugins configured: preserve the pre-existing default_model-first # priority exactly (changing it would be a silent behavior change for # every non-plugin user, not just a security fix). @@ -1809,12 +1811,14 @@ class ComplexityRouter(CustomLogger): routed_model = await self._pick_model_for_tier( ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) + fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, routing_decision=self._build_routing_decision( routed_model=routed_model, cause="default_fallback", + tier=fallback_tier, conversation_continuing=conversation_continuing, ), ) diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 6c8fb96a729..6626dea6849 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -2,6 +2,7 @@ Types for auto-router management endpoints """ +from collections.abc import Mapping from typing import Final from pydantic import BaseModel, Field, field_validator @@ -120,6 +121,16 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals): router_name: str = Field(description="The auto-router alias requests were sent to") router_type: str = Field(description="complexity, adaptive or quality") + tier_turns: Mapping[str, int] = Field( + default_factory=dict, + description="Turns per tier, keyed by the tier name the routing decision recorded at " + "request time (never re-derived at read time, since the tier-to-model mapping is " + "mutable config). Tier names are scoped to this group's router_type and are not " + "comparable across types: a complexity router reports 'simple'/'medium'/'complex'/" + "'reasoning', a quality router reports its numeric quality tier, and an adaptive router " + "records no tier at all. Turns no tier served (the classifier fell back to default_model) " + "are absent rather than pooled under a sentinel key, so the values may sum to less than turns", + ) class AutoRouterBenchmarksResponse(BaseModel): diff --git a/schema.prisma b/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index aa734ee22cc..65b70f13a3b 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -38,11 +38,13 @@ async def _turn( tokens: int = 100, spend: float = 0.01, saved: float = 0.02, + tier: "str | None" = None, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( UPSERT_AUTOROUTER_SESSION_SQL, key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched, + tier, ) @@ -195,6 +197,106 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db assert [(row["router_type"], row["sessions"]) for row in matching] == [("complexity", 1), ("quality", 1)] +async def test_tier_turns_count_each_tier_that_served_a_turn(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, tier="simple") + await _turn(db, key, "B", T0 + timedelta(seconds=10), tier="complex") + await _turn(db, key, "A", T0 + timedelta(seconds=20), tier="simple") + + assert (await _row(db, key))["tier_turns"] == {"simple": 2, "complex": 1} + + +async def test_an_untiered_turn_increments_no_tier_counter(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, tier=None) + assert (await _row(db, key))["tier_turns"] == {} + + await _turn(db, key, "A", T0 + timedelta(seconds=10), tier="medium") + await _turn(db, key, "A", T0 + timedelta(seconds=20), tier=None) + row = await _row(db, key) + assert row["tier_turns"] == {"medium": 1} + assert row["turns"] == 3 + + +async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_the_map(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, router_type="complexity", tier="medium") + await _turn(db, key, "A", T0 + timedelta(seconds=10), router_type="quality", tier="2") + await _turn(db, key, "A", T0 + timedelta(seconds=20), router_type="complexity", tier="medium") + + row = await _row(db, key) + assert row["tier_turns"] == {"medium": 2} + assert row["turns"] == 3 + + +async def test_an_out_of_order_turn_still_counts_toward_its_tier(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple") + await _turn(db, key, "A", T0, tier="simple") + + row = await _row(db, key) + assert row["tier_turns"] == {"simple": 2} + assert row["unordered_turns"] == 1 + + +async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier="simple") + await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, tier="simple") + await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{uuid.uuid4()}", router=router, tier="complex") + await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + grouped = next(row for row in rows if row["router_name"] == router) + assert grouped["tier_turns"] == {"simple": 2, "complex": 1} + assert grouped["turns"] == 4 + + +async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn( + db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity", tier="medium" + ) + await _turn( + db, + key, + "A", + T0 + timedelta(seconds=10), + session_id=f"s-{uuid.uuid4()}", + router=router, + router_type="quality", + tier="2", + ) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} + assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} + + +async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + grouped = next(row for row in rows if row["router_name"] == router) + assert grouped["tier_turns"] == {} + + async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db): key = f"k-{uuid.uuid4()}" await _turn(db, key, "A", T0, ttl=300) diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index aa7d01bc880..0df11f224a2 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -93,6 +93,19 @@ class TestBuildTransaction: def test_requests_without_a_routing_decision_are_skipped(self, metadata: dict): assert _build(metadata=metadata) is None + def test_the_tier_the_decision_recorded_is_carried_onto_the_transaction(self): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": "reasoning"})) + assert transaction is not None and transaction.tier == "reasoning" + + @pytest.mark.parametrize("tier", [None, "", 3, {"tier": "medium"}]) + def test_a_decision_without_a_usable_tier_records_no_tier(self, tier: object): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": tier})) + assert transaction is not None and transaction.tier is None + + def test_a_decision_that_never_mentions_tier_records_no_tier(self): + transaction = _build() + assert transaction is not None and transaction.tier is None + def test_router_name_falls_back_to_the_payload_model_group(self): transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"})) assert transaction is not None and transaction.router_name == "live-auto" @@ -167,7 +180,11 @@ class _FakeClient: self.db = _FakeDB(failures, poison_session) -def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0)) -> AutoRouterTurnTransaction: +def _transaction( + session_id: str = "s1", + at: datetime = datetime(2026, 8, 1, 12, 0, 0), + tier: str | None = "medium", +) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( api_key="k1", session_id=session_id, @@ -182,6 +199,7 @@ def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, cache_hit=False, cache_ttl_seconds=None, cache_touched=False, + tier=tier, ) @@ -201,7 +219,7 @@ class TestFlush: assert sql == UPSERT_AUTOROUTER_SESSION_SQL assert params == ( "k1", "s1", "live-auto", "complexity", "bedrock/haiku", - "2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, + "2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, "medium", ) def test_a_connect_error_retries_the_same_statement(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 888db031515..3a995e27697 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -292,6 +292,7 @@ class TestAutoRouterBenchmarks: ROW = _SessionAggRow( router_name="live-auto", router_type="complexity", + tier_turns={}, sessions=4, turns=40, unordered_turns=1, @@ -377,6 +378,21 @@ class TestAutoRouterBenchmarks: assert totals.avg_turns_per_session == 10.0 assert totals.spend == 10.0 + def test_tier_names_stay_scoped_to_the_router_type_that_recorded_them(self): + quality = self.ROW.model_copy( + update={"router_name": "quality-auto", "router_type": "quality", "tier_turns": {"2": 7}} + ) + complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}}) + assert complexity.tier_turns == {"medium": 7} + assert quality.tier_turns == {"2": 7} + + def test_summed_totals_carry_no_tier_map_because_names_are_router_scoped(self): + from litellm.proxy.management_endpoints.auto_router_endpoints import _summed_agg_row + + quality = self.ROW.model_copy(update={"router_type": "quality", "tier_turns": {"2": 7}}) + complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}}) + assert _summed_agg_row([complexity, quality]).tier_turns == {} + @pytest.mark.asyncio async def test_non_admin_roles_cannot_read_benchmarks(self): from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -427,3 +443,26 @@ class TestAutoRouterBenchmarks: assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "wire_value, expected", [({"simple": 24, "complex": 16}, {"simple": 24, "complex": 16}), ({}, {})] + ) + async def test_the_tier_map_reaches_the_response_as_the_jsonb_column_returns_it( + self, wire_value: dict, expected: dict, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks + + class _DB: + async def query_raw(self, sql: str, *params: object): + return [{**TestAutoRouterBenchmarks.ROW.model_dump(), "tier_turns": wire_value}] + + monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})()) + + response = await get_auto_router_benchmarks( + user_api_key_dict=ADMIN, + start_date="2026-07-01", + end_date="2026-08-01", + ) + assert response.groups[0].tier_turns == expected diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8e9e32f5898..356556f3563 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3477,6 +3477,23 @@ class TestSessionAffinity: # Pinned to the first turn's model, not re-classified down to SIMPLE. assert second.model == "o1-preview" + @pytest.mark.asyncio + async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = self._request_kwargs("session-1") + await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + pinned = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert pinned.routing_decision["tier"] == "REASONING" + @pytest.mark.asyncio async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config): mock_router_instance.cache = DualCache() @@ -4362,7 +4379,24 @@ class TestRoutingDecisionContents: assert decision is not None assert decision["cause"] == "default_fallback" assert decision["routed_model"] == response.model - assert "tier" not in decision + assert decision.get("tier") == "MEDIUM" + + @pytest.mark.asyncio + async def test_a_default_model_fallback_claims_no_tier(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "default_model": "gpt-4o"}, + ) + response = await router.async_pre_routing_hook( + model="test-complexity-router", + request_kwargs={}, + messages=[{"role": "system", "content": "be nice"}], + ) + assert response is not None + assert response.routing_decision is not None + assert response.routing_decision["cause"] == "default_fallback" + assert "tier" not in response.routing_decision @pytest.mark.asyncio async def test_session_pin_decision(self, mock_router_instance, basic_config): @@ -5907,11 +5941,21 @@ class TestConversationShapeDiscriminator: ) builds = source.split("self._build_routing_decision(")[1:] assert builds - missing = [ - i - for i, block in enumerate(builds) - if "conversation_continuing=conversation_continuing" not in block.split("),")[0] - ] + missing = [] + for i, block in enumerate(builds): + depth = 0 + end = 0 + for j, char in enumerate(block): + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth < 0: + end = j + break + extracted = block[:end] + if "conversation_continuing=conversation_continuing" not in extracted: + missing.append(i) assert not missing, f"routing decisions {missing} do not carry the conversation shape" From e50a42051c531f8d95a0dbf905917c991f8ed8f5 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 7 Aug 2026 17:28:55 -0700 Subject: [PATCH 42/74] fix(websearch): restore snippet text in native web_search_tool_result blocks (LIT-5315) (#36228) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(websearch): restore snippet text in native web_search_tool_result blocks (LIT-5315) The build_web_search_tool_result_block method copied url/title/page_age but hardcoded encrypted_content to empty string, never reading SearchResult.snippet. This left every native block content-free, forcing clients to web_fetch each result to recover evidence—the reported symptom. The Anthropic spec carries page text only in encrypted_content (an opaque server-issued blob we cannot mint), so snippet is emitted as an additive key alongside the spec fields. encrypted_content stays empty rather than holding plaintext, which would assert encryption semantics that don't hold. The anthropic SDK's BaseModel sets extra='allow', so the additive snippet key survives SDK parsing. litellm has no typed model for web_search_result at all, so nothing drops it internally. Turn-2 replay behavior is unaffected: the empty encrypted_content already exists today. Tests: - Updated test_shape_with_results to assert snippet present - Added test_snippet_carried_for_every_result to cover multi-result ordering - Added test_missing_snippet_degrades_to_empty_string for edge case - Mutation check: reverting source-only yields 3 test failures, restored to 117 passed Fixes: LIT-5315 Co-Authored-By: Claude * fix(websearch): make synthesized web_search blocks replayable by native clients Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(websearch): flatten a resultless replayed search block so Bedrock accepts the next turn The flatten added for LIT-5315 bails when the replayed web_search_tool_result carries an empty content list, but that is exactly what the interceptor emits when a search legitimately returns nothing and when a search raises. The block survived into the outbound body, Bedrock rejected the tag, and the conversation died on the following turn just as it did before the flatten existed. An empty content list has no encrypted_content to respect and no evidence to preserve, so it flattens safely, and its paired server_tool_use goes with it. The rendered text now says so explicitly rather than emitting a bare header. Adds the multi-turn replay coverage that existed nowhere: the outbound Bedrock invoke body is asserted free of both block types, parametrized over the results-present and resultless cases, and built from the interceptor's own builder so the fixture cannot drift from what it emits. Resolves LIT-5320 * test(websearch): pin flatten idempotency for the agentic-loop re-entry The agentic loop re-enters the same /v1/messages entry point for its follow-up call and hands it the original client history, so the flatten runs again over already-flattened messages once per iteration. Bedrock always takes that path, since its config reports web search as natively handled and the short-circuit is skipped. A pass that appended the rendered text instead of replacing the block would duplicate the evidence on every iteration and re-ship the unsupported tag, and no existing single-pass test sees it. Mutation checked: keeping the original block alongside the rendered text fails this test on its own. --------- Co-authored-by: Claude Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Yassin Kortam --- .../websearch_interception/handler.py | 56 ++++-- .../websearch_interception/transformation.py | 10 + litellm/llms/anthropic/common_utils.py | 151 +++++++++++++- .../messages/handler.py | 3 + .../integrations/websearch_interception.py | 23 ++- .../test_websearch_native_blocks.py | 47 ++++- ...erimental_pass_through_messages_handler.py | 55 ++++++ .../anthropic/test_anthropic_common_utils.py | 186 ++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 65 ++++++ 9 files changed, 572 insertions(+), 24 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 71388134e98..9748db2dcd2 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -9,7 +9,7 @@ server-side using litellm router's search tools. import asyncio import math import uuid -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -37,6 +37,8 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, ) from litellm.types.integrations.websearch_interception import ( + AnthropicSearchQuery, + AnthropicServerToolUseBlock, WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues @@ -833,22 +835,48 @@ class WebSearchInterceptionLogger(CustomLogger): def _build_native_result_blocks( tool_calls: list[dict], structured_results: list[SearchResponse | None], - ) -> list[dict[str, object]]: - """Build one ``web_search_tool_result`` block per tool_call.""" - blocks: Final[list[dict[str, object]]] = [] - for i, tool_call in enumerate(tool_calls): - tool_use_id = tool_call.get("id") or "" - structured = structured_results[i] if i < len(structured_results) else None - blocks.append( - WebSearchTransformation.build_web_search_tool_result_block( - tool_use_id=tool_use_id, - search_response=structured, - ) + ) -> tuple[Mapping[str, object], ...]: + """ + Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call. + + The pair is what Anthropic's spec requires: a bare result block, or one + keyed by the model's ``toolu_...`` id instead of a ``srvtoolu_...`` one, + is rejected on replay ("String should match pattern '^srvtoolu_'") and + leaves native clients without a search to attach the sources to. + """ + return tuple( + block + for i, tool_call in enumerate(tool_calls) + for block in WebSearchInterceptionLogger._native_result_pair( + query=WebSearchInterceptionLogger._tool_call_query(tool_call), + search_response=structured_results[i] if i < len(structured_results) else None, ) - return blocks + ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: list[dict[str, object]]) -> Any: + def _tool_call_query(tool_call: Mapping[str, object]) -> str: + tool_input: Final = tool_call.get("input") + if not isinstance(tool_input, Mapping): + return "" + query: Final = tool_input.get("query") + return query if isinstance(query, str) else "" + + @staticmethod + def _native_result_pair( + query: str, + search_response: SearchResponse | None, + ) -> tuple[Mapping[str, object], Mapping[str, object]]: + tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" + return ( + AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(), + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id=tool_use_id, + search_response=search_response, + ), + ) + + @staticmethod + def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 795810a7c40..199ab020559 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -412,6 +412,15 @@ class WebSearchTransformation: block that should accompany the model's text reply when the original request used a native ``web_search_*`` tool. + The spec'd shape carries page text only in ``encrypted_content``, an + opaque server-issued blob that we cannot mint. Emitting the four spec + fields alone would drop the snippet entirely, leaving the client (and + the model, on any replayed follow-up turn) with URLs and titles but no + evidence to answer from, forcing a fetch per result. So the snippet is + carried in an additive ``snippet`` key alongside the spec fields. + ``encrypted_content`` stays empty rather than holding plaintext, which + would assert encryption semantics that do not hold. + Spec reference: https://docs.anthropic.com/en/api/web-search-tool @@ -438,6 +447,7 @@ class WebSearchTransformation: "title": title, "page_age": page_age, "encrypted_content": "", + "snippet": getattr(r, "snippet", "") or "", } ) return { diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 314cfef6d84..9aa5a4f465f 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -4,9 +4,12 @@ This file contains common utils for anthropic calls. import copy import re -from typing import Any, Final +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Any, Final, Literal import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -1057,6 +1060,152 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any return out +class _ReplayedSearchQuery(BaseModel): + model_config = ConfigDict(extra="allow") + + query: str = "" + + +class _ReplayedWebSearchResult(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["web_search_result"] + url: str = "" + title: str = "" + snippet: str = "" + encrypted_content: str = "" + + +class _ReplayedWebSearchToolResult(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["web_search_tool_result"] + tool_use_id: str + content: tuple[_ReplayedWebSearchResult, ...] + + +class _ReplayedServerToolUse(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["server_tool_use"] + id: str + input: _ReplayedSearchQuery = _ReplayedSearchQuery() + + +class _TextBlock(BaseModel): + type: Literal["text"] = "text" + text: str + + +_WEB_SEARCH_TOOL_RESULT_ADAPTER: Final = TypeAdapter(_ReplayedWebSearchToolResult) +_SERVER_TOOL_USE_ADAPTER: Final = TypeAdapter(_ReplayedServerToolUse) + + +def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchToolResult | None: + """ + The parsed block when it is a ``web_search_tool_result`` carrying no + ``encrypted_content``, else None for anything Anthropic itself issued. + + An empty ``content`` list is flattenable too. It is what the interceptor emits + when a search legitimately returns nothing and when a search raises, and it + carries neither evidence to preserve nor an ``encrypted_content`` to respect, + so leaving it in place only buys the 400 this whole function exists to avoid. + """ + try: + parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block) + except ValidationError: + return None + if any(result.encrypted_content for result in parsed.content): + return None + return parsed + + +def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None: + try: + return _SERVER_TOOL_USE_ADAPTER.validate_python(block) + except ValidationError: + return None + + +def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str: + header: Final = f"Web search results for '{query}':" if query else "Web search results:" + if not results: + return f"{header}\n\nNo results were returned." + body: Final = "\n\n".join( + "\n".join( + line + for line in ( + f"Title: {result.title}" if result.title else "", + f"URL: {result.url}" if result.url else "", + f"Snippet: {result.snippet}" if result.snippet else "", + ) + if line + ) + for result in results + ) + return f"{header}\n\n{body}" if body else header + + +def _rewrite_replayed_web_search_block( + block: object, + flattenable: Mapping[str, _ReplayedWebSearchToolResult], + queries: Mapping[str, str], +) -> object | None: + parsed_result: Final = _flattenable_web_search_tool_result(block) + if parsed_result is not None: + return _TextBlock( + text=_render_web_search_results(queries.get(parsed_result.tool_use_id, ""), parsed_result.content) + ).model_dump() + parsed_use: Final = _replayed_server_tool_use(block) + if parsed_use is not None and parsed_use.id in flattenable: + return None + return block + + +def _flatten_web_search_results_in_message(message: object) -> object: + if not isinstance(message, Mapping) or not isinstance(message.get("content"), Sequence): + return message + content: Final = message["content"] + if isinstance(content, str): + return message + flattenable: Final = MappingProxyType( + { + parsed.tool_use_id: parsed + for parsed in (_flattenable_web_search_tool_result(block) for block in content) + if parsed is not None + } + ) + if not flattenable: + return message + queries: Final = MappingProxyType( + { + parsed.id: parsed.input.query + for parsed in (_replayed_server_tool_use(block) for block in content) + if parsed is not None + } + ) + rewritten: Final = tuple(_rewrite_replayed_web_search_block(block, flattenable, queries) for block in content) + return {**message, "content": [b for b in rewritten if b is not None]} # mutable-ok: JSON wire format + + +def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: as sibling sanitizers + messages: list[Any], +) -> list[Any]: + """ + Return a new message list with replayed ``web_search_tool_result`` blocks that + carry no ``encrypted_content`` rewritten into plain ``text`` blocks holding the + same title / url / snippet evidence. + + ``encrypted_content`` is an opaque blob only Anthropic's own search backend can + mint, so blocks synthesized by LiteLLM (websearch interception against a search + provider) are rejected with ``Invalid encrypted_content in search_result block`` + when a native client loops them back as history. Flattening them keeps the + evidence in the conversation instead of 400ing the follow-up turn, and leaves + genuine Anthropic-issued blocks untouched. + """ + return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format + + def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 3ef298aa336..c4b5cc628e2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -14,6 +14,7 @@ from typing import Any, Final, cast import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, sanitize_tool_use_ids_in_anthropic_messages, strip_empty_text_blocks_from_anthropic_messages, ) @@ -222,6 +223,7 @@ async def anthropic_messages( # Replay of cross-provider tool history (e.g. kimi -> Anthropic) may carry # ids like ``functions.Bash:0`` that violate Anthropic's id pattern. messages = sanitize_tool_use_ids_in_anthropic_messages(messages) + messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages) from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, @@ -413,6 +415,7 @@ def anthropic_messages_handler( if not kwargs.pop("_litellm_messages_presanitized", False): messages = strip_empty_text_blocks_from_anthropic_messages(messages) messages = sanitize_tool_use_ids_in_anthropic_messages(messages) + messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages) from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 05537da67d7..90713b270be 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -2,7 +2,28 @@ Type definitions for WebSearch Interception integration. """ -from typing import TypedDict +from typing import Literal, TypedDict + +from pydantic import BaseModel + + +class AnthropicSearchQuery(BaseModel): + """``input`` of an Anthropic ``server_tool_use`` block for a web search.""" + + query: str + + +class AnthropicServerToolUseBlock(BaseModel): + """ + The ``server_tool_use`` block that must accompany a ``web_search_tool_result``. + + Anthropic requires the pair, with a ``srvtoolu_``-prefixed id shared by both. + """ + + type: Literal["server_tool_use"] = "server_tool_use" + id: str + name: Literal["web_search"] = "web_search" + input: AnthropicSearchQuery class WebSearchInterceptionConfig(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py index 544abab8dcf..c859f9b2f55 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -134,6 +134,30 @@ class TestBuildWebSearchToolResultBlock: assert first["title"] == "LiteLLM Docs" assert first["page_age"] == "2025-01-15" assert first["encrypted_content"] == "" + assert first["snippet"] == "Unified interface for LLMs." + + def test_snippet_carried_for_every_result(self): + # The snippet is the only field carrying page text. Losing it leaves the + # client and the model with nothing to answer from, forcing a fetch per + # result. + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=_make_search_response(), + ) + assert [r["snippet"] for r in block["content"]] == [ + "Unified interface for LLMs.", + "Pay-per-use pricing model.", + ] + + def test_missing_snippet_degrades_to_empty_string(self): + response = SearchResponse( + results=[SearchResult(title="T", url="https://x/", snippet="")] + ) + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=response, + ) + assert block["content"][0]["snippet"] == "" def test_handles_none_search_response(self): block = WebSearchTransformation.build_web_search_tool_result_block( @@ -223,11 +247,15 @@ class TestBuildPlanAttachesBlocks: ) blocks = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) - assert isinstance(blocks, list) - assert len(blocks) == 1 - assert blocks[0]["type"] == "web_search_tool_result" - assert blocks[0]["tool_use_id"] == "toolu_one" - assert blocks[0]["content"][0]["url"] == "https://docs.litellm.ai/" + assert isinstance(blocks, tuple) + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + ] + assert blocks[0]["id"].startswith("srvtoolu_") + assert blocks[0]["input"] == {"query": "what is litellm"} + assert blocks[1]["tool_use_id"] == blocks[0]["id"] + assert blocks[1]["content"][0]["url"] == "https://docs.litellm.ai/" @pytest.mark.asyncio async def test_metadata_does_not_carry_blocks_when_flag_absent(self): @@ -479,6 +507,9 @@ class TestLegacyPathMatchesNewPath: kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, ) - assert out["content"][0]["type"] == "web_search_tool_result" - assert out["content"][0]["tool_use_id"] == "toolu_legacy" - assert out["content"][1]["type"] == "text" + assert [b["type"] for b in out["content"]] == [ + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert out["content"][1]["tool_use_id"] == out["content"][0]["id"] 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 df3db3d2c57..f11324ca376 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 @@ -732,6 +732,61 @@ def test_handler_skips_strip_when_presanitized(): assert result is not None +def test_handler_flattens_replayed_unencrypted_web_search_results(): + """Synthesized search blocks replayed as history must reach the provider as text.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + captured = {} + + def fake_base_handler(*args, **kwargs): + captured.update(kwargs) + return "stub" + + with patch.object( + handler.base_llm_http_handler, + "anthropic_messages_handler", + side_effect=fake_base_handler, + ): + handler.anthropic_messages_handler( + max_tokens=10, + messages=[ + {"role": "user", "content": "latest litellm version?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "latest litellm version"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": None, + "encrypted_content": "", + "snippet": "Latest release v1.95.0", + } + ], + }, + ], + }, + {"role": "user", "content": "which version?"}, + ], + model="anthropic/claude-3-5-sonnet-20241022", + custom_llm_provider="anthropic", + ) + + replayed = captured["messages"][1]["content"] + assert [b["type"] for b in replayed] == ["text"] + assert "Snippet: Latest release v1.95.0" in replayed[0]["text"] + + def test_presanitized_flag_not_leaked_to_provider_params(): """The private sentinel must be popped, never forwarded as a request param.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler 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 6ab0f2c08ab..9df72108332 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -10,8 +10,10 @@ Verifies that: - ANTHROPIC_API_KEY / ANTHROPIC_API_BASE take precedence over their aliases. """ +import json import os import sys +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -1457,6 +1459,190 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_text_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["tool_result"] + def test_flatten_unencrypted_web_search_results_keeps_snippet_evidence(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + {"role": "user", "content": "latest litellm version?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "latest litellm version"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": None, + "encrypted_content": "", + "snippet": "Latest release v1.95.0", + } + ], + }, + {"type": "text", "text": "v1.95.0"}, + ], + }, + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + assert [b["type"] for b in out[1]["content"]] == ["text", "text"] + flattened = out[1]["content"][0]["text"] + assert "Web search results for 'latest litellm version':" in flattened + assert "URL: https://github.com/BerriAI/litellm/releases" in flattened + assert "Snippet: Latest release v1.95.0" in flattened + assert msgs[1]["content"][0]["type"] == "server_tool_use" + + @pytest.mark.parametrize("results", [[], None], ids=["empty_list", "search_raised"]) + def test_flatten_unencrypted_web_search_results_flattens_a_resultless_search(self, results): + """A search that found nothing, or that raised, still has to be flattened. + + Both cases reach the client as ``content: []``, and leaving that block in + place ships an unsupported tag to Bedrock on the next turn just as surely + as a populated one does. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=None if results is None else SimpleNamespace(results=results), + ) + assert block["content"] == [], "fixture drifted from what the interceptor emits" + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "who won"}, + }, + block, + {"type": "text", "text": "I could not find that."}, + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert [b["type"] for b in out[0]["content"]] == ["text", "text"] + assert out[0]["content"][0]["text"] == ("Web search results for 'who won':\n\nNo results were returned.") + + @pytest.mark.parametrize("results", [[SimpleNamespace(title="Rome", url="u", snippet="s", date=None)], []]) + def test_flatten_unencrypted_web_search_results_is_idempotent(self, results): + """Flattening twice must equal flattening once. + + The agentic loop re-enters the same entry point for its follow-up call and + hands it the original history, so this runs again on already-flattened + messages once per iteration. A pass that appended instead of replacing + would duplicate the evidence on every loop. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "when"}}, + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=SimpleNamespace(results=results), + ), + {"type": "text", "text": "753 BC."}, + ], + } + ] + + once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once) + + assert [b["type"] for b in once[0]["content"]] == ["text", "text"] + assert json.dumps(twice) == json.dumps(once) + + def test_flatten_unencrypted_web_search_results_preserves_real_anthropic_blocks(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "q"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://example.com", + "title": "Example", + "page_age": None, + "encrypted_content": "EqgfCioIARgBIiQ4", + } + ], + }, + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + + def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "web_search_tool_result_error", + "error_code": "max_uses_exceeded", + }, + } + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + def test_sanitize_tool_use_ids_in_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( sanitize_tool_use_ids_in_anthropic_messages, diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 3b8b4af78d9..76bb11cc26d 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -4,6 +4,7 @@ import json import os import sys from datetime import datetime +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -2515,3 +2516,67 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( assert thinking.get("type") == "enabled" assert isinstance(thinking.get("budget_tokens"), int) assert "output_config" not in flipped + + +@pytest.mark.parametrize( + "search_results, expected_evidence", + [ + pytest.param( + [SimpleNamespace(title="Rome", url="https://ex.com/rome", snippet="Founded 753 BC.", date=None)], + "Snippet: Founded 753 BC.", + id="search_returned_results", + ), + pytest.param([], "No results were returned.", id="search_returned_nothing"), + ], +) +def test_replayed_intercepted_search_turn_leaves_no_unsupported_block_for_bedrock(search_results, expected_evidence): + """A native client replaying an intercepted search turn must not 400 on Bedrock. + + ``websearch_interception`` hands Claude Desktop an Anthropic-native + ``server_tool_use`` + ``web_search_tool_result`` pair, and Anthropic's protocol + obliges the client to replay that assistant turn verbatim on every later turn. + Bedrock's Anthropic schema defines neither tag, so both have to be gone from the + outbound body by the time it is signed, with the search evidence carried forward + as text instead. Built from the real builder rather than a hand-written fixture + so the two cannot drift apart. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + from litellm.types.router import GenericLiteLLMParams + + replayed_turn = [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "when was Rome founded"}, + }, + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=SimpleNamespace(results=search_results), + ), + {"type": "text", "text": "Rome was founded in 753 BC."}, + ] + messages = [ + {"role": "user", "content": [{"type": "text", "text": "When was Rome founded?"}]}, + {"role": "assistant", "content": replayed_turn}, + {"role": "user", "content": [{"type": "text", "text": "Repeat the year."}]}, + ] + + body = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=flatten_unencrypted_web_search_results_in_anthropic_messages(messages), + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + serialized = json.dumps(body) + assert "web_search_tool_result" not in serialized + assert "server_tool_use" not in serialized + assert expected_evidence in serialized + assert "Rome was founded in 753 BC." in serialized From 1a45bf9afebfb26e656b09c16df54a850fd035e3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:45:30 -0700 Subject: [PATCH 43/74] fix(proxy): resolve entity access groups in the model listing endpoints (#36230) * fix(proxy): resolve entity access groups in the model listing endpoints Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): reuse the fetched team object when listing models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover key-level access group resolution in model listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: shivam Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 5 +- litellm/proxy/auth/model_checks.py | 7 +- litellm/proxy/utils.py | 136 +++++++++++++++--- .../proxy/auth/test_model_checks.py | 25 ++++ .../proxy/utils/helpers/test_model_access.py | 117 +++++++++++++++ 5 files changed, 262 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3da899a5610..d07ac0c5586 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,6 +13,7 @@ import asyncio import math import re import time +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from fastapi import HTTPException, Request, status @@ -2834,7 +2835,7 @@ async def get_org_object( async def _get_resources_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, @@ -2893,7 +2894,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index a6e5eb2a0a0..ff9211742f3 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,6 +1,7 @@ # What is this? ## Common checks for /v1/models and `/model/info` import copy +from collections.abc import Sequence from typing import Any, Final import litellm @@ -178,8 +179,8 @@ def get_team_models( def get_complete_model_list( - key_models: list[str], - team_models: list[str], + key_models: Sequence[str], + team_models: Sequence[str], proxy_model_list: list[str], user_model: str | None, infer_model_from_keys: bool | None, @@ -203,7 +204,7 @@ def get_complete_model_list( def append_unique(models): for model in models: - if model not in unique_models: + if model not in unique_models and model != SpecialModelNames.no_default_models.value: unique_models.append(model) if key_models: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 605455a2f73..e59c6adaf22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -165,6 +165,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -6419,6 +6420,74 @@ def construct_database_url_from_env_vars() -> str | None: return None +async def _get_validated_team_object( + user_api_key_dict: "UserAPIKeyAuth", + team_id: str, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", +) -> "LiteLLM_TeamTableCachedObj": + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.management_endpoints.team_endpoints import validate_membership + + team_object: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) + return team_object + + +async def _get_team_object_for_access_groups( + team_id: str | None, + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> Optional["LiteLLM_TeamTableCachedObj"]: + from litellm.proxy.auth.auth_checks import get_team_object + + if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: + return None + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id) + return None + + +async def _get_access_group_models( + user_api_key_dict: "UserAPIKeyAuth", + team_object: Optional["LiteLLM_TeamTableCachedObj"], + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> tuple[str, ...]: + from litellm.proxy.auth.auth_checks import ( + _get_models_from_access_groups, + get_authorized_resources_from_key_access_groups, + ) + + team_group_models: Final = await _get_models_from_access_groups( + access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + key_group_models: Final = await get_authorized_resources_from_key_access_groups( + valid_token=user_api_key_dict, + team_object=team_object, + resource_field="access_model_names", + ) + return tuple(dict.fromkeys((*team_group_models, *key_group_models))) + + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -6450,13 +6519,11 @@ async def get_available_models_for_user( Returns: List of model names available to the user """ - from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.model_checks import ( get_complete_model_list, get_key_models, get_team_models, ) - from litellm.proxy.management_endpoints.team_endpoints import validate_membership # Get proxy model list and access groups if llm_router is None: @@ -6466,31 +6533,33 @@ async def get_available_models_for_user( proxy_model_list = llm_router.get_model_names() model_access_groups = llm_router.get_model_access_groups() - # Get key models - key_models = get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - ) - - # Get team models - team_models: list[str] = user_api_key_dict.team_models - - # If specific team_id is provided, validate and get team models - if team_id and prisma_client and proxy_logging_obj and user_api_key_cache: - key_models = [] - team_object: Final = await get_team_object( + requested_team_object: Final = ( + await _get_validated_team_object( + user_api_key_dict=user_api_key_dict, team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) - team_models = team_object.models + if team_id and prisma_client and proxy_logging_obj and user_api_key_cache + else None + ) - team_models = get_team_models( - team_models=team_models, + key_models: Final[Sequence[str]] = ( + () + if requested_team_object is not None + else get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + ) + ) + + team_models: Final = get_team_models( + team_models=( + requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models + ), proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, @@ -6498,10 +6567,31 @@ async def get_available_models_for_user( effective_team_id: Final = team_id or user_api_key_dict.team_id + access_group_models: Final = ( + await _get_access_group_models( + user_api_key_dict=user_api_key_dict, + team_object=requested_team_object + or await _get_team_object_for_access_groups( + team_id=effective_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if key_models or team_models + else () + ) + + granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models + granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models + # Get complete model list all_models: Final = get_complete_model_list( - key_models=key_models, - team_models=team_models, + key_models=granted_key_models, + team_models=granted_team_models, proxy_model_list=proxy_model_list, user_model=user_model, infer_model_from_keys=general_settings.get("infer_model_from_keys", False), diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f56b8e113a9..e6c0eaee3c4 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -735,3 +735,28 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.vertex_language_models.discard(fake_model) litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_get_complete_model_list_drops_no_default_models_sentinel(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models", "model-a"], + team_models=[], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == ["model-a"] + + +def test_get_complete_model_list_sentinel_only_grants_nothing(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models"], + team_models=["no-default-models"], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == [] diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 59268e1427b..5fb4392eec6 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -9,6 +9,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, get_available_models_for_user, + hash_token, is_known_model, is_known_vector_store_index, model_dump_with_preserved_fields, @@ -404,3 +405,119 @@ async def test_get_available_models_for_user_error_path_complete_list_raises( general_settings={}, user_model=None, ) + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_team_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + team = LiteLLM_TeamTableCachedObj( + team_id="team-1", + models=["no-default-models"], + access_group_ids=["ag-1"], + ) + access_group = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="repro-group", + access_model_names=["model-a", "model-b"], + assigned_team_ids=["team-1"], + ) + + async def _get_team_object(**_kwargs): + return team + + async def _get_access_object(**_kwargs): + return access_group + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert sorted(result) == ["model-a", "model-b"] + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_without_access_groups_grants_nothing( + monkeypatch, +): + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == [] + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_key_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + async def _get_access_object(**_kwargs): + return LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="key-group", + access_model_names=["model-b"], + assigned_key_ids=[hash_token("sk-test-key")], + ) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["no-default-models"], + team_models=["no-default-models"], + access_group_ids=["ag-1"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == ["model-b"] From 2a9aac70045282b82c67b87b52881c8af0521db1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 17:45:50 -0700 Subject: [PATCH 44/74] fix(ui): let access groups be a team's only model source, with hover provenance (#36234) * feat(proxy): return per-group model provenance on /team/info /team/info now carries access_group_details, one entry per resolved access group with its id, name, and model list, so the UI can attribute each inherited model to the group granting it. The batch resolver returns the access group rows keyed by id instead of a stringly dict of lists, and the team member budget helper returns a copy instead of mutating its parameter. Type discipline and basedpyright budgets ratchet down accordingly. * feat(ui): allow group-only teams and show model provenance on hover Team create and edit no longer require a model selection: an empty selection is saved as the no-default-models sentinel, never as a bare empty list, since an empty team model list means unrestricted access. The team info Models card now renders every badge with a hover tooltip naming how the team got that model: directly, via named access groups, or both, and group-granted badges stay visible when the direct list is empty or a sentinel. * refactor(proxy): dedupe access group ids and return copies instead of mutating Duplicate access_group_ids no longer amplify the /team/info response: ids collapse order-preserving before provenance is built, pinned by a regression test. The resolver returns a model_copy rather than mutating its parameter, and the team create call sends a new object instead of reassigning formValues.models. Budgets ratchet down further with the mutation removal. --- basedpyright-code-budget.json | 6 +- litellm/proxy/_types.py | 7 ++ .../management_endpoints/team_endpoints.py | 88 ++++++++++--------- .../test_team_endpoints.py | 79 +++++++++++++++-- type-discipline-budget.json | 6 +- .../src/components/Teams.test.tsx | 28 ++++++ ui/litellm-dashboard/src/components/Teams.tsx | 11 +-- .../src/components/team/TeamInfo.tsx | 41 +++++---- .../components/team/teamModelAccess.test.ts | 86 ++++++++++++++++++ .../src/components/team/teamModelAccess.ts | 82 +++++++++++++++++ 10 files changed, 357 insertions(+), 77 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts create mode 100644 ui/litellm-dashboard/src/components/team/teamModelAccess.ts diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c8765eb0bd0..32b8eb3d4d0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45110 + "limit": 45098 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39838 + "limit": 39826 }, "reportUnknownParameterType": { "limit": 20237 }, "reportUnknownVariableType": { - "limit": 31383 + "limit": 31371 }, "reportUnnecessaryCast": { "limit": 122 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b7dc3a7c73..1fc05ac4653 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3894,12 +3894,19 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): ########################################## +class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): + access_group_id: str + access_group_name: str + models: tuple[str, ...] + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None + access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe5a0e06d2e..0b99879f9fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy._types import ( SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, + TeamAccessGroupModelGrant, TeamAddMemberResponse, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, @@ -3829,7 +3830,7 @@ async def _add_team_member_budget_table( ) -> TeamInfoResponseObjectTeamTable: try: team_budget: Final = await _budget_db(prisma_client).find_unique(where={"budget_id": team_member_budget_id}) - team_info_response_object.team_member_budget_table = team_budget + return team_info_response_object.model_copy(update={"team_member_budget_table": team_budget}) except Exception: verbose_proxy_logger.info( "Team member budget table not found, passed team_member_budget_id=%s", team_member_budget_id @@ -3838,21 +3839,34 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None: - """Populate access_group_models / mcp_server_ids / agent_ids on the team - info response by resolving inherited resources from its access groups.""" +async def _resolve_team_access_group_resources( + _team_info: TeamInfoResponseObjectTeamTable, +) -> TeamInfoResponseObjectTeamTable: + """Return a copy of the team info with access_group_models / mcp_server_ids / + agent_ids / details resolved from its access groups.""" if not _team_info.access_group_ids: - return + return _team_info ag_lookup: Final = await _batch_resolve_access_group_resources(_team_info.access_group_ids) - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in _team_info.access_group_ids: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - _team_info.access_group_models = list(models) - _team_info.access_group_mcp_server_ids = list(mcp_ids) - _team_info.access_group_agent_ids = list(agent_ids) + resolved_groups: Final = tuple( + ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup + ) + return _team_info.model_copy( + update={ + "access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}), + "access_group_mcp_server_ids": list( + {s for group in resolved_groups for s in (group.access_mcp_server_ids or [])} + ), + "access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}), + "access_group_details": tuple( + TeamAccessGroupModelGrant( + access_group_id=group.access_group_id, + access_group_name=group.access_group_name, + models=tuple(group.access_model_names or ()), + ) + for group in resolved_groups + ), + } + ) @router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -3958,11 +3972,11 @@ async def team_info( ) # Resolve resources inherited from access groups - await _resolve_team_access_group_resources(_team_info) + resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=_team_info, + team_info=resolved_team_info, keys=keys, team_memberships=returned_tm, ) @@ -4391,32 +4405,21 @@ async def _build_team_list_where_conditions( async def _batch_resolve_access_group_resources( all_access_group_ids: list[str], -) -> dict[str, dict[str, list[str]]]: +) -> dict[str, LiteLLM_AccessGroupTable]: """ - Batch-fetch access groups in a single DB query and return a per-group - resource map. - - Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}. - Missing/invalid groups are silently omitted. + Batch-fetch access groups in a single DB query and return them keyed by + access_group_id. Missing/invalid groups are silently omitted. """ from litellm.proxy.proxy_server import prisma_client as _prisma_client if not all_access_group_ids or _prisma_client is None: return {} - unique_ids: Final = list(set(all_access_group_ids)) + unique_ids: Final = tuple(frozenset(all_access_group_ids)) rows: Final = await _access_group_db(_prisma_client).find_many( where={"access_group_id": {"in": unique_ids}}, ) - - result: Final[dict[str, dict[str, list[str]]]] = {} - for row in rows: - result[row.access_group_id] = { - "models": list(row.access_model_names or []), - "mcp_server_ids": list(row.access_mcp_server_ids or []), - "agent_ids": list(row.access_agent_ids or []), - } - return result + return {row.access_group_id: row for row in rows} def _convert_teams_to_response_models( @@ -4710,15 +4713,18 @@ async def list_team_v2( all_ag_ids: Final = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])] ag_lookup: Final = await _batch_resolve_access_group_resources(all_ag_ids) for team_item in team_items_with_ag: - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in team_item.access_group_ids or []: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - team_item.access_group_models = list(models) - team_item.access_group_mcp_server_ids = list(mcp_ids) - team_item.access_group_agent_ids = list(agent_ids) + team_groups = tuple( + ag_lookup[ag_id] for ag_id in (team_item.access_group_ids or []) if ag_id in ag_lookup + ) + team_item.access_group_models = list( + {m for group in team_groups for m in (group.access_model_names or [])} + ) + team_item.access_group_mcp_server_ids = list( + {s for group in team_groups for s in (group.access_mcp_server_ids or [])} + ) + team_item.access_group_agent_ids = list( + {a for group in team_groups for a in (group.access_agent_ids or [])} + ) return { "teams": team_list, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index eedffa1ea5f..a1cbc77b7a5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -8637,9 +8637,9 @@ class TestBatchResolveAccessGroupResources: with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1"]) - assert sorted(result["ag-1"]["models"]) == ["claude-3", "gpt-4"] - assert result["ag-1"]["mcp_server_ids"] == ["mcp-1"] - assert sorted(result["ag-1"]["agent_ids"]) == ["agent-1", "agent-2"] + assert sorted(result["ag-1"].access_model_names) == ["claude-3", "gpt-4"] + assert result["ag-1"].access_mcp_server_ids == ["mcp-1"] + assert sorted(result["ag-1"].access_agent_ids) == ["agent-1", "agent-2"] @pytest.mark.asyncio async def test_multiple_access_groups(self): @@ -8668,8 +8668,8 @@ class TestBatchResolveAccessGroupResources: with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"]) - assert result["ag-1"]["models"] == ["gpt-4"] - assert result["ag-2"]["models"] == ["gemini"] + assert result["ag-1"].access_model_names == ["gpt-4"] + assert result["ag-2"].access_model_names == ["gemini"] @pytest.mark.asyncio async def test_missing_access_group_omitted(self): @@ -8735,6 +8735,75 @@ class TestBatchResolveAccessGroupResources: assert "ag-1" in result +class TestResolveTeamAccessGroupResources: + """Tests for the per-team access group resolution on /team/info.""" + + @pytest.mark.asyncio + async def test_populates_flat_lists_and_per_group_details(self): + """access_group_details must attribute each model to the group granting it, + so the UI can show provenance on hover; flat lists stay for back-compat. + Duplicated ids must collapse to one entry (response amplification), and the + input object must stay untouched (resolution returns a copy).""" + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_group_name = "shared-models" + row1.access_model_names = ["gpt-4", "claude-3"] + row1.access_mcp_server_ids = ["mcp-1"] + row1.access_agent_ids = [] + + row2 = MagicMock() + row2.access_group_id = "ag-2" + row2.access_group_name = "extra-models" + row2.access_model_names = ["claude-3", "gemini"] + row2.access_mcp_server_ids = [] + row2.access_agent_ids = ["agent-1"] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( + return_value=[row1, row2] + ) + + team_info = TeamInfoResponseObjectTeamTable( + team_id="team-1", access_group_ids=["ag-1", "ag-2", "ag-1", "ag-missing"] + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resolved = await _resolve_team_access_group_resources(team_info) + + assert team_info.access_group_details is None + assert sorted(resolved.access_group_models or []) == [ + "claude-3", + "gemini", + "gpt-4", + ] + assert resolved.access_group_mcp_server_ids == ["mcp-1"] + assert resolved.access_group_agent_ids == ["agent-1"] + assert [ + (d.access_group_id, d.access_group_name, d.models) + for d in (resolved.access_group_details or []) + ] == [ + ("ag-1", "shared-models", ("gpt-4", "claude-3")), + ("ag-2", "extra-models", ("claude-3", "gemini")), + ] + + @pytest.mark.asyncio + async def test_no_access_groups_leaves_details_unset(self): + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + team_info = TeamInfoResponseObjectTeamTable(team_id="team-1", access_group_ids=[]) + resolved = await _resolve_team_access_group_resources(team_info) + + assert resolved.access_group_details is None + assert resolved.access_group_models is None + + @pytest.mark.asyncio async def test_verify_team_access_denies_unauthorized_user(): """ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 991f8eaa934..0a0cfe9a617 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23245 + "limit": 23235 }, "LIT002": { - "limit": 27179 + "limit": 27176 }, "LIT003": { "limit": 269 @@ -30,6 +30,6 @@ "limit": 16769 }, "LIT011": { - "limit": 5602 + "limit": 5598 } } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 2bda0f72cec..e8331294972 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -613,6 +613,34 @@ describe("Teams - access_group_ids in team create", () => { ); }); }); + + it("creates a team with no models selected, sending the no-default-models sentinel instead of an empty list", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Group Only Team" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Group Only Team", + models: ["no-default-models"], + }), + ); + }); + }); }); describe("Teams - metadata key-value pairs in team create", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 95642b93019..42ff8bcc4c2 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -42,6 +42,7 @@ interface TeamProps { import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { teamCreateCall } from "./networking"; +import { normalizeTeamModelSelection } from "./team/teamModelAccess"; import { ModelSelect } from "./ModelSelect/ModelSelect"; const canCreateOrManageTeams = ( @@ -351,7 +352,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } } - await teamCreateCall(accessToken, formValues); + await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) }); NotificationsManager.success("Team created"); await refreshTeams(); form.resetFields(); @@ -618,17 +619,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser label={ Models{" "} - + } - rules={[ - { - required: true, - message: "Please select at least one model", - }, - ]} name="models" > = new Set([ "disable_global_guardrails", ]); +const TEAM_MODEL_BADGE_COLORS: Record = { + "all-proxy": "red", + "no-default": "gray", + direct: "blue", + "access-group": "green", +}; + export interface TeamMembership { user_id: string; team_id: string; @@ -132,6 +145,7 @@ export interface TeamData { access_group_models?: string[]; access_group_mcp_server_ids?: string[]; access_group_agent_ids?: string[]; + access_group_details?: TeamAccessGroupModelGrant[]; router_settings?: Record; guardrails?: string[]; policies?: string[]; @@ -483,7 +497,7 @@ const TeamInfoView: React.FC = ({ const updateData: any = { team_id: teamId, team_alias: values.team_alias, - models: values.models, + models: normalizeTeamModelSelection(values.models), tpm_limit: sanitizeNumeric(values.tpm_limit), rpm_limit: sanitizeNumeric(values.rpm_limit), model_tpm_limit: modelTpmLimit, @@ -764,21 +778,14 @@ const TeamInfoView: React.FC = ({ Models
- {info.models.length === 0 || info.models.includes("all-proxy-models") ? ( - All proxy models - ) : ( - <> - {info.models.map((model: string, index: number) => ( - - {model} - - ))} - {(info.access_group_models || []).map((model: string, index: number) => ( - - {model} - - ))} - + {computeTeamModelBadges(info.models, info.access_group_models || [], info.access_group_details).map( + (badge, index) => ( + + + {badge.label} + + + ), )}
@@ -982,7 +989,7 @@ const TeamInfoView: React.FC = ({ { + it("substitutes the no-default-models sentinel for an empty selection", () => { + expect(normalizeTeamModelSelection([])).toEqual(["no-default-models"]); + expect(normalizeTeamModelSelection(undefined)).toEqual(["no-default-models"]); + }); + + it("passes a non-empty selection through untouched", () => { + expect(normalizeTeamModelSelection(["gpt-4o-mini"])).toEqual(["gpt-4o-mini"]); + expect(normalizeTeamModelSelection(["all-proxy-models"])).toEqual(["all-proxy-models"]); + }); +}); + +describe("computeTeamModelBadges", () => { + it("attributes group-only models to the groups granting them", () => { + const badges = computeTeamModelBadges(["sonnet-direct"], [], GRANTS); + expect(badges).toEqual([ + { + label: "sonnet-direct", + kind: "direct", + tooltip: "Granted directly in the team's model list", + }, + { label: "haiku", kind: "access-group", tooltip: "Granted via access groups shared, extra" }, + { label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" }, + { label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" }, + ]); + }); + + it("marks a model both direct and group-granted on the direct badge, without a duplicate badge", () => { + const badges = computeTeamModelBadges(["haiku"], [], GRANTS); + expect(badges).toEqual([ + { + label: "haiku", + kind: "direct", + tooltip: "Granted directly in the team's model list, and also via access groups shared, extra", + }, + { label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" }, + { label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" }, + ]); + }); + + it("shows the no-default-models sentinel as its own badge and keeps group badges visible", () => { + const badges = computeTeamModelBadges(["no-default-models"], [], [GRANTS[0]]); + expect(badges.map((b) => [b.label, b.kind])).toEqual([ + ["No default models", "no-default"], + ["haiku", "access-group"], + ["gpt-4o-mini", "access-group"], + ]); + }); + + it("still shows group badges when the empty model list grants everything", () => { + const badges = computeTeamModelBadges([], [], [GRANTS[0]]); + expect(badges[0]).toEqual({ + label: "All proxy models", + kind: "all-proxy", + tooltip: "The team's model list is empty, so it can access every model on the proxy", + }); + expect(badges.slice(1).map((b) => b.label)).toEqual(["haiku", "gpt-4o-mini"]); + }); + + it("distinguishes the all-proxy-models sentinel from an empty list in the tooltip", () => { + const badges = computeTeamModelBadges(["all-proxy-models"], [], []); + expect(badges).toEqual([ + { + label: "All proxy models", + kind: "all-proxy", + tooltip: "Granted by the All Proxy Models entry in the team's model list", + }, + ]); + }); + + it("falls back to the flat access_group_models list when per-group details are absent", () => { + const badges = computeTeamModelBadges(["direct-model"], ["haiku"], undefined); + expect(badges).toEqual([ + { label: "direct-model", kind: "direct", tooltip: "Granted directly in the team's model list" }, + { label: "haiku", kind: "access-group", tooltip: "Granted via an access group" }, + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts new file mode 100644 index 00000000000..91ddf0f4045 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts @@ -0,0 +1,82 @@ +export const ALL_PROXY_MODELS = "all-proxy-models"; +export const NO_DEFAULT_MODELS = "no-default-models"; + +export interface TeamAccessGroupModelGrant { + access_group_id: string; + access_group_name: string; + models: string[]; +} + +export type TeamModelBadgeKind = "all-proxy" | "no-default" | "direct" | "access-group"; + +export interface TeamModelBadge { + label: string; + kind: TeamModelBadgeKind; + tooltip: string; +} + +export function normalizeTeamModelSelection(models: string[] | undefined): string[] { + return models && models.length > 0 ? models : [NO_DEFAULT_MODELS]; +} + +const describeGroups = (names: string[]): string => + names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`; + +export function computeTeamModelBadges( + models: string[], + accessGroupModels: string[], + accessGroupDetails: TeamAccessGroupModelGrant[] | undefined, +): TeamModelBadge[] { + const grants = accessGroupDetails ?? []; + const groupNamesFor = (model: string): string[] => + grants.filter((g) => g.models.includes(model)).map((g) => g.access_group_name); + const viaGroups = (model: string): string => { + const names = groupNamesFor(model); + return names.length > 0 ? describeGroups(names) : "an access group"; + }; + + const allProxy = models.length === 0 || models.includes(ALL_PROXY_MODELS); + const directModels = allProxy ? [] : models.filter((m) => m !== NO_DEFAULT_MODELS); + const groupModels = [...new Set(grants.length > 0 ? grants.flatMap((g) => g.models) : accessGroupModels)].filter( + (m) => !directModels.includes(m), + ); + + const allProxyBadge: TeamModelBadge = { + label: "All proxy models", + kind: "all-proxy", + tooltip: models.includes(ALL_PROXY_MODELS) + ? "Granted by the All Proxy Models entry in the team's model list" + : "The team's model list is empty, so it can access every model on the proxy", + }; + const noDefaultBadge: TeamModelBadge = { + label: "No default models", + kind: "no-default", + tooltip: "No models are granted directly. Access comes only from access groups", + }; + const headBadge = (): TeamModelBadge[] => { + if (allProxy) return [allProxyBadge]; + if (models.includes(NO_DEFAULT_MODELS)) return [noDefaultBadge]; + return []; + }; + + return [ + ...headBadge(), + ...directModels.map( + (m): TeamModelBadge => ({ + label: m, + kind: "direct", + tooltip: + groupNamesFor(m).length > 0 + ? `Granted directly in the team's model list, and also via ${viaGroups(m)}` + : "Granted directly in the team's model list", + }), + ), + ...groupModels.map( + (m): TeamModelBadge => ({ + label: m, + kind: "access-group", + tooltip: `Granted via ${viaGroups(m)}`, + }), + ), + ]; +} From 96a8b7f488d48d338fa2ba1007fda6a45d380499 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 02:09:53 +0000 Subject: [PATCH 45/74] chore(ui): regenerate dashboard api types for tier_turns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f1660e77ad9..8e950874a10 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21391,6 +21391,13 @@ export interface components { * @description What the routed traffic actually cost */ spend: number; + /** + * Tier Turns + * @description Turns per tier, keyed by the tier name the routing decision recorded at request time (never re-derived at read time, since the tier-to-model mapping is mutable config). Tier names are scoped to this group's router_type and are not comparable across types: a complexity router reports 'simple'/'medium'/'complex'/'reasoning', a quality router reports its numeric quality tier, and an adaptive router records no tier at all. Turns no tier served (the classifier fell back to default_model) are absent rather than pooled under a sentinel key, so the values may sum to less than turns + */ + tier_turns?: { + [key: string]: number; + }; /** Turns */ turns: number; }; From a2bb97fa7c80a04be85f4b26f7ffb70a988c1a84 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 02:13:27 +0000 Subject: [PATCH 46/74] fix(websearch_interception): satisfy type discipline gate for search metadata forwarding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../websearch_interception/handler.py | 40 ++++++++++++------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 6be64c89828..3e24c64ad82 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1288,14 +1288,11 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None search_litellm_params: dict[str, Any] = {} - search_tool_name: str | None = None + search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) search_provider = search_litellm_params.get("search_provider") - selected_tool_name = search_tool.get("search_tool_name") - if isinstance(selected_tool_name, str) and selected_tool_name: - search_tool_name = selected_tool_name # Fallback to perplexity if no router or no search tools configured if not search_provider: @@ -1318,14 +1315,20 @@ class WebSearchInterceptionLogger(CustomLogger): ) ) search_kwargs: Final = { - **{ - key: value - for key, value in search_litellm_params.items() - if key != "search_provider" and value is not None - }, - **({} if search_metadata is None else {"litellm_metadata": search_metadata}), + key: value + for key, value in search_litellm_params.items() + if key != "search_provider" and value is not None } - result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + result: Final = ( + await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + if search_metadata is None + else await litellm.asearch( + query=query, + search_provider=search_provider, + litellm_metadata=search_metadata, + **search_kwargs, + ) + ) # Format using transformation function search_result_text: Final = WebSearchTransformation.format_search_response(result) @@ -1386,7 +1389,7 @@ class WebSearchInterceptionLogger(CustomLogger): def _build_search_request_metadata( user_api_key_auth: "UserAPIKeyAuth", search_tool_name: str | None, - ) -> dict[str, object]: + ) -> Mapping[str, object]: """ Spend-tracking metadata for the intercepted search, so its provider cost is logged and billed against the key/user/team that made the originating LLM request instead @@ -1394,16 +1397,23 @@ class WebSearchInterceptionLogger(CustomLogger): """ from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - user_api_key_metadata: StandardLoggingUserAPIKeyMetadata = ( + user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) ) - return { + return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches **user_api_key_metadata, - **({} if search_tool_name is None else {"model_group": search_tool_name}), + "model_group": search_tool_name, "user_api_key": user_api_key_auth.api_key, "user_api_key_auth": user_api_key_auth, } + @staticmethod + def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None: + if search_tool is None: + return None + search_tool_name: Final = search_tool.get("search_tool_name") + return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None + @staticmethod def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": if not kwargs: From 0791dd941b3d5261a037834c3c139034d4fdc83a Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 02:19:57 +0000 Subject: [PATCH 47/74] test(proxy): assert the copy _add_team_member_budget_table returns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_team_endpoints.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a1cbc77b7a5..1e47010b57c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -932,8 +932,11 @@ async def test_add_team_member_budget_table_success(): ) # Verify the result - assert result == team_info_response assert result.team_member_budget_table == mock_budget_record + assert result == team_info_response.model_copy( + update={"team_member_budget_table": mock_budget_record} + ) + assert team_info_response.team_member_budget_table is None # Verify database call was made correctly mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( From 0a606cb258f731ddbddf71655a1af6b6b48e6213 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 7 Aug 2026 19:22:35 -0700 Subject: [PATCH 48/74] fix(otel): name the RPC system and upstream on MCP tool-call spans (#35857) * fix(otel): name the RPC system and upstream on MCP tool-call spans An MCP tool-call span carried only gen_ai.*, mcp.* and litellm.* attributes. A CLIENT span holding none of the http/db/messaging/rpc families is unclassifiable, so Elastic APM indexed these spans as span.type=unknown with no span.subtype at all, and its span-links API then rejected the whole trace with "Missing required fields (span.subtype)". MCP frames every message as JSON-RPC 2.0, so the tool-call span now names rpc.system. It names server.address and server.port alongside it, derived from the already-redacted mcp_server_resource origin: naming the RPC system makes a consumer treat the span as a downstream dependency and key that dependency off the server address, so emitting one without the other labels the dependency ":0". The tools/list span is left alone. It reaches the callbacks with no upstream identity, and a listing can span several upstreams, so it has no address to attach and would produce exactly that ":0" node. The wire is untouched: streamable MCP still returns HTTP 200 with isError: true. * fix(otel): drop rpc.system when no MCP upstream address resolved server.address and server.port come from mcp_server_resource, which is absent whenever the tool name resolves to no registered server, is None for a stdio transport that has no host to log, and parses to no host for an IPv6 origin the redactor rebuilds without its brackets. rpc.system was stamped unconditionally, so each of those paths emitted it alone and named the dependency ":0", the outcome the address pair exists to prevent. Gating the system attribute on a resolved address makes the pairing structural rather than leaving it to the two extractors happening to agree. * fix(otel): require a full MCP destination before naming the RPC system The gate gave rpc.system a resolved address, but not a resolved port. A host-bearing scheme outside the HTTP(S) default-port map resolves an address alone, and mcp_servers[].url is not scheme-validated, so an origin like mcp://host or ws://host reaches the mapper and names the dependency host:0 instead of the :0 the previous commit removed. Gating on the complete pair closes it, and covers a port of 0 as well. _upstream_address_port also gets a direct contract test, including the IPv6 origin the redactor rebuilds without brackets. * fix(otel): do not raise when an MCP origin has an unparseable port _redact_mcp_resource_url rebuilds the origin without its IPv6 brackets, so a zone-scoped address leaves a truthy hostname behind that the host check admits: http://[fe80::1%25eth0]:80 becomes http://fe80::1%25eth0:80, whose hostname is fe80 and whose port raises ValueError. That propagated out of MCPToolCallSpanData.from_standard_logging_payload and cost the span. Reading both halves inside a guard degrades an unparseable origin to no address, which is already how the mapper treats an unresolvable upstream, and matches the guard the redactor puts around the same split. The scheme default port drops the dict literal so the LIT002 ceiling stays put. --- litellm/integrations/otel/__init__.py | 2 + litellm/integrations/otel/mappers/genai.py | 5 ++ litellm/integrations/otel/model/payloads.py | 33 ++++++++ litellm/integrations/otel/model/semconv.py | 14 ++++ .../integrations/otel/test_otel_v2_logger.py | 79 +++++++++++++++++++ .../otel/test_otel_v2_sources_of_truth.py | 38 ++++++++- 6 files changed, 170 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/otel/__init__.py b/litellm/integrations/otel/__init__.py index 94442e96adb..9c1205bb277 100644 --- a/litellm/integrations/otel/__init__.py +++ b/litellm/integrations/otel/__init__.py @@ -57,6 +57,7 @@ from litellm.integrations.otel.model.semconv import ( Metric, Network, NetworkTransport, + RpcSystem, Server, resolve_operation, resolve_provider, @@ -102,6 +103,7 @@ __all__ = [ "ProxyRequestSpanData", "RequestContext", "RequestIdentity", + "RpcSystem", "Server", "ServerInfo", "ServiceSpanData", diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index af56734bec1..032441535e0 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -31,7 +31,9 @@ from litellm.integrations.otel.model.semconv import ( MCP, Error, GenAI, + JsonRpc, LiteLLM, + RpcSystem, Server, ) from litellm.integrations.otel.model.spans import db_system @@ -94,11 +96,14 @@ class GenAIMapper: _MCP_ATTRS: dict[str, Callable[[MCPToolCallSpanData], AttrValue | None]] = { GenAI.OPERATION_NAME: lambda d: d.operation.value, + JsonRpc.SYSTEM: lambda d: RpcSystem.JSONRPC.value if d.server_address and d.server_port else None, MCP.METHOD_NAME: lambda d: d.method, MCP.SESSION_ID: lambda d: d.session_id, GenAI.TOOL_NAME: lambda d: d.tool_name or None, GenAI.TOOL_CALL_ARGUMENTS: lambda d: d.arguments_json, GenAI.TOOL_CALL_RESULT: lambda d: d.result_json, + Server.ADDRESS: lambda d: d.server_address, + Server.PORT: lambda d: d.server_port, LiteLLM.MCP_SERVER_NAME: lambda d: d.server_name, LiteLLM.CALL_ID: lambda d: d.identity.call_id or None, f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 02499010624..aba9cc80240 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -364,6 +364,34 @@ class LLMCallSpanData: # --- the MCP tool-call model ------------------------------------------------- # +def _upstream_address_port(resource: str | None) -> tuple[str | None, int | None]: + """Split a redacted MCP server origin into ``server.address`` / ``server.port``. + + ``mcp_server_resource`` is a scheme + host + port origin with userinfo, path, + query and fragment already stripped. The port falls back to the scheme default + when the origin omits it, because a consumer that keys a downstream dependency + off the address renders a missing port as ``0``. + + The origin is rebuilt without its IPv6 brackets upstream, so reading the port can + raise on an address the host check still admits: a zone-scoped ``fe80::1%25eth0`` + leaves a truthy hostname of ``fe80`` behind. Both halves are read inside the guard + so an unparseable origin yields no address rather than propagating out of span + construction, matching how the redactor guards the same split. + """ + if not resource: + return None, None + try: + parsed: Final = urlsplit(resource) + hostname: Final = parsed.hostname + port: Final = parsed.port + except ValueError: + return None, None + if not hostname: + return None, None + default_port: Final = 443 if parsed.scheme == "https" else 80 if parsed.scheme == "http" else None + return hostname, port or default_port + + @dataclass(frozen=True) class MCPToolCallSpanData: """One MCP ``tools/call`` execution, parsed from a closed request's payload. @@ -378,6 +406,8 @@ class MCPToolCallSpanData: method: str tool_name: str server_name: str | None + server_address: str | None + server_port: int | None session_id: str | None arguments_json: str | None result_json: str | None @@ -390,11 +420,14 @@ class MCPToolCallSpanData: cls, payload: StandardLoggingPayload, capture_content: bool = False ) -> MCPToolCallSpanData: meta: Final = _mcp_tool_call_metadata(cast(Mapping[str, object], payload)) + address, port = _upstream_address_port(as_str(meta.get("mcp_server_resource")) or None) return cls( operation=resolve_operation(as_str(payload.get("call_type"))), method=MCPMethod.TOOLS_CALL.value, tool_name=as_str(meta.get("name")) or "", server_name=as_str(meta.get("mcp_server_name")), + server_address=address, + server_port=port, session_id=as_str(meta.get("mcp_session_id")), arguments_json=( _json_or_none(meta.get("arguments")) if capture_content and meta.get("arguments") is not None else None diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 24f0b947b08..3d585c36b67 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -130,11 +130,25 @@ class JsonRpc: """JSON-RPC keys carried on MCP spans. The error/status code lives in the ``rpc.*`` namespace per semconv, not ``jsonrpc.*``.""" + SYSTEM: Final = "rpc.system" REQUEST_ID: Final = "jsonrpc.request.id" PROTOCOL_VERSION: Final = "jsonrpc.protocol.version" RESPONSE_STATUS_CODE: Final = "rpc.response.status_code" +class RpcSystem(str, Enum): + """Well-known values for ``rpc.system``. MCP frames every message as JSON-RPC 2.0. + + Naming the system also classifies the span: a CLIENT span carrying none of the + ``rpc.*``/``http.*``/``db.*``/``messaging.*`` families records no span type or + subtype in backends that derive those from the attribute family. It is emitted + only alongside ``server.address``/``server.port``, since a backend that reads it + as a downstream dependency names that dependency from the server address. + """ + + JSONRPC = "jsonrpc" + + class NetworkTransport(str, Enum): """Well-known values for ``network.transport``.""" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 41c02501acc..2573ad5a375 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -356,6 +356,7 @@ def _mcp_payload(**overrides): "arguments": {"city": "Paris"}, "result": {"temp_c": 21}, "mcp_server_name": "weather-mcp", + "mcp_server_resource": "https://weather.example.com", "mcp_session_id": "sess-abc123", }, }, @@ -452,6 +453,8 @@ def test_mcp_tool_call_failure_marks_error(): assert span.name == "tools/call get_weather" assert span.status.status_code is StatusCode.ERROR assert span.attributes["error.type"] == "MCPError" + assert span.attributes["rpc.system"] == "jsonrpc" + assert span.attributes["server.address"] == "weather.example.com" def test_mcp_tool_call_deduped_on_repeat(): @@ -531,6 +534,82 @@ _MCP_SPAN_CASES = [ ] +def test_mcp_tool_call_names_its_rpc_system_and_upstream(): + """A tool-call span names the RPC system, and always alongside the upstream it called. + + A CLIENT span holding none of the ``rpc.*``/``http.*``/``db.*``/``messaging.*`` + families is unclassifiable, so a backend deriving a span type from them has nothing + to derive from: Elastic APM indexed these spans as ``span.type=unknown`` with no + ``span.subtype`` at all, and its span-links API then rejected the whole trace with + ``Missing required fields (span.subtype)``. + + The two assertions are one invariant, not two. Naming the RPC system makes a + consumer treat the span as a downstream dependency and key that dependency off + ``server.address``/``server.port``; emitting the first without the second names the + dependency ``:0``, which is worse than leaving the span unclassified. + """ + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_payload()}, None, None, None + ) + ) + (span,) = exporter.get_finished_spans() + assert span.attributes["rpc.system"] == "jsonrpc" + assert span.attributes["server.address"] == "weather.example.com" + assert span.attributes["server.port"] == 443 + + +@pytest.mark.parametrize( + "resource", + [None, "mcp://weather.example.com", "ws://weather.example.com"], + ids=["no-resource", "scheme-with-no-default-port", "ws-scheme"], +) +def test_mcp_tool_call_omits_rpc_system_without_a_complete_upstream(resource): + """A tool call drops the RPC system unless the full destination resolved. + + ``mcp_server_resource`` is absent whenever the tool name resolves to no registered + server, and it is ``None`` for a transport with no host to log at all (stdio). A + host-bearing scheme outside the HTTP(S) default-port map resolves an address but no + port, and the ``url`` field is not scheme-validated, so that state is reachable from + config. Each case names the dependency ``:0`` or ``host:0`` if ``rpc.system`` ships + on its own, so the pairing is enforced here rather than left to the extractors + happening to agree. + """ + logger, exporter = _logger() + payload = _mcp_payload() + if resource is None: + del payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_resource"] + else: + payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_resource"] = resource + asyncio.run( + logger.async_log_success_event({"standard_logging_object": payload}, None, None, None) + ) + (span,) = exporter.get_finished_spans() + assert "rpc.system" not in span.attributes + assert "server.port" not in span.attributes + assert span.attributes["mcp.method.name"] == "tools/call" + + +def test_mcp_list_tools_omits_rpc_system_without_an_upstream(): + """The discovery span carries no upstream identity, so it must not claim to be RPC. + + ``tools/list`` reaches the callbacks with no ``mcp_tool_call_metadata``, so there is + no ``server.address`` to attach and a listing can span several upstreams anyway. + Naming ``rpc.system`` here would buy a ``span.subtype`` at the cost of a bogus ``:0`` + dependency node in every consumer that aggregates on it. + """ + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_list_payload()}, None, None, None + ) + ) + (span,) = exporter.get_finished_spans() + assert "rpc.system" not in span.attributes + assert "server.address" not in span.attributes + + @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) def test_mcp_span_nests_under_transport_without_propagated_context( make_payload, span_name diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 612ac1e5113..19d0cfc0b18 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -24,7 +24,11 @@ from litellm.integrations.otel import ( resolve_provider, ) from litellm.integrations.otel.model import spans as spans_mod -from litellm.integrations.otel.model.payloads import LLMCallSpanData, RequestIdentity +from litellm.integrations.otel.model.payloads import ( + LLMCallSpanData, + RequestIdentity, + _upstream_address_port, +) from litellm.integrations.otel.model.spans import ( SPAN_REGISTRY, LiteLLMSpanKind, @@ -177,6 +181,7 @@ def test_mcp_attribute_vocabulary_is_complete(): "mcp.resource.uri", "jsonrpc.request.id", "jsonrpc.protocol.version", + "rpc.system", "rpc.response.status_code", "gen_ai.operation.name", "gen_ai.tool.name", @@ -808,3 +813,34 @@ def test_promoted_baggage_is_bounded_allowlist(): # http.* is never a promoted key assert HTTP.ROUTE not in promoted assert HTTP.REQUEST_METHOD not in promoted + + +@pytest.mark.parametrize( + "resource, expected", + [ + ("https://weather.example.com", ("weather.example.com", 443)), + ("http://weather.example.com", ("weather.example.com", 80)), + ("https://weather.example.com:8443", ("weather.example.com", 8443)), + ("mcp://weather.example.com", ("weather.example.com", None)), + ("http://::1:8080", (None, None)), + ("http://fe80::1%25eth0:80", (None, None)), + (None, (None, None)), + ("", (None, None)), + ], + ids=[ + "https-default", + "http-default", + "explicit-port", + "no-default-port", + "ipv6-unbracketed", + "ipv6-zone-scoped", + "none", + "empty", + ], +) +def test_upstream_address_port(resource, expected): + """The redacted MCP origin resolves to the address and port a consumer names its + dependency from. A scheme outside the default-port map yields no port, and an IPv6 + origin yields nothing at all because the redactor rebuilds it without its brackets; + both are why the mapper gates ``rpc.system`` on the complete pair.""" + assert _upstream_address_port(resource) == expected From 09a98f55052f03d78aeea4a6fe6e3916078b8220 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 7 Aug 2026 19:36:30 -0700 Subject: [PATCH 49/74] test(e2e): settle control-plane writes across every replica, not just one The suite already waits for a new model or agent to become servable before handing it back, but that wait returns on the first successful read. Every request opens a fresh connection (e2e_http calls requests.* with no Session), so a load-balanced Service routes each one independently: one successful read proves one replica converged, and the caller's next request re-rolls and can land on a replica that has not reloaded yet. At replicaCount: 2 this surfaced as 30 failures on a SHA that is green at 1 replica -- 400 "Invalid model name passed", 404 "Guardrail not found", "no healthy deployments for this model", and a /model/info listing that contained one of two models created moments apart. Add PROPAGATION_TIMEOUT (default 15s, override E2E_PROPAGATION_TIMEOUT) and settle_propagation(), sized off the proxy's proxy_config_reload_interval_seconds (30s by default, 7s on the e2e stack) plus margin, and settle after every control-plane create whose object the suite then uses: - ProxyClient.create_model and A2AClient.register_agent, after their existing polls -- the poll still fails loudly if the object never appears at all - GuardrailsClient.register, which had no barrier; create_content_filter_guardrail and create_bedrock_guardrail now route through it instead of POSTing directly - the guardrail creates in mcp_client and logging_client - the vertex passthrough model, whose body cannot go through create_model Left alone: the /model/new calls that assert a 403 or read back a status code, since they never use the model. --- tests/e2e/a2a/a2a_client.py | 7 +- tests/e2e/e2e_config.py | 26 +++++++ tests/e2e/guardrails/guardrails_client.py | 76 ++++++++----------- .../test_vertex_passthrough_e2e.py | 14 +++- tests/e2e/logging/logging_client.py | 3 +- tests/e2e/mcp/mcp_client.py | 5 +- tests/e2e/proxy_client.py | 22 ++++-- 7 files changed, 98 insertions(+), 55 deletions(-) diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index e83897025a3..605dd8fb7e5 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field +from e2e_config import settle_propagation from e2e_http import NoBody, Result, Success, get_external, is_ok from proxy_client import ProxyClient @@ -298,7 +299,9 @@ class A2AClient: the next DB reload. A card read or message/send issued the instant this returns can therefore 404 on the agent it just created. Waiting here keeps every caller from having to poll, the same way ProxyClient.create_model - waits for a new model to become servable. + waits for a new model to become servable -- including the settle that + covers the other replicas, since one successful card read only proves the + replica that answered it has the agent. """ result = self.proxy.transport.post( "/v1/agents", @@ -307,7 +310,9 @@ class A2AClient: response_type=AgentResponse, ) if isinstance(result, Success): + written_at = time.monotonic() self._await_agent_servable(result.data.agent_id) + settle_propagation(written_at) return result def _await_agent_servable(self, agent_id: str) -> None: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 34770bc596b..277478eebaf 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -7,6 +7,7 @@ environment so the same tests run against localhost or a deployed proxy. from __future__ import annotations import os +import time import uuid from pathlib import Path @@ -75,6 +76,18 @@ POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) +# How long a control-plane write (/model/new, /guardrails, /v1/agents) may take to +# reach EVERY replica. Distinct from POLL_TIMEOUT, which is sized for spend-row +# flush; this one is sized for the proxy's config reload +# (`proxy_config_reload_interval_seconds`, 30s by default and 7s on the e2e stack) +# plus margin. +# +# The barriers below wait this out instead of returning on first sight, because a +# single successful read only proves ONE replica converged: every request opens a +# fresh connection, so a load-balanced Service routes each one independently and +# the next call re-rolls. See ProxyClient._await_model_servable. +PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15")) + EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") # Deliberately modest concurrency. The suite shares its proxy with every other @@ -137,3 +150,16 @@ def unique_marker() -> str: """A short unique token per call/run, so concurrent runs and the shared response cache never collide on prompts, tags, or customer ids.""" return uuid.uuid4().hex[:12] + + +def settle_propagation(written_at: float) -> None: + """Block until PROPAGATION_TIMEOUT has elapsed since `written_at`, a + `time.monotonic()` stamp taken the moment a control-plane write returned. + + Callers that already polled for the object still need this: the poll proves one + replica has it, not all of them. Waiting out the config-reload budget is what + makes the object safe to use on whichever replica the next request lands on. + """ + remaining = PROPAGATION_TIMEOUT - (time.monotonic() - written_at) + if remaining > 0: + time.sleep(remaining) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 93861d19922..85964529ada 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -11,7 +11,7 @@ from typing import Literal from pydantic import BaseModel -from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker from e2e_http import NoBody, Result, Success, unwrap from lifecycle import ResourceManager from models import ( @@ -104,25 +104,14 @@ class GuardrailsClient: proxy: ProxyClient def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str: - return unwrap( - self.proxy.transport.post( - "/guardrails", - headers=self.proxy.transport.master, - json=GuardrailCreateBody( - guardrail=GuardrailSpecBody( - guardrail_name=name, - litellm_params=ContentFilterParamsBody( - mode="pre_call", - default_on=True, - blocked_words=[ - BlockedWordBody(keyword=blocked_keyword, action="BLOCK") - ], - ), - ) - ), - response_type=GuardrailCreateResponse, - ) - ).guardrail_id + return self.register( + name, + ContentFilterParamsBody( + mode="pre_call", + default_on=True, + blocked_words=[BlockedWordBody(keyword=blocked_keyword, action="BLOCK")], + ), + ) def create_bedrock_guardrail( self, @@ -141,24 +130,15 @@ class GuardrailsClient: test takes out whatever else is running. Callers select the guardrail per-request instead, which keeps the blast radius to the test that wants it. """ - return unwrap( - self.proxy.transport.post( - "/guardrails", - headers=self.proxy.transport.master, - json=GuardrailCreateBody( - guardrail=GuardrailSpecBody( - guardrail_name=name, - litellm_params=BedrockGuardrailParamsBody( - mode="pre_call", - default_on=default_on, - guardrailIdentifier=identifier, - guardrailVersion=version, - ), - ) - ), - response_type=GuardrailCreateResponse, - ) - ).guardrail_id + return self.register( + name, + BedrockGuardrailParamsBody( + mode="pre_call", + default_on=default_on, + guardrailIdentifier=identifier, + guardrailVersion=version, + ), + ) def create_backend_model(self, resources: ResourceManager, prefix: str = "e2e-guard-backend") -> str: """Register a gemini chat deployment for a guardrail test to run against @@ -174,11 +154,19 @@ class GuardrailsClient: return model_name def register(self, name: str, params: GuardrailParamsBody) -> str: - """Register any guardrail via POST /guardrails and return its id. New - built-ins register with default_on=False and are opted into per request - via the chat body's `guardrails` list, so one guardrail under test never - intercepts unrelated traffic on the shared proxy.""" - return unwrap( + """Register any guardrail via POST /guardrails and return its id, once every + replica can be expected to serve it. New built-ins register with + default_on=False and are opted into per request via the chat body's + `guardrails` list, so one guardrail under test never intercepts unrelated + traffic on the shared proxy. + + /guardrails is a control-plane route and guardrails reach the data plane on + the config reload, so a request naming this guardrail the instant the POST + returns can 404 with "Guardrail not found" on a replica that has not + reloaded. There is no data-plane read that lists guardrails, so unlike + ProxyClient.create_model this settles on the propagation budget alone with + nothing to poll first.""" + guardrail_id = unwrap( self.proxy.transport.post( "/guardrails", headers=self.proxy.transport.master, @@ -188,6 +176,8 @@ class GuardrailsClient: response_type=GuardrailCreateResponse, ) ).guardrail_id + settle_propagation(time.monotonic()) + return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: _ = self.proxy.transport.delete( diff --git a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py index b6d90b7f6a2..5e9c9f614e5 100644 --- a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py @@ -23,11 +23,12 @@ model, spend > 0), correlated by the x-litellm-call-id header. """ import os +import time import pytest from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import settle_propagation, unique_marker from e2e_http import NoBody, require_successful_call, unwrap from lifecycle import ResourceManager from models import SpendLogRow @@ -90,7 +91,14 @@ class _ModelDeleteBody(BaseModel): def _add_vertex_passthrough_model( client: PassthroughClient, model_name: str, project: str, credentials: str ) -> str: - return unwrap( + """Register the passthrough deployment and settle before the caller uses it. + + This body carries `use_in_pass_through` and a pinned `model_info.id`, so it + cannot go through ProxyClient.create_model -- but it needs that helper's + propagation settle just the same, or the passthrough call can land on a replica + that has not reloaded yet. + """ + model_id = unwrap( client.proxy.transport.post( "/model/new", headers=client.proxy.transport.master, @@ -108,6 +116,8 @@ def _add_vertex_passthrough_model( response_type=_ModelNewResponse, ) ).model_id + settle_propagation(time.monotonic()) + return model_id def _delete_model(client: PassthroughClient, model_id: str) -> None: diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 7053bd1dfd2..d76f7b356b2 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -23,7 +23,7 @@ from typing import Callable, Literal import pytest from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError -from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation from proxy_client import ProxyClient from e2e_http import ( URL, @@ -409,6 +409,7 @@ class LoggingClient: ) guardrail_id = response.guardrail_id assert guardrail_id, f"create guardrail returned no id: {response!r}" + settle_propagation(time.monotonic()) return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 33ec557c339..73453478e5a 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel +from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap from models import KeyGenerateBody, ObjectPermission from proxy_client import ProxyClient @@ -330,7 +331,7 @@ class McpClient: tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is unique per test, so default_on only ever intercepts this test's own banned tool call on the shared proxy.""" - return unwrap( + guardrail_id = unwrap( self.proxy.transport.post( "/guardrails", headers=self.proxy.transport.master, @@ -345,6 +346,8 @@ class McpClient: response_type=GuardrailCreateResponse, ) ).guardrail_id + settle_propagation(time.monotonic()) + return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: _ = self.proxy.transport.delete( diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..2627bdb8038 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -70,6 +70,7 @@ from e2e_config import ( POLL_TIMEOUT, PROXY_BASE_URL, REQUEST_TIMEOUT, + settle_propagation, ) from transport import HttpTransport, SplitTransport, Transport @@ -161,13 +162,18 @@ class ProxyClient: """Register a deployment under `model_name` and return its proxy-assigned model_id, once the model is actually servable on the data plane. - /model/new is a control-plane route; in a split control/data-plane - deployment the gateway (data plane, which serves /chat, /ocr, ...) only - picks the new model up on its next DB reload, so a call issued the instant - this returns can race the reload and 400 with "Invalid model name passed". - We therefore poll the data-plane /v1/models until the model appears before - handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request.""" + /model/new is a control-plane route; the data plane (which serves /chat, + /ocr, ...) only picks the new model up on its next DB reload, so a call + issued the instant this returns can race the reload and 400 with "Invalid + model name passed". We poll the data-plane /v1/models until the model + appears, then settle for the remainder of the propagation budget. + + Both steps are needed, and the second is the one that matters at >1 replica. + The poll proves *a* replica is serving the model; it cannot prove they all + are, because every request opens a fresh connection and a load-balanced + Service routes each one independently -- so the caller's next request + re-rolls and can land on a replica that has not reloaded yet. Waiting out + PROPAGATION_TIMEOUT is what makes the model safe to use anywhere.""" model_id = unwrap( self.transport.post( "/model/new", @@ -180,7 +186,9 @@ class ProxyClient: response_type=ModelNewResponse, ) ).model_id + written_at = time.monotonic() self._await_model_servable(model_name) + settle_propagation(written_at) return model_id def _await_model_servable(self, model_name: str) -> None: From d4dc2c39e7ab3f560cd67d0e974629fc2fab79a6 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 7 Aug 2026 19:44:24 -0700 Subject: [PATCH 50/74] fix(guardrails): chunk oversized Bedrock ApplyGuardrail requests instead of failing (#36119) * feat(guardrails): chunk oversized Bedrock ApplyGuardrail requests instead of failing AWS's ApplyGuardrail API rejects requests whose content exceeds the account's per-request "maximum input size in text units" quota with a 400 ValidationException. That cap is account/region/policy-dependent and cannot be predicted from config, so it can only be reacted to. _make_apply_guardrail_request now tries the whole-content call first (no behavior change for requests that already fit). On a too-large ValidationException it bisects the flat content list and retries each half sequentially, recursing until every piece fits or cannot be split further, then merges the per-chunk responses (action, assessments, outputs, usage) into one so callers cannot tell chunking happened. A real guardrail block on any (sub-)chunk still raises immediately. Contextual-grounding requests are never chunked: grounding scores the response holistically against the whole reference source, so fragmenting it would produce misleading scores. Each chunk call also gets a small exponential backoff retry on AWS ThrottlingException (429), since chunking increases the number of per-second API calls and can trade a 400 for a 429. All new state is local to a single request's call stack (no shared cache, no cross-process coordination), so this is safe for single-pod, multi-pod, and cache-less LiteLLM proxy deployments alike. * fix(guardrails): address Bedrock ApplyGuardrail chunking review feedback Fixes three issues flagged in review of the chunking fallback: a single oversized content item couldn't be split (only list-length bisection was supported), a chunked request that got recovered still logged a stray failure telemetry entry alongside the real outcome, and flattening chunk outputs without positional bookkeeping could misalign masked text onto the wrong original message once a chunk had nothing to mask. * test(guardrails): add regression test for multi-level Bedrock guardrail chunking Confirms the too-large bisection recursion isn't capped at a single split: a payload that is still oversized after the first halving keeps splitting until every piece fits, converging on however many chunks it takes rather than only ever producing two. * fix(guardrails): hybrid bin-pack+bisection chunking, whitespace-safe splits Rework Bedrock ApplyGuardrail chunking from pure reactive bisection to a hybrid strategy: bin-pack content into fixed-budget batches up front as the fast path, falling back to the existing recursive bisection only for a batch AWS still rejects as too large. Avoids paying O(log n) round trips on every oversized request when a single pass would do. Also switch single-item text splitting from a raw character midpoint to the nearest whitespace boundary, so a fragment never starts or ends mid-word. Closes the accidental-severing case from review; the residual gap (a multi-word denied phrase deliberately straddling the boundary) is documented as an accepted limitation, since fixing it would require an overlap window reconciled against masked output with no documented length-preservation guarantee from AWS. * chore(ui): regenerate dashboard API types * fix(guardrails): don't retry an oversized Bedrock guardrail call as a throttle AWS reports an ApplyGuardrail request that exceeds the per-request text-unit cap as a ThrottlingException (429), not only as the documented ValidationException (400). Verified against a live guardrail with an active content-filter policy: a 3273-text-unit request comes back as "Input text size (3273 text units) exceeds the maximum allowed (1000 text units) for the content filter policy (Classic tier)". The throttle retry keyed off status 429 alone, so every oversized chunk burned the full backoff-retry budget - each attempt a billed AWS call preceded by a sleep - before the bisection fallback got a chance, at every level of the recursion. A size error is not transient; re-posting the same content can never succeed. It now short-circuits straight to bisection. Also rename _is_input_too_large_validation_error to _is_input_too_large_error (it never keyed off the status code, and the error is not always a ValidationException), correct the docstrings that asserted a 400, and log at warning level when a split happens so the recovery is visible without --detailed_debug. * Revert "chore(ui): regenerate dashboard API types" This reverts commit ebf8ba2fd57f13bccf7aa6c5dfcac41c74db1ed9. * fix(guardrails): group all fragments of one item and stop double-logging Two defects found in review, both invisible to the existing tests. Fragment grouping assumed a split content item always produces exactly two adjacent fragments. That holds for one bisection level but not two: an item split twice yields four fragments, which were regrouped in fixed pairs into two output entries for a single message. Since masking walks the merged outputs by a running index across the original, unchunked message list, that message was written back truncated to its first half and every later message shifted. Fragments now carry the size of the group they belong to, so any number of them collapse back into exactly one output entry. Telemetry was also double-counted. AsyncHTTPHandler.post calls raise_for_status(), so every non-200 from Bedrock reaches _sign_and_post's error path, which logged guardrail_failed_to_respond before re-raising as an HTTPException that the consolidating caller then logged again. A request recovered by chunking reported one failure per rejected attempt plus a success. The ApplyGuardrail path now opts out of that per-attempt logging, since it owns consolidated per-request logging; the connection-level branch still logs, as nothing else records it. The existing tests missed both because their mocks return a non-200 response object, while the real client raises. Added a helper that raises a genuine httpx.HTTPStatusError so these paths are covered the way production hits them, plus a case asserting an unrecoverable failure still logs exactly once rather than zero times. * refactor(guardrails): move Bedrock chunking rationale into docstrings The chunking work explained itself with inline comment blocks, which this repo's conventions do not want. Folded that reasoning into the docstrings of the functions it describes and dropped the comments, including the module-level constant blocks and the test-file banner. No behavior change. The banner also claimed AWS rejects an oversized request with a 400 ValidationException, which live testing disproved, so removing it drops a stale claim as well as an internal ticket reference from a public repo. * feat(guardrails): match AWS default chunk budget and make it configurable ApplyGuardrail's default quota is 25 text units, roughly 25,000 characters, per second. Chunking has to respect that throughput limit rather than just the per-request size, otherwise splitting an oversized request trades a size error for a throttle. The budget now defaults to 25,000 to match that default for every user, up from an arbitrary 20,000. Accounts with raised quotas can spend fewer calls by setting chunk_budget_chars on the guardrail. A value AWS still rejects as too large is bisected automatically, so an over-large setting costs an extra round trip rather than failing the request. * fix(guardrails): never split a Bedrock text into an empty fragment _nearest_whitespace_split_index could return len(text) when the only space at or after the midpoint was the final character, so the first fragment came back identical to the text AWS had just rejected as too large and the second came back empty. AWS rejects the unchanged fragment again, and each retry re-splits it into the same fragment, so an oversized single item shaped like a long unbroken token with one trailing space exhausted the stack with a RecursionError instead of scanning or surfacing Bedrock's error. Candidate boundaries that would leave either side empty are now discarded, and the raw midpoint is used when none remain. The midpoint is always safe because _split_bedrock_content only calls this for text of at least two characters. * style(guardrails): move chunking rationale out of comments and into docstrings * fix(guardrails): raise 500 when Bedrock reports a failure inside a 200 body Also types the credentials parameter on the new chunking helpers and rebuilds fragment grouping without mutating a list or rebinding an index * fix(guardrails): raise 500 when Bedrock reports a failure inside a 200 body Restores the source changes intended for a08e4cf309, which landed with only the test. Also types the credentials parameter on the new chunking helpers and rebuilds fragment grouping without mutating a list or rebinding an index * style(guardrails): sort the constants import into the first-party block * refactor(guardrails): bring the Bedrock chunking path under the LIT lint budgets Annotates never-rebound locals with Final, replaces the retry counter and the two branch-assigned locals with single bindings, and moves the internal chunking chain to Sequence parameters and tuple returns. Collections that reach the logged payload stay lists on purpose: redact_nested_match_and_regex_keys only traverses dict and list, so a tuple would carry PII past redaction. The remaining constructions are contract-bound and carry inline reasons * fix(guardrails): keep the pre-chunking contract for failures reported inside a 200 Reverts the 500 this branch introduced for an AWS 200 whose body carries an Output.__type exception marker: the request proceeds as it did before chunking existed. The logged status is now derived from the merged response instead of being hardcoded to success, so that shape is still reported as guardrail_failed_to_respond. The consolidated failure logger also goes back to logging a dict rather than a bare string, matching both the pre-chunking code and the InvokeGuardrailChecks path in this file * docs(guardrails): correct the docstring for failures reported inside a 200 body The raise was reverted, so the docstring no longer describes the code. Records that the request proceeds by design and points at LIT-5338 for closing the fail-open path behind the existing unreachable_fallback setting --------- Co-authored-by: spencer-burridge <265588760+spencer-burridge@users.noreply.github.com> --- litellm/constants.py | 1 + .../guardrail_hooks/bedrock_guardrails.py | 857 ++++++- .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 10 + .../guardrail_hooks/bedrock_guardrails.py | 3 + .../test_bedrock_guardrails.py | 1989 ++++++++++++----- .../proxy/guardrails/test_init_guardrails.py | 35 + 7 files changed, 2326 insertions(+), 570 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6f0e9e7afe2..30d3bb1f26e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -280,6 +280,7 @@ TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECO GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int( os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) ) +BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000 # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index e9e729fb118..eecbce57468 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -9,11 +9,14 @@ import os import sys sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path +import asyncio import copy import json +import re import sys -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import datetime, timezone +from itertools import accumulate, groupby from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast import httpx @@ -23,6 +26,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys @@ -46,6 +50,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailOutput, BedrockGuardrailQualifier, BedrockGuardrailResponse, + BedrockGuardrailUsage, BedrockRequest, BedrockTextContent, ) @@ -53,6 +58,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from botocore.awsrequest import AWSPreparedRequest + from botocore.credentials import Credentials from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -71,6 +77,17 @@ from litellm.types.utils import ( GUARDRAIL_NAME: Final = "bedrock" _BEDROCK_DYNAMIC_BODY_DENYLIST: Final = frozenset({"content", "source"}) +_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = ( + "text unit", + "maximum input size", + "content size", + "too long", + "too large", + "exceeds the maximum", +) +_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3 +_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5 +_BEDROCK_WHITESPACE: Final = re.compile(r"\s") # Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required). _BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke" # InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with @@ -118,6 +135,29 @@ class GuardrailMessageFilterResult(NamedTuple): target_indices: list[int] | None +class BedrockContentChunkResult(NamedTuple): + """One chunk's ApplyGuardrail response, paired with enough bookkeeping to + reconstruct global masked-output positions once every chunk is back. + + `content` is the exact content items this chunk was called with -- needed + so an all-clear chunk (empty `outputs`) can still contribute one unmasked + placeholder per item it covers, keeping every later chunk's masked text + aligned to its original global position. `fragment_group_size` is 1 for an + ordinary chunk, and otherwise the total number of consecutive chunk results + that together make up ONE original content item's own text (split because a + list of length 1 could not be bisected by list length). All of them must be + concatenated back into that one item's masked output rather than treated as + separate items. It is a count rather than a boolean because one item can be + bisected more than once: two levels of splitting produce four fragments for + a single item, not two, and grouping them in fixed pairs would emit two + outputs for one message and shift every later message's masked text. + """ + + response: BedrockGuardrailResponse + content: tuple[BedrockContentItem, ...] + fragment_group_size: int + + class ApplyGuardrailMessageSelection(NamedTuple): """Messages selected for an apply_guardrail scan + write-back metadata.""" @@ -168,12 +208,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): content_filter_threshold: float | None = 0.5, prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, + chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" + self.chunk_budget_chars = chunk_budget_chars self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` @@ -759,12 +801,35 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None = None, logging_event_type: GuardrailEventHooks | None = None, ) -> BedrockGuardrailResponse: + """Scan `messages`/`response` with ApplyGuardrail, chunking if it is too large. + + Content is bin-packed into budget-sized batches and each batch posted + sequentially, every batch independently falling back to bisection if AWS + rejects it. The per-batch responses are merged so callers cannot tell whether + chunking happened. + + Content using contextual grounding opts out of chunking entirely: grounding is + scored holistically against the whole reference source, so bisecting it would + fragment that evaluation and yield misleading scores. Such a request keeps the + old behavior of surfacing a too-large error rather than being split. + + `logging_event_type` drives what UI and spend logs report. It is distinct from + Bedrock's `source`, which is INPUT vs OUTPUT for the API body and must not be + confused with the proxy hook (pre_call / during_call / post_call); when omitted, + the legacy source-derived mapping is kept for backward compatibility. + + A guardrail *block* is logged where it happens, in + `_post_apply_guardrail_content`, because chunking stops immediately and there is + no later merged response to log instead. Everything else that fails out of the + chunking flow (an unrecoverable too-large error, a non-size validation error, + exhausted throttle retries) is a genuine end-to-end failure of this one logical + guardrail call and is logged exactly once here. + """ start_time: Final = datetime.now(timezone.utc) credentials, aws_region_name = self._load_credentials() bedrock_request_data: Final[dict] = dict( self.convert_to_bedrock_format(source=source, messages=messages, response=response) ) - bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse() api_key: str | None = None if request_data: dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data) @@ -778,6 +843,257 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if request_data.get("api_key") is not None: api_key = request_data["api_key"] + event_type: Final = ( + logging_event_type + if logging_event_type is not None + else (GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call) + ) + + content: Final[tuple[BedrockContentItem, ...]] = tuple(bedrock_request_data.get("content") or ()) + allow_chunking: Final = not self._content_uses_contextual_grounding(content) + + try: + responses: Final = await self._apply_guardrail_content_with_chunking( + content=content, + base_request_data=bedrock_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + except HTTPException as exc: + if not isinstance(exc.detail, dict): + self._log_apply_guardrail_failure( + detail=exc.detail, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + raise + merged_response: Final = self._merge_bedrock_guardrail_responses(responses) + self._log_apply_guardrail_success( + merged_response=merged_response, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + return merged_response + + async def _apply_guardrail_content_with_chunking( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + allow_chunking: bool, + ) -> tuple[BedrockContentChunkResult, ...]: + """Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large. + + Tries `content` as a single call first. AWS's per-request "maximum input + size in text units" quota is account/region/policy-dependent and cannot be + predicted ahead of time, so it is only ever discovered reactively: on an + error whose message indicates the input was too large (a ThrottlingException + in practice, a ValidationException per the docs -- see + ``_is_input_too_large_error``), the content is re-sent in smaller pieces. + + Probing with the whole payload first is what keeps a request AWS would have + accepted at exactly one call. Packing into fixed batches up front instead + would split conversations AWS was happy to take whole, multiplying billed + calls and guardrail latency on traffic that never had a size problem, and + no fixed budget can avoid that because the real cap is unknown here. + + Once a rejection proves the payload is over the cap, a multi-item payload is + re-sent as ``chunk_budget_chars``-sized batches rather than bisected: that + reaches a working size in one step instead of paying an O(log n) ladder of + rejected calls. Bisection remains the fallback for anything bin-packing + cannot make smaller, which is what makes the recursion terminate: a batch + already inside the budget packs back to itself, so it falls through to the + split below. A single oversized + content item (one very long message) is split by its own text instead of + by list length, since a list of length 1 has no items left to bisect -- + the resulting fragments all carry a ``fragment_group_size`` so the merge + step can recombine them into the one content item they came from, rather + than treating each fragment as its own item when reconstructing positions + for masking. That count covers however many fragments the item ended up + split into, not just two, since it can be bisected repeatedly: the + outermost single-item split stamps the total leaf count on every leaf + below it, overwriting any smaller count an inner split had set. A real + guardrail block on any (sub-)chunk raises immediately + -- callers must not lose that signal by continuing to post the remaining + chunks. + """ + try: + response: Final = await self._post_apply_guardrail_content_with_retry( + content=content, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + return ( + BedrockContentChunkResult( + response=response, + content=tuple(content), + fragment_group_size=1, + ), + ) + except HTTPException as exc: + if allow_chunking and self._is_input_too_large_error(exc.detail): + batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars) + if len(batches) > 1: + verbose_proxy_logger.warning( + "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; " + "re-sending as %d batches of at most %d characters", + len(content), + len(batches), + self.chunk_budget_chars, + ) + batch_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below + await self._apply_guardrail_content_with_chunking( + content=batch, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + for batch in batches + ] + return tuple(result for results in batch_results for result in results) + split_content: Final = self._split_bedrock_content(content) + if split_content is None: + raise + first_half, second_half = split_content + is_single_item_text_split: Final = len(content) == 1 + verbose_proxy_logger.warning( + "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; " + "splitting into %d + %d and retrying each", + len(content), + len(first_half), + len(second_half), + ) + first_results: Final = await self._apply_guardrail_content_with_chunking( + content=first_half, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + second_results: Final = await self._apply_guardrail_content_with_chunking( + content=second_half, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + combined_results: Final = tuple(first_results) + tuple(second_results) + if is_single_item_text_split: + return tuple( + result._replace(fragment_group_size=len(combined_results)) for result in combined_results + ) + return combined_results + raise + + async def _post_apply_guardrail_content_with_retry( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> BedrockGuardrailResponse: + """Post one ApplyGuardrail call for `content`, retrying with exponential + backoff on AWS ThrottlingException (HTTP 429). + + Chunking already trades one oversized call for several smaller ones, so + retries here are capped low -- they must not multiply per-request latency + by an order of magnitude when the account's per-second text-unit quota is + the binding constraint rather than the per-request size quota. + + A too-large rejection is deliberately excluded from the retry. AWS reports + it as a ThrottlingException (429), not only as a ValidationException, but + unlike a genuine throttle it is not transient: re-posting the same + oversized content can never succeed. Retrying it would burn every backoff + sleep and every (billed) attempt before the caller's bisection gets a + chance to split the content, at every level of the recursion. + """ + for attempt in range(_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + 1): + try: + return await self._post_apply_guardrail_content( + content=content, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + except HTTPException as exc: + if ( + exc.status_code != 429 + or self._is_input_too_large_error(exc.detail) + or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + ): + raise + await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt)) + raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") + + async def _post_apply_guardrail_content( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> BedrockGuardrailResponse: + """Make exactly one signed ApplyGuardrail HTTP call for `content` and + parse the result. Raises HTTPException on a guardrail block or any + non-200 response (including 429, handled by the retry wrapper above). + + AWS also reports some failures inside a 200 body, tagging ``Output.__type`` + with an Exception marker. Those deliberately do NOT raise: the request proceeds, + matching the behaviour of this code before chunking existed. The marker survives + the merge, so the one consolidated log entry still records + ``guardrail_failed_to_respond`` rather than a success. Making that path fail + closed is a separate change, tracked apart from this PR, and belongs behind the + existing ``unreachable_fallback`` setting rather than a hardcoded status. + + A block is logged here rather than by the caller: it ends the whole chunking + flow immediately, with no further chunks attempted, so there is no later + merged response for the caller to log instead. + """ + bedrock_request_data: Final = { # mutable-ok: outbound JSON request body + **base_request_data, + "content": content, + } # mutable-ok: outbound JSON request body prepared_request: Final = self._prepare_request( credentials=credentials, data=bedrock_request_data, @@ -792,42 +1108,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - # UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API - # body, which must not be confused with the proxy hook (pre_call / during_call / - # post_call). When omitted, keep legacy mapping for backward compatibility. - if logging_event_type is not None: - event_type = logging_event_type - else: - event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call - httpx_response: Final = await self._sign_and_post( prepared_request=prepared_request, request_data=request_data, event_type=event_type, start_time=start_time, + log_transport_failure=False, ) - ######################################################### - # Add guardrail information to request trace - ######################################################### - _json_response: Final = httpx_response.json() - tracing_detail: Final = self._build_tracing_detail(_json_response) - - # Raw Bedrock JSON is passed here; match/regex redaction runs once inside - # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_provider=self.guardrail_provider, - guardrail_json_response=_json_response, - request_data=request_data or {}, - guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), - start_time=start_time.timestamp(), - end_time=datetime.now(timezone.utc).timestamp(), - duration=(datetime.now(timezone.utc) - start_time).total_seconds(), - event_type=event_type, - tracing_detail=tracing_detail or None, - ) - ######################################################### if httpx_response.status_code == 200: + _json_response: Final = httpx_response.json() # check if the response was flagged verbose_proxy_logger.debug( "Bedrock AI response : %s", @@ -835,19 +1125,462 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): + self._log_apply_guardrail_attempt( + httpx_response=httpx_response, + json_response=_json_response, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) raise self._get_http_exception_for_blocked_guardrail( bedrock_guardrail_response, request_data=request_data ) - else: - status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) - verbose_proxy_logger.error( - "Bedrock AI: error in response. Status code: %s, response: %s", - httpx_response.status_code, - httpx_response.text, - ) - raise HTTPException(status_code=status_code, detail=detail_message) + return bedrock_guardrail_response - return bedrock_guardrail_response + status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) + verbose_proxy_logger.error( + "Bedrock AI: error in response. Status code: %s, response: %s", + httpx_response.status_code, + httpx_response.text, + ) + raise HTTPException(status_code=status_code, detail=detail_message) + + def _log_apply_guardrail_attempt( + self, + httpx_response: httpx.Response, + json_response: dict, # mutable-ok: raw AWS JSON payload + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log a single ApplyGuardrail HTTP attempt as-is (its own status, + derived from its own response). Used only for the blocked-content + case, which ends the whole chunking flow immediately.""" + tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response)) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=json_response, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + tracing_detail=tracing_detail or None, + ) + + def _log_apply_guardrail_success( + self, + merged_response: BedrockGuardrailResponse, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log one logical ApplyGuardrail call -- possibly several chunk calls + under the hood -- using its final merged response, so a chunked + request produces exactly one telemetry entry, the same as an + unchunked one would. + + AWS can report a failure inside an HTTP 200 body by tagging + ``Output.__type`` with an exception marker. That marker survives the merge, + so the status is derived from the merged response rather than assumed to be + a success, which is what the pre-chunking code reported for that shape.""" + tracing_detail: Final = self._build_tracing_detail(merged_response) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status=( + "guardrail_failed_to_respond" + if "Exception" in str((merged_response.get("Output") or {}).get("__type", "")) + else "success" + ), + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + tracing_detail=tracing_detail or None, + ) + + def _log_apply_guardrail_failure( + self, + detail: object, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log one logical ApplyGuardrail call that failed end-to-end (an + unrecoverable too-large error, a non-size validation error, or + exhausted throttle retries) as a single failure, rather than logging + every failed attempt chunking made along the way.""" + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + ) + + @staticmethod + def _content_uses_contextual_grounding(content: Sequence[BedrockContentItem]) -> bool: + """True if any content item carries a contextual-grounding qualifier + (``grounding_source``, ``query``, or the ``guard_content`` the response + itself is tagged with once grounding is present).""" + for item in content: + if (item.get("text") or {}).get("qualifiers"): # mutable-ok: read-only empty fallback + return True + return False + + @staticmethod + def _bin_pack_bedrock_content( + content: Sequence[BedrockContentItem], + budget: int, + ) -> tuple[tuple[BedrockContentItem, ...], ...]: + """Pack whole content items, in order, into batches whose combined text + length stays within `budget`, in a single pass that carries the running + total rather than re-summing the open batch per item. + + This is the fast-path half of the hybrid chunking strategy: bin-packing + at a conservative fixed budget keeps the common case at O(n / budget) + ApplyGuardrail calls instead of the O(log n) round trips pure reactive + bisection pays on every oversized request. An item whose own text + already exceeds `budget` is not split here -- it becomes its own + (still oversized) batch and is sent as-is; if AWS rejects that batch as + too large, `_apply_guardrail_content_with_chunking`'s existing + recursive-bisection fallback takes over for that batch only. + + `budget` comes from the guardrail's ``chunk_budget_chars`` setting and + defaults to 25,000, matching ApplyGuardrail's default quota of 25 text + units (roughly 1,000 characters each) per second. Packing to that size and + posting sequentially is what keeps chunking from tripping the rate quota + and trading a size error for a throttle. Accounts with raised quotas can + configure a larger budget to spend fewer calls. + + The budget is not a correctness dependency either way. AWS's effective cap + varies by account, region, and policy, is not a fixed character count, and + cannot be read from config, so any batch it still rejects falls back to + bisection, which self-corrects however wrong the value was. An over-large + budget therefore costs one extra probe-and-bisect round trip rather than + failing the request. + """ + if not content: + return (tuple(content),) + + lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content) + + def assign(carried: tuple[int, int], length: int) -> tuple[int, int]: + batch_index, used = carried + if used + length <= budget: + return batch_index, used + length + return batch_index + 1, length + + batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:]) + return tuple( + tuple(item for _, item in group) + for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0]) + ) + + @staticmethod + def _split_bedrock_content( + content: Sequence[BedrockContentItem], + ) -> tuple[tuple[BedrockContentItem, ...], tuple[BedrockContentItem, ...]] | None: + """Bisect `content` into two roughly-equal, non-empty halves. + + When `content` already holds more than one item, it is split by list + length. When it holds exactly one item, that item's own text is split + instead (a list of length 1 has no items left to bisect, but one very + long message is still a single content item) -- at the whitespace + character nearest the midpoint rather than a raw character index, so + the cut never lands inside a word/token. This is a plain, lossless + cut with no overlap: concatenating the two fragments in order always + reproduces the original text exactly, so merging back at + ``_merge_logical_unit_outputs`` needs no reconciliation step. + + Known, accepted limitation: whitespace splitting only guards against + *accidentally* severing a single token (one denied word, one PII + pattern) across the cut. It does not, and cannot without an overlap + window, stop a *multi-word* denied phrase deliberately positioned to + straddle the boundary -- each fragment can scan clean on its own and + still reassemble into the flagged phrase. AWS's own guidance on this + API acknowledges the same gap for input chunking ("a critical piece of + text could span two (or more) chunks if not carefully divided") with + no documented resolution, and overlap-and-reconcile was evaluated and + rejected for this PR: AWS's masking output has no documented + length-preservation guarantee, so reconciling an overlap region against + masked text is not sound in general. Out of scope for this PR. + + Returns None when there is nothing left to split -- a single item + whose text is too short to halve into two non-empty pieces -- so the + caller can give up and propagate the original too-large error instead + of recursing forever. + """ + if len(content) > 1: + midpoint: Final = max(1, len(content) // 2) + return tuple(content[:midpoint]), tuple(content[midpoint:]) + + text_content: Final = content[0].get("text") or BedrockTextContent() + text: Final = text_content.get("text") or "" + if len(text) < 2: + return None + split_at: Final = BedrockGuardrail._nearest_whitespace_split_index(text) + qualifiers: Final = text_content.get("qualifiers") + + def fragment(piece: str) -> BedrockContentItem: + block: Final = ( + BedrockTextContent(text=piece, qualifiers=qualifiers) if qualifiers else BedrockTextContent(text=piece) + ) + return BedrockContentItem(text=block) + + return (fragment(text[:split_at]),), (fragment(text[split_at:]),) + + @staticmethod + def _nearest_whitespace_split_index(text: str) -> int: + """Return the index nearest `text`'s midpoint that falls on a whitespace + boundary, so splitting `text[:i]` / `text[i:]` there never severs a word. + + Any Unicode whitespace counts, not just an ASCII space. Matching only `" "` + would leave the boundary unguarded for exactly the payloads that get large + enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited, so a deny-listed word sitting at the midpoint of + one would be cut in half, scan clean on both fragments, and reassemble + intact. + + The returned index always leaves both sides non-empty, which is what makes + the caller's recursion terminate. A boundary that would put the split at 0 + or at ``len(text)`` is discarded: it would hand back a fragment identical to + the text just rejected as too large, AWS would reject that again, and each + retry would re-split it into the same unchanged fragment until the stack ran + out. The dangerous shape is a text whose only space at or after the midpoint + is its final character. + + Falls back to the raw midpoint when no usable whitespace boundary exists, either + because `text` has none at all (a single giant token) or because the only + candidates were degenerate. That is still a correct, lossless split, just no + longer guaranteed word-safe for those cases. `text` must be at least two + characters, which `_split_bedrock_content` guarantees, so the midpoint itself + is never degenerate. + """ + midpoint: Final = len(text) // 2 + before: Final = max((found.end() for found in _BEDROCK_WHITESPACE.finditer(text, 0, midpoint)), default=None) + after_match: Final = _BEDROCK_WHITESPACE.search(text, midpoint) + candidates: Final = sorted( + (split for split in (before, after_match.end() if after_match else None) if split is not None), + key=lambda split: abs(split - midpoint), + ) + return next((split for split in candidates if 0 < split < len(text)), midpoint) + + @staticmethod + def _is_input_too_large_error(detail: object) -> bool: + """True if `detail` is an AWS error message for input exceeding the + per-request text-unit quota. + + Matched on the message rather than the status code on purpose: AWS is not + consistent about which error it raises for this. Observed against a live + guardrail with an active content-filter policy, an oversized request comes + back as a *ThrottlingException* (429) reading ``Input text size (3273 text + units) exceeds the maximum allowed (1000 text units) for the content filter + policy (Classic tier)``, while the documented failure mode is a + ValidationException (400). Keying off the message covers both. + + A guardrail *block* is also raised as an HTTPException with status 400, + but its ``detail`` is always a dict (built by + ``_get_http_exception_for_blocked_guardrail``); a non-200 API error's + ``detail`` is always the plain string returned by + ``_parse_bedrock_guardrail_error_response``. Checking ``isinstance(detail, + str)`` is therefore sufficient to never mistake a real block for a + too-large error. + """ + if not isinstance(detail, str): + return False + lowered: Final = detail.lower() + return any(substring in lowered for substring in _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS) + + @staticmethod + def _merge_bedrock_guardrail_responses( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> BedrockGuardrailResponse: + """Merge the per-chunk ApplyGuardrail responses of a chunked request into + one, so a caller cannot tell whether chunking happened. + + Only ever called with responses that all passed (a block raises + immediately from ``_apply_guardrail_content_with_chunking`` and is never + added to this list). ``action`` is only set on the merged response when + at least one chunk's raw response included it, and left absent otherwise + -- mirroring a real single-call response and matching what + ``_build_tracing_detail`` treats as "Bedrock didn't report an action". + + Fields this merge has no opinion on (``actionReason``, ``guardrailCoverage``, + ``blockedResponse``, anything AWS adds later) are carried over from the chunk + responses rather than dropped, so the response and the logged telemetry keep + the shape a single unchunked call returned. The merged keys below win. + + Per AWS's documented ApplyGuardrail contract, a single call's ``outputs`` + is positionally parallel to the ``content`` items *of that call*: an + entry per item when anything in the call was masked, or an empty list + when nothing in the whole call was masked. Downstream masking + (``_apply_masking_to_messages``) walks the merged ``outputs`` by a single + running index across the *original, unchunked* message list, so a later + chunk's masked text must land at the same global position it would have + if chunking had never happened. Naively concatenating each chunk's + ``outputs`` breaks that whenever a chunk had nothing masked (its empty + list would otherwise silently swallow its items' slots, shifting every + later chunk's masked text left onto the wrong message). So every + item -- masked or not -- always contributes exactly one entry here, + falling back to that item's own original (unmasked) text when its + chunk returned no output for it; a wholly-untouched result is then + collapsed back to an empty ``outputs`` list to match a real single-call + no-op response. A chunk that returns a nonzero output count not equal + to its item count is passed through as-is instead of guessed at, since + AWS's docs don't cover partial masking within one multi-item call. + """ + logical_units: Final = BedrockGuardrail._group_fragment_units(chunk_results) + per_unit_outputs: Final = tuple(BedrockGuardrail._merge_logical_unit_outputs(unit) for unit in logical_units) + merged_outputs: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list + output for outputs, _ in per_unit_outputs for output in outputs + ] + any_masked: Final = any(masked for _, masked in per_unit_outputs) + + actions: Final = tuple( + chunk_result.response.get("action") + for chunk_result in chunk_results + if isinstance(chunk_result.response.get("action"), str) + ) + merged_action: Final = ( + "GUARDRAIL_INTERVENED" if "GUARDRAIL_INTERVENED" in actions else (actions[-1] if actions else None) + ) + merged_assessments: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list + assessment + for chunk_result in chunk_results + for assessment in (chunk_result.response.get("assessments") or []) # mutable-ok: logged payload + ] + any_usage_reported: Final = any(chunk_result.response.get("usage") for chunk_result in chunk_results) + + merged: Final[BedrockGuardrailResponse] = cast( # cast-ok: TypedDict assembled from a comprehension + BedrockGuardrailResponse, + { # mutable-ok: builds the TypedDict payload + key: value for chunk_result in chunk_results for key, value in chunk_result.response.items() + }, + ) + if merged_action is not None: + merged["action"] = merged_action + if merged_outputs and any_masked: + merged["outputs"] = merged_outputs + merged["output"] = merged_outputs + if merged_assessments: + merged["assessments"] = merged_assessments + if any_usage_reported: + merged["usage"] = BedrockGuardrail._sum_bedrock_guardrail_usage(chunk_results) + return merged + + @staticmethod + def _sum_bedrock_guardrail_usage( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> BedrockGuardrailUsage: + """Sum each chunk's ``usage`` counters field-by-field into one totals dict. + + Keys are taken from the responses rather than from a fixed list, so a counter + this code does not know about (AWS has added several) is still summed and + reported instead of being silently dropped to zero.""" + chunk_usages: Final = tuple( + chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback + for chunk_result in chunk_results + ) + return cast( # cast-ok: TypedDict assembled from a comprehension + BedrockGuardrailUsage, + { # mutable-ok: builds the TypedDict payload + key: sum(usage.get(key) or 0 for usage in chunk_usages) + for key in dict.fromkeys(key for usage in chunk_usages for key in usage) + }, + ) + + @staticmethod + def _group_fragment_units( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> tuple[tuple[BedrockContentChunkResult, ...], ...]: + """Group consecutive text-fragment chunk results back into the one content + item each group came from, leaving every ordinary chunk result as a unit of + one. + + The group size is read off the results themselves rather than assumed, + because a single content item can be bisected repeatedly: two levels of + splitting yield four fragments for one item, not two. Assuming a fixed pair + here would emit two outputs for one message and shift every later message's + masked text onto the wrong message.""" + + def advance(carried: tuple[int, bool], result: BedrockContentChunkResult) -> tuple[int, bool]: + remaining, _ = carried + if remaining == 0: + return max(1, result.fragment_group_size) - 1, True + return remaining - 1, False + + starts: Final = tuple( + index + for index, (_, starts_unit) in enumerate(tuple(accumulate(chunk_results, advance, initial=(0, False)))[1:]) + if starts_unit + ) + return tuple(tuple(chunk_results[start:end]) for start, end in zip(starts, starts[1:] + (len(chunk_results),))) + + @staticmethod + def _merge_logical_unit_outputs( + unit: tuple[BedrockContentChunkResult, ...], + ) -> tuple[tuple[BedrockGuardrailOutput, ...], bool]: + """Reduce one logical unit (a fragment group of any size, or a single chunk + result) to the ``BedrockGuardrailOutput`` entries it contributes to the + merged response, plus whether any masking actually happened in it. + + Per AWS's documented ApplyGuardrail contract, a single call's + ``outputs`` is positionally parallel to the ``content`` items *of that + call*: an entry per item when anything in the call was masked, or an + empty list when nothing in the whole call was masked. Downstream + masking (``_apply_masking_to_messages``) walks the merged ``outputs`` + by a single running index across the *original, unchunked* message + list, so a later chunk's masked text must land at the same global + position it would have if chunking had never happened. So every item + -- masked or not -- always contributes exactly one entry here, falling + back to that item's own original (unmasked) text when its chunk + returned no output for it. A chunk that returns a nonzero output count + not equal to its item count is passed through as-is instead of guessed + at, since AWS's docs don't cover partial masking within one multi-item + call. + + A unit holding more than one result is a fragment group: every result in it + is one fragment of a single content item's text, so the group collapses to + one entry built from each fragment's masked text (or that fragment's own + original text where it came back unmasked), concatenated in order. This + holds for any group size, not only two. + """ + if len(unit) > 1: + + def fragment_outputs(result: BedrockContentChunkResult) -> tuple[BedrockGuardrailOutput, ...]: + return tuple(result.response.get("outputs") or result.response.get("output") or ()) + + def fragment_text(result: BedrockContentChunkResult) -> str: + source: Final = (result.content[0].get("text") or {}).get( # mutable-ok: read-only fallback + "text" + ) or "" + outputs: Final = fragment_outputs(result) + masked: Final = outputs[0].get("text") if outputs else None + return masked if masked is not None else source + + merged_text: Final = "".join(fragment_text(result) for result in unit) + any_masked: Final = any(fragment_outputs(result) for result in unit) + return (BedrockGuardrailOutput(text=merged_text),), any_masked + + (chunk_result,) = unit + chunk_outputs: Final = chunk_result.response.get("outputs") or chunk_result.response.get("output") or () + if len(chunk_outputs) == len(chunk_result.content): + return tuple(chunk_outputs), bool(chunk_outputs) + if not chunk_outputs: + return tuple( + BedrockGuardrailOutput( + text=(item.get("text") or {}).get("text") or "" # mutable-ok: read-only fallback + ) + for item in chunk_result.content + ), False + return tuple(chunk_outputs), True async def _sign_and_post( self, @@ -855,6 +1588,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, event_type: GuardrailEventHooks, start_time: "datetime", + log_transport_failure: bool = True, ) -> httpx.Response: """POST a signed Bedrock request, logging+raising on network/HTTP errors. @@ -862,6 +1596,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): transport-error handling cannot drift. Returns the raw ``httpx.Response`` on success (including non-2xx that httpx did not raise on); the 200-path logging, status and tracing stay with each caller because the two APIs report differently. + + ``log_transport_failure=False`` suppresses the ``guardrail_failed_to_respond`` + entry for a non-200 that is re-raised as an ``HTTPException``, for callers that + own consolidated per-request logging. The ApplyGuardrail path needs this: + ``AsyncHTTPHandler.post`` calls ``raise_for_status()``, so every non-200 lands + in this handler, and one logical request can legitimately produce several of + them (a too-large probe, then each rejected bisection level) while still + succeeding overall. Logging per attempt would report a recovered request as + several failures plus a success. + + The connection-level branch below (timeout, endpoint down) still logs + unconditionally: it re-raises the original exception rather than an + ``HTTPException``, so no consolidating caller catches it, and suppressing it + would drop the only record of the failure. """ try: return await self.async_handler.post( @@ -882,16 +1630,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): status_code, detail_message, ) = self._parse_bedrock_guardrail_error_response(err_response) - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_provider=self.guardrail_provider, - guardrail_json_response={"error": detail_message}, - request_data=request_data or {}, - guardrail_status="guardrail_failed_to_respond", - start_time=start_time.timestamp(), - end_time=datetime.now(timezone.utc).timestamp(), - duration=(datetime.now(timezone.utc) - start_time).total_seconds(), - event_type=event_type, - ) + if log_transport_failure: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={ # mutable-ok: logging helper requires a dict + "error": detail_message + }, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + ) raise HTTPException(status_code=status_code, detail=detail_message) from e except HTTPException: raise @@ -900,7 +1651,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1027,7 +1778,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": detail_message}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1043,7 +1794,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1061,7 +1812,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response), - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status=self._get_invoke_checks_status(bool(violations)), start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 190d19f3d52..0d23e19f88d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -20,6 +20,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): content_filter_threshold=litellm_params.content_filter_threshold, prompt_attack_threshold=litellm_params.prompt_attack_threshold, pii_confidence_threshold=litellm_params.pii_confidence_threshold, + chunk_budget_chars=litellm_params.chunk_budget_chars, default_on=litellm_params.default_on, disable_exception_on_block=litellm_params.disable_exception_on_block, mask_request_content=litellm_params.mask_request_content, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6b354a39101..bbb6d758814 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,6 +5,7 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Required, TypedDict +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( AktoConfigModel, ) @@ -525,6 +526,15 @@ class BedrockGuardrailConfigModel(BaseModel): description="InvokeGuardrailChecks: block when any sensitiveInformation confidenceScore " ">= this value (scores are in [0,1]). Set to null to make PII detection detect-only.", ) + chunk_budget_chars: int = Field( + default=BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + gt=0, + description="ApplyGuardrail: batch size, in characters, used to re-send content after AWS " + "has rejected a request as too large. Requests AWS accepts are always sent in a single " + "call, so this has no effect until a rejection happens. Defaults to 25,000; a batch AWS " + "still rejects is bisected automatically, so this value only trades round trips against " + "batch size and cannot fail a request on its own.", + ) class LakeraV2GuardrailConfigModel(BaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d97bdc3532f..8d66b624341 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -28,6 +28,9 @@ class BedrockGuardrailUsage(TypedDict, total=False): sensitiveInformationPolicyUnits: int | None sensitiveInformationPolicyFreeUnits: int | None contextualGroundingPolicyUnits: int | None + contentPolicyImageUnits: int | None + automatedReasoningPolicyUnits: int | None + automatedReasoningPolicies: int | None class BedrockGuardrailOutput(TypedDict, total=False): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 76a695ce3fd..837fb93d331 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -7,6 +7,7 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException @@ -15,12 +16,18 @@ sys.path.insert(0, os.path.abspath("../../../../../..")) import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentChunkResult, BedrockGuardrail, _redact_pii_matches, ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentItem, + BedrockTextContent, +) from litellm.types.utils import CallTypes, ModelResponse @@ -53,9 +60,7 @@ async def test__redact_pii_matches_function(): redacted_response = _redact_pii_matches(response_with_pii) # Verify that PII matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" @@ -173,12 +178,8 @@ async def test__redact_pii_matches_multiple_assessments(): redacted_response = _redact_pii_matches(response_multiple_assessments) # Verify all PII in all assessments are redacted - assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ - "piiEntities" - ] + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"]["piiEntities"] assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" @@ -199,9 +200,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII mock_bedrock_response = MagicMock() @@ -239,20 +238,11 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" - ) as mock_debug, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug") as mock_debug, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method that should log the redacted response @@ -275,37 +265,23 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): bedrock_response_log_call = call break - assert ( - bedrock_response_log_call is not None - ), "Should have logged Bedrock AI response" + assert bedrock_response_log_call is not None, "Should have logged Bedrock AI response" # Extract the logged response data - logged_response = bedrock_response_log_call[0][ - 1 - ] # Second argument to debug call + logged_response = bedrock_response_log_call[0][1] # Second argument to debug call # Verify that the logged response has redacted PII assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "[REDACTED]" + logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) # Verify other fields are preserved assert logged_response["action"] == "GUARDRAIL_INTERVENED" - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["type"] - == "PHONE" - ) + assert logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["type"] == "PHONE" slg_list = request_data["metadata"]["standard_logging_guardrail_information"] assert ( - slg_list[0]["guardrail_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + slg_list[0]["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) @@ -319,9 +295,7 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII original_response_data = { @@ -361,17 +335,10 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method @@ -385,19 +352,12 @@ async def test_bedrock_guardrail_original_response_not_modified(): # (The json() method should return the original data) original_data = mock_bedrock_response.json() assert ( - original_data["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + original_data["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" ) # Verify that the returned BedrockGuardrailResponse contains original data - assert ( - result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "+1 412 555 1212" - ) + assert result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" print("Original response not modified test passed") @@ -454,18 +414,14 @@ async def test__redact_pii_matches_preserves_non_pii_entities(): redacted_response = _redact_pii_matches(response_with_mixed_data) # Verify that PII entity matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" # Verify that regex matches are also redacted (updated behavior) - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" @@ -496,9 +452,7 @@ async def test_pii_redaction_matches_debug_output_format(): "assessments": [ { "invocationMetrics": { - "guardrailCoverage": { - "textCharacters": {"guarded": 84, "total": 84} - }, + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, "guardrailProcessingLatency": 322, "usage": { "contentPolicyImageUnits": 0, @@ -553,9 +507,7 @@ async def test_pii_redaction_matches_debug_output_format(): redacted_response = _redact_pii_matches(original_response) # Verify the redacted response matches your expected debug output - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] # All PII matches should be redacted assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" @@ -570,34 +522,19 @@ async def test_pii_redaction_matches_debug_output_format(): assert pii_entities[0]["detected"] == True # Verify that the original response is unchanged - original_pii_entities = original_response["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"] - assert ( - original_pii_entities[0]["match"] == "John Smith" - ), "Original should be unchanged" - assert ( - original_pii_entities[1]["match"] == "324-12-3212" - ), "Original should be unchanged" - assert ( - original_pii_entities[2]["match"] == "607-456-7890" - ), "Original should be unchanged" + original_pii_entities = original_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert original_pii_entities[0]["match"] == "John Smith", "Original should be unchanged" + assert original_pii_entities[1]["match"] == "324-12-3212", "Original should be unchanged" + assert original_pii_entities[2]["match"] == "607-456-7890", "Original should be unchanged" # Verify all other metadata is preserved in redacted response assert redacted_response["action"] == "GUARDRAIL_INTERVENED" assert redacted_response["actionReason"] == "Guardrail blocked." assert redacted_response["blockedResponse"] == "Input blocked by PII policy" - assert ( - redacted_response["assessments"][0]["invocationMetrics"][ - "guardrailProcessingLatency" - ] - == 322 - ) + assert redacted_response["assessments"][0]["invocationMetrics"]["guardrailProcessingLatency"] == 322 print("PII redaction matches debug output format test passed") - print( - f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" - ) + print(f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}") print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") @@ -632,14 +569,10 @@ async def test__redact_pii_matches_with_regex_matches(): redacted_response = _redact_pii_matches(response_with_regex) # Verify that regex matches are redacted - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" - assert ( - regexes[1]["match"] == "[REDACTED]" - ), "Credit card regex match should be redacted" + assert regexes[1]["match"] == "[REDACTED]", "Credit card regex match should be redacted" # Verify other fields are preserved assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" @@ -648,13 +581,9 @@ async def test__redact_pii_matches_with_regex_matches(): assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" # Verify original response is unchanged - original_regexes = response_with_regex["assessments"][0][ - "sensitiveInformationPolicy" - ]["regexes"] + original_regexes = response_with_regex["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" - assert ( - original_regexes[1]["match"] == "4111-1111-1111-1111" - ), "Original should be unchanged" + assert original_regexes[1]["match"] == "4111-1111-1111-1111", "Original should be unchanged" print("Regex matches redaction test passed") @@ -690,31 +619,17 @@ async def test__redact_pii_matches_with_custom_words(): # Verify that custom word matches are redacted custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "First custom word match should be redacted" - assert ( - custom_words[1]["match"] == "[REDACTED]" - ), "Second custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "First custom word match should be redacted" + assert custom_words[1]["match"] == "[REDACTED]", "Second custom word match should be redacted" # Verify other fields are preserved - assert ( - custom_words[0]["action"] == "BLOCKED" - ), "Custom word action should be preserved" - assert ( - custom_words[1]["action"] == "ANONYMIZED" - ), "Custom word action should be preserved" + assert custom_words[0]["action"] == "BLOCKED", "Custom word action should be preserved" + assert custom_words[1]["action"] == "ANONYMIZED", "Custom word action should be preserved" # Verify original response is unchanged - original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ - "customWords" - ] - assert ( - original_custom_words[0]["match"] == "confidential_data" - ), "Original should be unchanged" - assert ( - original_custom_words[1]["match"] == "secret_information" - ), "Original should be unchanged" + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"]["customWords"] + assert original_custom_words[0]["match"] == "confidential_data", "Original should be unchanged" + assert original_custom_words[1]["match"] == "secret_information", "Original should be unchanged" print("Custom words redaction test passed") @@ -750,41 +665,21 @@ async def test__redact_pii_matches_with_managed_words(): redacted_response = _redact_pii_matches(response_with_managed_words) # Verify that managed word matches are redacted - managed_words = redacted_response["assessments"][0]["wordPolicy"][ - "managedWordLists" - ] + managed_words = redacted_response["assessments"][0]["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "First managed word match should be redacted" - assert ( - managed_words[1]["match"] == "[REDACTED]" - ), "Second managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "First managed word match should be redacted" + assert managed_words[1]["match"] == "[REDACTED]", "Second managed word match should be redacted" # Verify other fields are preserved - assert ( - managed_words[0]["action"] == "BLOCKED" - ), "Managed word action should be preserved" - assert ( - managed_words[0]["type"] == "PROFANITY" - ), "Managed word type should be preserved" - assert ( - managed_words[1]["action"] == "ANONYMIZED" - ), "Managed word action should be preserved" - assert ( - managed_words[1]["type"] == "HATE_SPEECH" - ), "Managed word type should be preserved" + assert managed_words[0]["action"] == "BLOCKED", "Managed word action should be preserved" + assert managed_words[0]["type"] == "PROFANITY", "Managed word type should be preserved" + assert managed_words[1]["action"] == "ANONYMIZED", "Managed word action should be preserved" + assert managed_words[1]["type"] == "HATE_SPEECH", "Managed word type should be preserved" # Verify original response is unchanged - original_managed_words = response_with_managed_words["assessments"][0][ - "wordPolicy" - ]["managedWordLists"] - assert ( - original_managed_words[0]["match"] == "inappropriate_word" - ), "Original should be unchanged" - assert ( - original_managed_words[1]["match"] == "offensive_term" - ), "Original should be unchanged" + original_managed_words = response_with_managed_words["assessments"][0]["wordPolicy"]["managedWordLists"] + assert original_managed_words[0]["match"] == "inappropriate_word", "Original should be unchanged" + assert original_managed_words[1]["match"] == "offensive_term", "Original should be unchanged" print("Managed words redaction test passed") @@ -841,9 +736,7 @@ async def test__redact_pii_matches_comprehensive_coverage(): # PII entities pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] - assert ( - pii_entities[0]["match"] == "[REDACTED]" - ), "PII entity match should be redacted" + assert pii_entities[0]["match"] == "[REDACTED]", "PII entity match should be redacted" # Regex matches regexes = assessment["sensitiveInformationPolicy"]["regexes"] @@ -851,15 +744,11 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Custom words custom_words = assessment["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "Custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "Custom word match should be redacted" # Managed words managed_words = assessment["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "Managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "Managed word match should be redacted" # Verify all other fields are preserved assert pii_entities[0]["type"] == "EMAIL" @@ -868,21 +757,10 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Verify original response is unchanged original_assessment = comprehensive_response["assessments"][0] - assert ( - original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] - == "user@example.com" - ) - assert ( - original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] - == "555-123-4567" - ) - assert ( - original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" - ) - assert ( - original_assessment["wordPolicy"]["managedWordLists"][0]["match"] - == "inappropriate" - ) + assert original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "user@example.com" + assert original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] == "555-123-4567" + assert original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + assert original_assessment["wordPolicy"]["managedWordLists"][0]["match"] == "inappropriate" print("Comprehensive coverage redaction test passed") @@ -914,9 +792,7 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method to avoid actual AWS credential loading - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -926,10 +802,12 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + ) print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") @@ -944,9 +822,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) # Create guardrail without explicit aws_bedrock_runtime_endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -960,9 +836,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -972,10 +846,10 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint from environment is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, f"Expected URL to contain env endpoint. Got: {prepped_request.url}" print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") @@ -988,9 +862,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) # Create guardrail without any custom endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -1004,9 +876,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey aws_region_name = "us-west-2" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1017,9 +887,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey # Verify that the default endpoint is used expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected default URL. Got: {prepped_request.url}" + assert prepped_request.url == expected_url, f"Expected default URL. Got: {prepped_request.url}" print(f"Default endpoint test passed. URL: {prepped_request.url}") @@ -1057,9 +925,7 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1069,10 +935,12 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch ) # Verify that the parameter takes precedence over environment variable - expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + expected_url = ( + f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + ) print(f"Parameter precedence test passed. URL: {prepped_request.url}") @@ -1081,14 +949,10 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" # Create a BedrockGuardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the make_bedrock_api_request method - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api_request: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api_request: # Test the apply_guardrail method with tool_calls in response inputs = { "texts": [], @@ -1115,14 +979,9 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): assert guardrailed_inputs is not None assert "tool_calls" in guardrailed_inputs assert len(guardrailed_inputs["tool_calls"]) == 1 - assert ( - guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" - ) + assert guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" - assert ( - guardrailed_inputs["tool_calls"][0]["function"]["arguments"] - == '{"location":"São Paulo"}' - ) + assert guardrailed_inputs["tool_calls"][0]["function"]["arguments"] == '{"location":"São Paulo"}' # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") @@ -1136,14 +995,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): policies (e.g. PII on model output) then returned action=NONE for non-streaming completions that go through unified_guardrail -> process_output_response. """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1168,14 +1023,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): @pytest.mark.asyncio async def test_bedrock_apply_guardrail_request_uses_INPUT_source(): """input_type='request' must call Bedrock with source=INPUT and user messages.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1258,12 +1109,8 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): # Mock AWS-related methods with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -1431,9 +1278,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" def _create_guardrail(self) -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") @pytest.mark.asyncio async def test_should_handle_all_null_policy_sub_lists(self): @@ -1554,9 +1399,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: { "sensitiveInformationPolicy": { "piiEntities": None, - "regexes": [ - {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} - ], + "regexes": [{"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"}], }, } ], @@ -1611,18 +1454,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_none_texts_in_inputs(self): """inputs[\"texts\"] is explicitly None — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {"texts": None} # Explicit None mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1645,18 +1484,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_missing_texts_key(self): """inputs has no \"texts\" key at all — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {} # No "texts" key mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1677,9 +1512,7 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Test 1: ANONYMIZED action should NOT raise exception anonymized_response = { @@ -1700,9 +1533,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): ], } - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) + should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" # Test 2: BLOCKED action should raise exception @@ -1710,13 +1541,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "I can't provide that information."}], "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } + {"topicPolicy": {"topics": [{"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}]}} ], } @@ -1738,19 +1563,13 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): } ] }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"}]}, } ], } should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" + assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" # Test 4: NONE action should not raise exception none_response = { @@ -1782,9 +1601,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): When logging_event_type is set, it must be forwarded to standard guardrail logging. When omitted, INPUT maps to pre_call (legacy). """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1795,13 +1612,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): mock_bedrock_response.json.return_value = { "action": "NONE", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ], } @@ -1811,12 +1622,8 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -1831,15 +1638,13 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): request_data=request_data, logging_event_type=GuardrailEventHooks.during_call, ) - assert ( - mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call - ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call # Raw Bedrock JSON is forwarded; redaction runs once in # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. assert ( - mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] == "GG" ) @@ -1855,9 +1660,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): @pytest.mark.asyncio async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1873,15 +1676,9 @@ async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): prepared_request.headers = {} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), - patch.object( - guardrail, "_prepare_request", return_value=prepared_request - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=prepared_request) as mock_prepare_request, patch.object( guardrail, "get_guardrail_dynamic_request_body_params", @@ -1933,9 +1730,7 @@ async def test_during_call_hook_invokes_bedrock_async_moderation_hook(): "model": "gpt-4", "messages": [{"role": "user", "content": "test"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="test_key", user_id="test_user" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), call_type="completion", ) finally: @@ -1991,11 +1786,7 @@ def test_extract_blocked_assessments_multiple_policies(): "action": "GUARDRAIL_INTERVENED", "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Investment", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Investment", "type": "DENY", "action": "BLOCKED"}]}, "contentPolicy": { "filters": [ { @@ -2006,9 +1797,7 @@ def test_extract_blocked_assessments_multiple_policies(): } ] }, - "wordPolicy": { - "customWords": [{"match": "forbidden", "action": "BLOCKED"}] - }, + "wordPolicy": {"customWords": [{"match": "forbidden", "action": "BLOCKED"}]}, } ], } @@ -2023,13 +1812,7 @@ def test_extract_blocked_assessments_only_anonymized_returns_empty(): response = { "action": "GUARDRAIL_INTERVENED", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } assert g._extract_blocked_assessments(response) == [] @@ -2049,23 +1832,14 @@ def test_get_http_exception_includes_assessments_and_identifier(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "BLOCKED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "BLOCKED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) assert isinstance(exc, HTTPException) assert exc.status_code == 400 assert exc.detail["error"] == "Violated guardrail policy" - assert ( - exc.detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question." - ) + assert exc.detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question." assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" assert exc.detail["guardrailVersion"] == "1" assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy" @@ -2088,15 +1862,11 @@ def test_extract_violation_category_names_mixed_policies(): {"name": "Tax Advice", "action": "BLOCKED"}, ] }, - "contentPolicy": { - "filters": [{"type": "VIOLENCE", "action": "BLOCKED"}] - }, + "contentPolicy": {"filters": [{"type": "VIOLENCE", "action": "BLOCKED"}]}, "wordPolicy": { "managedWordLists": [{"type": "PROFANITY", "action": "BLOCKED"}], }, - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}] - }, + "sensitiveInformationPolicy": {"piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}]}, } ], } @@ -2120,13 +1890,9 @@ def test_extract_violation_category_names_does_not_leak_user_input(): "assessments": [ { "wordPolicy": { - "customWords": [ - {"match": "secret-codeword-abc-123", "action": "BLOCKED"} - ], - }, - "sensitiveInformationPolicy": { - "regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}] + "customWords": [{"match": "secret-codeword-abc-123", "action": "BLOCKED"}], }, + "sensitiveInformationPolicy": {"regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}]}, } ], } @@ -2166,13 +1932,7 @@ def test_extract_violation_category_names_skips_anonymized(): g = _make_guardrail() response = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}] - } - } - ], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}]}}], } assert g._extract_violation_category_names(response) == [] @@ -2190,9 +1950,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): ``tracing_detail`` so downstream loggers (OTEL, ...) can surface the raw provider verdict as a queryable attribute without re-parsing the redacted guardrail_response blob.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2202,13 +1960,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}]}}], } request_data = { @@ -2217,12 +1969,8 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2253,9 +2001,7 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): """If the Bedrock response omits ``action`` (older / partial payloads), the field must be left off ``tracing_detail`` rather than written as ``None`` — downstream code expects strings or absence, not nulls.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2266,12 +2012,8 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): mock_bedrock_response.json.return_value = {"assessments": []} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2300,13 +2042,7 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "blocked"}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) @@ -2370,9 +2106,7 @@ async def test_streaming_post_call_only_runs_output_scan(): yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: out = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), @@ -2382,18 +2116,11 @@ async def test_streaming_post_call_only_runs_output_scan(): out.append(chunk) assert len(out) >= 1 - output_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT" - ] + output_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"] assert len(output_calls) == 1 assert output_calls[0].kwargs.get("request_data") is request_data - assert ( - output_calls[0].kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) - input_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT" - ] + assert output_calls[0].kwargs.get("logging_event_type") == GuardrailEventHooks.post_call + input_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT"] assert len(input_calls) == 0 @@ -2432,9 +2159,7 @@ async def test_streaming_post_call_output_only_path_passes_request_data_to_make_ yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: async for _ in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), response=mock_stream(), @@ -2488,9 +2213,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): ) minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), @@ -2499,10 +2222,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): sources = [c.kwargs.get("source") for c in mock_make.call_args_list] assert sources == ["OUTPUT"] - assert ( - mock_make.call_args.kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) + assert mock_make.call_args.kwargs.get("logging_event_type") == GuardrailEventHooks.post_call # --------------------------------------------------------------------------- @@ -2522,9 +2242,7 @@ _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo." def _grounding_guardrail() -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") def _grounding_messages() -> list: @@ -2556,27 +2274,19 @@ def _model_response(content: str) -> ModelResponse: # Expected OUTPUT content blocks, keyed by their grounding qualifier, so the # per-test assertions read as the block sequence they expect. -_GROUNDING_SOURCE_BLOCK = { - "text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]} -} +_GROUNDING_SOURCE_BLOCK = {"text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]}} _QUERY_BLOCK = {"text": {"text": _GROUNDING_QUERY_TEXT, "qualifiers": ["query"]}} -_GUARD_BLOCK = { - "text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]} -} +_GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]}} def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="INPUT", messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="OUTPUT", response=response, messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2600,9 +2310,7 @@ def test_grounding_input_leaves_existing_guarded_text_unqualified(): """An existing guarded_text input block keeps its legacy unqualified payload.""" expected_request = {"source": "INPUT", "content": [{"text": {"text": "policy"}}]} - actual_request = _input_request( - [{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}] - ) + actual_request = _input_request([{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}]) assert actual_request == expected_request @@ -2615,9 +2323,7 @@ def test_grounding_output_assembles_source_query_and_response(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } - actual_request = _output_request( - _grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(_grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2629,9 +2335,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): "content": [{"text": {"text": "Hi there."}}], } - actual_request = _output_request( - [{"role": "user", "content": "hello"}], _model_response("Hi there.") - ) + actual_request = _output_request([{"role": "user", "content": "hello"}], _model_response("Hi there.")) assert actual_request == expected_request @@ -2639,9 +2343,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." - uk_source_block = { - "text": {"text": uk_source_text, "qualifiers": ["grounding_source"]} - } + uk_source_block = {"text": {"text": uk_source_text, "qualifiers": ["grounding_source"]}} messages = [ { "role": "system", @@ -2662,9 +2364,7 @@ def test_grounding_output_combines_multiple_sources(): ], } - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2709,9 +2409,7 @@ def test_grounding_source_trusted_only_from_app_roles(role, is_trusted): if is_trusted: expected_content = [_GROUNDING_SOURCE_BLOCK, *expected_content] - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == {"source": "OUTPUT", "content": expected_content} @@ -2748,12 +2446,8 @@ async def test_grounding_output_blocked_raises_400(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -2791,13 +2485,7 @@ def _blocked_bedrock_httpx_response() -> MagicMock: response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}]}}], } return response @@ -2820,12 +2508,8 @@ async def test_make_bedrock_api_request_block_raises_modify_response_when_flag_s mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2857,12 +2541,8 @@ async def test_make_bedrock_api_request_block_raises_http_400_when_flag_unset(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2903,12 +2583,8 @@ async def test_async_pre_call_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2953,12 +2629,8 @@ async def test_async_moderation_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2998,12 +2670,8 @@ async def test_async_post_call_success_hook_attaches_original_response_on_block( mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -3032,9 +2700,7 @@ async def test_apply_guardrail_propagates_modify_response_on_block(): disable_exception_on_block=True, ) - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.side_effect = ModifyResponseException( message="Sorry, the model cannot answer this question.", model="bedrock-nova-micro", @@ -3276,6 +2942,1160 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n assert response is not None +def _too_large_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = { + "message": "Input is too long. Content size exceeds the maximum input size in text units." + } + response.text = json.dumps(response.json.return_value) + return response + + +def _other_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = {"message": "guardrailIdentifier is not valid"} + response.text = json.dumps(response.json.return_value) + return response + + +def _throttling_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 429 + response.json.return_value = {"message": "Rate exceeded"} + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _too_large_throttling_httpx_response() -> MagicMock: + """The shape AWS actually returns for an oversized ApplyGuardrail request when + the guardrail has an active content-filter policy: a 429 ThrottlingException, + not the documented 400 ValidationException. Message taken from a live call.""" + response = MagicMock() + response.status_code = 429 + response.json.return_value = { + "message": ( + "Input text size (3273 text units) exceeds the maximum allowed " + "(1000 text units) for the content filter policy (Classic tier)." + ) + } + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _passing_bedrock_httpx_response(marker: str) -> MagicMock: + """A successful ApplyGuardrail response tagged with `marker` so tests can + verify which chunk produced which output/usage after merging.""" + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "NONE", + "outputs": [{"text": marker}], + "assessments": [], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _blocking_bedrock_httpx_response(marker: str) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": marker}], + "assessments": [{"topicPolicy": {"topics": [{"name": marker, "type": "DENY", "action": "BLOCKED"}]}}], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _bedrock_guardrail_for_chunk_tests() -> "BedrockGuardrail": + return BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunks_on_too_large_validation_error(): + """A too-large 400 on the whole-content call must trigger a bisect-and-retry, + and the two chunk responses must be merged (assessments concatenated, usage + summed, outputs concatenated) rather than losing either half's result.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first half of a very long message"}, + {"role": "user", "content": "second half of a very long message"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _blocking_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + detail = exc_info.value.detail + assert exc_info.value.status_code == 400 + assert "chunk-2" in detail["bedrock_guardrail_response"] + assert detail["assessments"][0]["matches"][0]["name"] == "chunk-2" + + +@pytest.mark.asyncio +async def test_apply_guardrail_merges_usage_and_outputs_across_chunks_when_both_pass(): + """When both chunks pass clean, the merged response must still carry both + chunks' outputs/usage forward (needed for accurate logging/telemetry) and + must not itself raise.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + assert result.get("usage", {}).get("contentPolicyUnits") == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_recurses_past_first_bisection_into_four_chunks(): + """A payload that is still too large after one bisection must keep splitting + -- chunking is not capped at two pieces. Four messages where both the + whole-content call AND both first-level halves are too large must recurse + one level deeper into four chunks that all fit, not give up after the + first split.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "message one"}, + {"role": "user", "content": "message two"}, + {"role": "user", "content": "message three"}, + {"role": "user", "content": "message four"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole content: [1,2,3,4] + _too_large_validation_httpx_response(), # first half: [1,2] + _passing_bedrock_httpx_response("message one"), + _passing_bedrock_httpx_response("message two"), + _too_large_validation_httpx_response(), # second half: [3,4] + _passing_bedrock_httpx_response("message three"), + _passing_bedrock_httpx_response("message four"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["message one", "message two", "message three", "message four"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_when_grounding_present(): + """Contextual-grounding requests are scored holistically against the whole + source; chunking them would silently produce misleading grounding scores. + A too-large error on a grounded request must propagate unchanged, not be + bisected.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_on_non_size_validation_error(): + """A 400 for an unrelated validation problem (e.g. a bad guardrail id) must + not trigger chunking -- retrying a bad-config error split into pieces would + just fail twice more and mask the real problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _other_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + assert "not valid" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_unsplittable_text_propagates_original_error(): + """A too-large error on content that has been bisected down to text too + short to split further (< 2 characters) must propagate the original error + rather than looping or crashing.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "a"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_single_item_splits_by_text_and_succeeds(): + """A too-large error on content that is already down to a single content + item must be bisected by that item's own text (not abandoned), so an + oversized single message can still be scanned successfully in halves.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "one giant single block of text"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + response = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 3 + assert response.get("action") == "NONE" + + +def _raised_bedrock_error(status_code: int, message: str) -> httpx.HTTPStatusError: + """A non-200 the way `AsyncHTTPHandler.post` actually surfaces it. + + That handler calls `response.raise_for_status()`, so in production a non-200 from + Bedrock arrives as a raised `httpx.HTTPStatusError` carrying the response, never + as a returned response object. Tests that return the response instead exercise a + branch real traffic never reaches. A real `httpx.Response` is used rather than a + MagicMock because the transport helper branches on + `isinstance(err_response, httpx.Response)`.""" + response = httpx.Response( + status_code=status_code, + json={"message": message}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail"), + ) + return httpx.HTTPStatusError(message, request=response.request, response=response) + + +_TOO_LARGE_MESSAGE = "Input is too long. Content size exceeds the maximum input size in text units." + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_once_when_client_raises_for_status(): + """The too-large attempt recovered by chunking must still produce exactly one + telemetry entry when the HTTP client raises for status, which is what really + happens: `AsyncHTTPHandler.post` calls `raise_for_status()`. + + Regression for per-attempt `guardrail_failed_to_respond` entries leaking out of + the transport helper on a request that ultimately succeeded, which made a + recovered request look like several failures plus a success.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["success"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_still_logs_once_when_client_raises(): + """Suppressing the transport helper's per-attempt logging must not swallow the only + record of a genuine failure: an unsplittable too-large request still has to produce + exactly one `guardrail_failed_to_respond` entry, not zero.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "x"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _post_side_effect(*_args, **_kwargs): + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["guardrail_failed_to_respond"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_single_item_split_twice_still_yields_one_output_per_item(): + """One oversized content item that needs two levels of text bisection ends up + as four text fragments, and all four must still collapse back into exactly + ONE output entry, because they all came from one original content item. + + Downstream masking (`_apply_masking_to_messages`) walks the merged outputs by + a running index across the original, unchunked message list, so emitting more + than one entry for a single message shifts every later message's masked text + onto the wrong message and drops the surplus. Regression for fragment + grouping assuming fragments only ever arrive as adjacent sibling *pairs*, + which holds for one bisection level but not for two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "aaaa bbbb cccc dddd eeee ffff gggg hhhh"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole single item + _too_large_validation_httpx_response(), # first half + _passing_bedrock_httpx_response("q1"), + _passing_bedrock_httpx_response("q2"), + _too_large_validation_httpx_response(), # second half + _passing_bedrock_httpx_response("q3"), + _passing_bedrock_httpx_response("q4"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["q1q2q3q4"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_retries_after_throttling_then_succeeds(): + """A chunk call throttled with a 429 must be retried with backoff and + eventually succeed, rather than surfacing the 429 to the caller.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _throttling_httpx_response() + if call_count == 3: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 4 + mock_sleep.assert_awaited() + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_exactly_once_as_success(): + """A too-large 400 that is recovered by chunking must not leave behind a + 'guardrail_failed_to_respond' telemetry entry for the initial oversized + attempt: the whole logical request (1 too-large attempt + 2 chunk + attempts here) must produce exactly one standard-logging entry, and it + must reflect the eventual success, not the transient too-large failure.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_logs_exactly_once_as_failed(): + """A too-large error that cannot be recovered (chunking disabled by + contextual grounding) must still log exactly once, as a failure -- not be + silently dropped by the chunking telemetry consolidation.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_merge_preserves_masking_position(): + """An earlier chunk that comes back clean (empty `outputs`) must not + shift a later chunk's masked text onto the wrong message. Regression for: + flattening outputs without positional metadata let a later chunk's PII + redaction get applied to the first message while the actual PII-bearing + message (in a later chunk) was forwarded unmasked.""" + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + mask_request_content=True, + ) + + request_data = { + "model": "bedrock-nova-micro", + "messages": [ + {"role": "user", "content": "clean chunk with nothing to mask"}, + {"role": "user", "content": "chunk with PII: John Doe"}, + ], + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + def _clean_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"action": "NONE", "assessments": []} + return response + + def _masked_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "chunk with PII: [NAME]"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "NAME", "match": "John Doe", "action": "ANONYMIZED"}] + } + } + ], + } + return response + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _clean_httpx_response() + return _masked_httpx_response() + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=request_data, + call_type="acompletion", + ) + + assert call_count == 3 + updated_messages = request_data["messages"] + assert updated_messages[0]["content"] == "clean chunk with nothing to mask" + assert updated_messages[1]["content"] == "chunk with PII: [NAME]" + + +@pytest.mark.asyncio +async def test_apply_guardrail_accepted_content_costs_exactly_one_call(): + """Content AWS accepts must cost exactly one ApplyGuardrail call, however far over + the chunk budget it is. Chunking is a recovery path, not something every request + pays for. Regression for: bin-packing eagerly on every request, which split + conversations AWS was happy to take whole and multiplied billed calls and guardrail + latency on traffic that never had a size problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [{"role": "user", "content": item_text} for _ in range(3)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + return _passing_bedrock_httpx_response(f"batch-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 1 + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_small_content_makes_exactly_one_call(): + """Content that fits entirely within the budget in a single batch must + make exactly one ApplyGuardrail call -- confirms bin-packing does not + introduce an extra probe call for the common (small-request) case.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "short message one"}, + {"role": "user", "content": "short message two"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _passing_bedrock_httpx_response("single-batch") + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_batch_under_budget_still_rejected_falls_back_to_bisection(): + """A batch that fits the budget guess but is still rejected by AWS as too large + (a lower real per-account/region/policy cap) must fall back to bisection for that + batch only, and any other batch from the same request that AWS already accepted + must not be re-sent. + + Three half-budget items pack into two batches once the whole-payload probe is + rejected, so the sequence is probe, batch one (rejected), its two halves, batch + two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [ + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count in (1, 2): + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 5 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-3", "chunk-4", "chunk-5"] + + +def test_split_bedrock_content_single_item_splits_on_whitespace_not_mid_word(): + """A single content item whose raw character midpoint would fall inside a + word must instead split at the nearest whitespace, so neither fragment + ends or begins mid-token. Regression for the Veria AI review finding: a + denied word/PII pattern straddling a raw character-midpoint cut could be + truncated on both fragments and scan clean on each, then reassemble into + the original unmasked text -- a detection bypass.""" + text = ("a" * 20) + " " + ("b" * 30) + raw_midpoint = len(text) // 2 + assert text[raw_midpoint] == "b" + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + + assert first_text + second_text == text + assert first_text == ("a" * 20) + " " + assert second_text == "b" * 30 + + +def test_split_bedrock_content_single_item_with_no_whitespace_falls_back_to_midpoint(): + """A single giant token with no whitespace anywhere has no safe split + point, so the split must fall back to the raw character midpoint rather + than failing or looping.""" + text = "a" * 40 + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + assert first_text + second_text == text + assert len(first_text) == 20 + assert len(second_text) == 20 + + +def test_bin_pack_bedrock_content_packs_minimal_batches_within_budget(): + """Many medium items should pack into the minimal number of in-order + batches that each stay within budget, not one batch per item.""" + items = [BedrockContentItem(text=BedrockTextContent(text="x" * 30)) for _ in range(10)] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert sum(len(batch) for batch in batches) == 10 + for batch in batches: + combined_len = sum(len(item["text"]["text"]) for item in batch) + assert combined_len <= 100 + assert len(batches) == 4 + + +def test_bin_pack_bedrock_content_oversized_single_item_becomes_its_own_batch(): + """An item whose own text already exceeds the budget must not be + pre-split here -- it becomes its own oversized batch, and only the + reactive bisection fallback (on an AWS rejection) may split it later.""" + small_item = BedrockContentItem(text=BedrockTextContent(text="short")) + oversized_item = BedrockContentItem(text=BedrockTextContent(text="x" * 200)) + items = [small_item, oversized_item, small_item] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert batches == ((small_item,), (oversized_item,), (small_item,)) + + +def test_bin_pack_bedrock_content_empty_content_makes_exactly_one_empty_batch(): + """Empty content must still pack into exactly one (empty) batch, matching + pre-bin-packing behavior of sending the content list as-is in one call -- + bin-packing must not turn an empty request into zero ApplyGuardrail calls.""" + assert BedrockGuardrail._bin_pack_bedrock_content([], budget=100) == ((),) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_reported_as_429_bisects_without_burning_retries(): + """AWS reports an oversized ApplyGuardrail request as a 429 ThrottlingException + (not the documented 400 ValidationException) when the guardrail has an active + content-filter policy. That is not a transient throttle -- re-posting the same + oversized content can never succeed -- so it must bisect immediately instead of + consuming the exponential-backoff retry budget first. + + Regression for a bug found against a live guardrail: because the throttle retry + only keyed off status 429, every oversized chunk burned all + _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES attempts (each a billed AWS call, + each preceded by a backoff sleep) before bisection got a chance, at every level + of the recursion.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_throttling_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_sleep.assert_not_awaited() + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["half-2", "half-3"] + + +def test_chunk_budget_defaults_to_apply_guardrail_per_second_quota(): + """The default budget must track ApplyGuardrail's default quota of 25 text units + (about 1,000 characters each) per second. Packing to that size and posting + sequentially is what stops chunking from trading a size error for a throttle, so + this default is a deliberate match to AWS behaviour rather than an arbitrary + number.""" + assert BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS == 25_000 + assert BedrockGuardrail(guardrailIdentifier="g", guardrailVersion="DRAFT").chunk_budget_chars == 25_000 + + +@pytest.mark.asyncio +async def test_configured_chunk_budget_changes_how_content_is_packed(): + """An account with raised quotas can set a larger `chunk_budget_chars` and have it + actually drive packing once AWS has rejected a payload, spending fewer + ApplyGuardrail calls for the same content instead of being pinned to the + conservative default. + + Four 20,000-character messages are 80,000 characters total, and every call here is + preceded by the one whole-payload probe AWS rejects. At the 25,000 default only one + message fits per batch, so it is the probe plus four; at 50,000 two fit per batch, + so it is the probe plus two.""" + messages = [{"role": "user", "content": "x" * 20_000} for _ in range(4)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _calls_made_with_budget(budget: int) -> int: + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + chunk_budget_chars=budget, + ) + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + posted = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal posted + posted += 1 + if posted == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response("ok") + + mock_post.side_effect = _post_side_effect + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + return mock_post.await_count + + assert await _calls_made_with_budget(BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS) == 5 + assert await _calls_made_with_budget(50_000) == 3 + + +def test_split_index_never_produces_an_empty_fragment(): + """Both fragments must be non-empty for every splittable text, so bisection always + makes progress. + + A text whose only qualifying whitespace is its final character is the dangerous + shape: taking that boundary puts the split at len(text), leaving the first fragment + identical to the input that was just rejected and the second empty. The recursion + would then resubmit the unchanged fragment forever and exhaust the stack instead of + scanning or surfacing Bedrock's error.""" + for text in ("ab ", "xxxx ", ("x" * 40) + " ", " ab", "a b", "ab", " "): + split_at = BedrockGuardrail._nearest_whitespace_split_index(text) + assert 0 < split_at < len(text), f"degenerate split {split_at} for {text!r}" + assert text[:split_at] and text[split_at:], f"empty fragment for {text!r}" + assert text[:split_at] + text[split_at:] == text + + +@pytest.mark.asyncio +async def test_oversized_single_item_with_trailing_space_gives_up_instead_of_recursing(): + """An oversized single item whose only space is trailing must bottom out and + surface Bedrock's error, not recurse forever. + + AWS is modelled the way it really behaves, rejecting every attempt, because the + danger is a fragment identical to the input that was just rejected: AWS would + reject it again, and each retry would split it into the same unchanged fragment. + A split that always shrinks the text terminates and re-raises; one that can return + the whole text raises RecursionError instead. The call-count bound is generous: + halving 41 characters down to unsplittable is a handful of attempts, nowhere near + a stack limit.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + messages = [{"role": "user", "content": ("x" * 40) + " "}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = lambda *_a, **_k: _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as excinfo: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert excinfo.value.status_code == 400 + assert mock_post.await_count < 200 + + class TestBedrockOnlyScanNewMessages: """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. @@ -3559,14 +4379,10 @@ class TestBedrockIncrementalFlagInteractions: session = {"litellm_session_id": "sess-flags-mask"} with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1 mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" @pytest.mark.asyncio @@ -3586,9 +4402,7 @@ class TestBedrockIncrementalFlagInteractions: assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" assert result["texts"] == ["MASKED q1"], "masked content must be applied" mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" @pytest.mark.asyncio @@ -3647,9 +4461,7 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should } with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.return_value = MagicMock( - action="NONE", output=[], outputs=[], assessments=[] - ) + mock_api.return_value = MagicMock(action="NONE", output=[], outputs=[], assessments=[]) await guardrail.async_moderation_hook( data=data, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u"), @@ -3707,3 +4519,146 @@ class TestScanOnlyToolResultsWithLatestRoleFilter: assert result["texts"] == ["TOOL-RESULT"] warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) assert "scan_only_tool_results" in warning_text + + +@pytest.mark.parametrize("separator", ["\n", "\t", "\r\n", "　"]) +def test_split_bedrock_content_splits_on_any_whitespace_not_just_space(separator): + """Regression: the midpoint split must land on any Unicode whitespace, not only an + ASCII space. + + Matching only " " left the boundary unguarded for exactly the payloads that grow + large enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited. A deny-listed word sitting at the midpoint of one was cut + in half, scanned clean on both fragments, and reassembled intact, which is the + single-token detection bypass the whitespace split exists to close.""" + text = separator.join(["aaaaaaa"] * 4) + separator + "BADWORDXYZ" + separator + separator.join(["bbbbbbb"] * 4) + + first, second = BedrockGuardrail._split_bedrock_content([BedrockContentItem(text=BedrockTextContent(text=text))]) + + first_text = first[0]["text"]["text"] + second_text = second[0]["text"]["text"] + assert first_text + second_text == text, "split must stay lossless" + assert "BADWORDXYZ" in first_text or "BADWORDXYZ" in second_text, "split severed the token" + + +def test_merge_bedrock_responses_preserves_fields_the_merge_has_no_opinion_on(): + """Regression: merging must not drop AWS response fields it does not itself merge. + + The merged response used to be rebuilt from an empty dict holding only action, + outputs, assessments and usage, so actionReason, guardrailCoverage and anything AWS + adds later vanished from the guardrail_json_response the Admin UI renders, on every + ApplyGuardrail request rather than only chunked ones.""" + chunk = BedrockContentChunkResult( + response={ + "action": "NONE", + "actionReason": "No action.", + "guardrailCoverage": {"textCharacters": {"guarded": 41, "total": 41}}, + "usage": {"contentPolicyUnits": 1}, + }, + content=[BedrockContentItem(text=BedrockTextContent(text="hello"))], + fragment_group_size=1, + ) + + merged = BedrockGuardrail._merge_bedrock_guardrail_responses([chunk]) + + assert merged["actionReason"] == "No action." + assert merged["guardrailCoverage"] == {"textCharacters": {"guarded": 41, "total": 41}} + + +def test_merge_bedrock_usage_sums_counters_not_on_the_known_list(): + """Regression: usage counters were summed from a hardcoded list of six keys, so the + ones AWS also returns (contentPolicyImageUnits, the automatedReasoning pair) were + reported as absent no matter what the chunks actually used.""" + chunks = [ + BedrockContentChunkResult( + response={"action": "NONE", "usage": {"contentPolicyImageUnits": units, "contentPolicyUnits": 1}}, + content=[BedrockContentItem(text=BedrockTextContent(text="x"))], + fragment_group_size=1, + ) + for units in (3, 4) + ] + + usage = BedrockGuardrail._merge_bedrock_guardrail_responses(chunks)["usage"] + + assert usage["contentPolicyImageUnits"] == 7 + assert usage["contentPolicyUnits"] == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_exception_inside_200_logs_failure_and_proceeds(): + """Regression: AWS can report a failure inside an HTTP 200 body via Output.__type, + and that must be logged as guardrail_failed_to_respond rather than success. + + Real AWS does this: an unrecognised operation path on bedrock-runtime answers + HTTP 200 with {"Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}}. + Consolidating telemetry had replaced the derived status with a hardcoded "success", + which reported a failed scan as a clean one. The request itself still proceeds, which + is the behaviour of the code before chunking existed.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + exception_response = MagicMock() + exception_response.status_code = 200 + exception_response.json.return_value = { + "Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}, + "Version": "1.0", + } + exception_response.text = json.dumps(exception_response.json.return_value) + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.return_value = exception_response + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + assert result is not None, "the request proceeds, as it did before chunking existed" + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): + """Regression: the consolidated failure logger must log guardrail_json_response as a + dict, the shape the pre-chunking code and the InvokeGuardrailChecks path both use. + + Consolidating telemetry had changed it to a bare string on the ApplyGuardrail path + only, which breaks any consumer that reads it as a mapping and leaves the two paths + in this file inconsistent.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.side_effect = _raised_bedrock_error(400, "guardrailIdentifier is not valid") + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_log.assert_called_once() + logged = mock_log.call_args.kwargs["guardrail_json_response"] + assert isinstance(logged, dict), f"expected a dict, got {type(logged).__name__}" + assert "error" in logged diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 71e775842e3..8edb56ce25e 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -38,6 +38,41 @@ def test_initialize_presidio_guardrail(): assert result["litellm_params"].mode == "pre_call" +def test_initialize_bedrock_forwards_chunk_budget_chars(): + """Regression: `chunk_budget_chars` set in config.yaml must reach the guardrail. + + The field lives on BedrockGuardrailConfigModel, so LitellmParams parsed it and the + Admin UI rendered it, but initialize_bedrock enumerates its kwargs explicitly and + dropped it. The setting validated and then silently did nothing. Asserting through + initialize_guardrail rather than the constructor is the point: constructing + BedrockGuardrail directly bypasses the only path a user can actually reach. + """ + import litellm + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + test_guardrail = { + "guardrail_name": "test_bedrock_chunk_budget", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.BEDROCK.value, + "mode": "pre_call", + "guardrailIdentifier": "test-guardrail", + "guardrailVersion": "DRAFT", + "chunk_budget_chars": 60_000, + }, + } + + guardrail_handler = InMemoryGuardrailHandler() + guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_chunk_budget" + ] + assert initialized, "bedrock guardrail was not registered as a callback" + assert initialized[-1].chunk_budget_chars == 60_000 + + def test_initialize_guardrail_preserves_guardrail_info(): """ Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the From f05d468769e26879ec10f3d4ec8367b0fa503cd8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:57:26 -0700 Subject: [PATCH 51/74] fix(responses): forward allowed_openai_params through the chat completions bridge (#35885) Resolves #35878 Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/main.py | 1 + .../test_responses_api_bridge_flag.py | 39 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f923702119c..7b02c1b8023 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1064,6 +1064,7 @@ def responses( extra_headers=extra_headers, extra_body=extra_body, timeout=timeout if timeout is not None else request_timeout, + allowed_openai_params=allowed_openai_params, **kwargs, ) diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index 463af6562f1..f94c31831bf 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -16,6 +16,7 @@ sys.path.insert( import litellm from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import Choices, Message, ModelResponse, Usage class TestUseResponsesApiBridgeFlag: @@ -130,6 +131,44 @@ class TestUseResponsesApiBridgeFlag: mock_bridge_handler.assert_called_once() + @patch("litellm.acompletion") + @patch( + "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + ) + async def test_allowed_openai_params_forwarded_through_bridge( + self, mock_get_config, mock_acompletion + ): + """allowed_openai_params is a named param of responses(), so it must be + explicitly forwarded to the bridge; otherwise litellm.acompletion raises + UnsupportedParamsError for params the caller explicitly allowed.""" + mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() + mock_acompletion.return_value = ModelResponse( + id="chatcmpl_123", + model="openai/my-custom-model", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="Answer"), + finish_reason="stop", + ) + ], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + await litellm.aresponses( + model="openai/my-custom-model", + input="Hello", + use_chat_completions_api=True, + allowed_openai_params=["reasoning_effort"], + reasoning={"effort": "high"}, + litellm_logging_obj=MagicMock(), + ) + + mock_acompletion.assert_called_once() + assert mock_acompletion.call_args.kwargs.get("allowed_openai_params") == [ + "reasoning_effort" + ] + @patch("litellm.responses.file_search.emulated_handler._call_aresponses") @patch( "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" From 4de7a7443ac5506f422efb36e96387aaae185607 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:27:50 +0000 Subject: [PATCH 52/74] refactor(types): declare mirrored pricing fields on ModelInfo Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/router.py | 18 +++--- litellm/types/utils.py | 21 +++++-- tests/test_litellm/types/test_router.py | 73 +++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/types/test_router.py diff --git a/litellm/types/router.py b/litellm/types/router.py index 8b4b547bdcc..487d95cd762 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -17,7 +17,12 @@ from .completion import CompletionRequest from .embedding import EmbeddingRequest from .llms.openai import OpenAIFileObject from .search import SearchProvider -from .utils import CustomPricingLiteLLMParams, ModelResponse, StandardLoggingRoutingDecision +from .utils import ( + CustomPricingLiteLLMParams, + MirroredPricingParams, + ModelResponse, + StandardLoggingRoutingDecision, +) class ConfigurableClientsideParamsCustomAuth(TypedDict): @@ -122,7 +127,7 @@ class UpdateRouterConfig(BaseModel): model_config = ConfigDict(protected_namespaces=()) -class ModelInfo(BaseModel): +class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: datetime.datetime | None = None @@ -424,14 +429,7 @@ class DeploymentTypedDict(TypedDict, total=False): model_info: dict -SPECIAL_MODEL_INFO_PARAMS = [ - "input_cost_per_token", - "output_cost_per_token", - "input_cost_per_character", - "output_cost_per_character", - "cache_read_input_token_cost", - "cache_creation_input_token_cost", -] +SPECIAL_MODEL_INFO_PARAMS: Final = tuple(MirroredPricingParams.model_fields) class Deployment(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 35d4250782f..18cf9461648 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3245,10 +3245,23 @@ class StandardCallbackDynamicParams(TypedDict, total=False): litellm_disabled_callbacks: list[str] | None -class CustomPricingLiteLLMParams(BaseModel): - ## CUSTOM PRICING ## +class MirroredPricingParams(BaseModel): + """Pricing overrides that ``Deployment.__init__`` mirrors from ``litellm_params`` + onto ``model_info``, so both blobs hold the same rate. + + Declared once and inherited by both sides of that mirror, so the two can't drift. + """ + input_cost_per_token: float | None = None output_cost_per_token: float | None = None + input_cost_per_character: float | None = None + output_cost_per_character: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + + +class CustomPricingLiteLLMParams(MirroredPricingParams): + ## CUSTOM PRICING ## input_cost_per_second: float | None = None output_cost_per_second: float | None = None output_cost_per_second_1080p: float | None = None @@ -3259,7 +3272,6 @@ class CustomPricingLiteLLMParams(BaseModel): # This allows any model_info parameter to be set in litellm_params input_cost_per_token_flex: float | None = None input_cost_per_token_priority: float | None = None - cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None cache_creation_input_token_cost_above_272k_tokens: float | None = None @@ -3268,7 +3280,6 @@ class CustomPricingLiteLLMParams(BaseModel): cache_creation_input_token_cost_flex: float | None = None cache_creation_input_token_cost_priority: float | None = None cache_creation_input_audio_token_cost: float | None = None - cache_read_input_token_cost: float | None = None cache_read_input_token_cost_flex: float | None = None cache_read_input_token_cost_priority: float | None = None cache_read_input_token_cost_above_200k_tokens: float | None = None @@ -3276,7 +3287,6 @@ class CustomPricingLiteLLMParams(BaseModel): cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None cache_read_input_audio_token_cost: float | None = None - input_cost_per_character: float | None = None input_cost_per_character_above_128k_tokens: float | None = None input_cost_per_audio_token: float | None = None input_cost_per_token_cache_hit: float | None = None @@ -3298,7 +3308,6 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_token_batches: float | None = None output_cost_per_token_flex: float | None = None output_cost_per_token_priority: float | None = None - output_cost_per_character: float | None = None output_cost_per_audio_token: float | None = None output_cost_per_token_above_128k_tokens: float | None = None output_cost_per_token_above_200k_tokens: float | None = None diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py new file mode 100644 index 00000000000..1b66863a82f --- /dev/null +++ b/tests/test_litellm/types/test_router.py @@ -0,0 +1,73 @@ +import pytest + +from litellm.types.router import ( + SPECIAL_MODEL_INFO_PARAMS, + Deployment, + LiteLLM_Params, + ModelInfo, +) +from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams + + +def test_model_info_declares_mirrored_pricing_fields(): + """The pricing keys Deployment mirrors onto model_info must be declared fields, not + extras that only survive because ModelInfo sets extra="allow".""" + for field in SPECIAL_MODEL_INFO_PARAMS: + assert field in ModelInfo.model_fields + + info = ModelInfo(id="x", input_cost_per_token=1e-06) + assert info.__pydantic_extra__ == {} + assert info.input_cost_per_token == 1e-06 + + +def test_special_model_info_params_cannot_drift_from_the_mirror(): + assert SPECIAL_MODEL_INFO_PARAMS == tuple(MirroredPricingParams.model_fields) + assert set(SPECIAL_MODEL_INFO_PARAMS) <= set(CustomPricingLiteLLMParams.model_fields) + assert set(SPECIAL_MODEL_INFO_PARAMS) <= set(LiteLLM_Params.model_fields) + + +def test_custom_pricing_params_keeps_every_field_it_had(): + """The mirrored fields moved to a base class; none of them may go missing from + CustomPricingLiteLLMParams, whose model_fields drive custom-pricing detection.""" + for field in ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_character", + "output_cost_per_character", + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "input_cost_per_second", + "cache_read_input_token_cost_flex", + "input_cost_per_character_above_128k_tokens", + "output_cost_per_audio_token", + ): + assert field in CustomPricingLiteLLMParams.model_fields + + +@pytest.mark.parametrize("field", SPECIAL_MODEL_INFO_PARAMS) +def test_deployment_mirrors_pricing_from_litellm_params_onto_model_info(field): + deployment = Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="gpt-4o", **{field: 3e-06}), + ) + assert getattr(deployment.model_info, field) == 3e-06 + assert deployment.model_info.model_dump(exclude_none=True)[field] == 3e-06 + + +def test_unset_pricing_is_still_absent_from_dumps(): + """/model/info responses and DB writes dump model_info with exclude_none=True, so + declaring the pricing fields must not start emitting ~6 null keys per deployment.""" + dumped = ModelInfo(id="x").model_dump(exclude_none=True) + assert [field for field in SPECIAL_MODEL_INFO_PARAMS if field in dumped] == [] + + +def test_pricing_strings_are_coerced_to_float(): + """Cost values arrive from the DB and the Admin UI as strings; they must land as + floats so cost calculation doesn't multiply a str.""" + info = ModelInfo(id="x", output_cost_per_token="0.000002") + assert info.output_cost_per_token == 2e-06 + + +def test_invalid_pricing_is_rejected(): + with pytest.raises(ValueError): + ModelInfo(id="x", input_cost_per_token="free") From 24ac999cf7f74d7b7495d07f7f7783d691877e50 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 20:06:41 +0000 Subject: [PATCH 53/74] fix(types): drop Final on SPECIAL_MODEL_INFO_PARAMS for star-import rebinding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/types/router.py b/litellm/types/router.py index 487d95cd762..4280da08cbb 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -429,7 +429,7 @@ class DeploymentTypedDict(TypedDict, total=False): model_info: dict -SPECIAL_MODEL_INFO_PARAMS: Final = tuple(MirroredPricingParams.model_fields) +SPECIAL_MODEL_INFO_PARAMS = tuple(MirroredPricingParams.model_fields) class Deployment(BaseModel): From f668c1060981cd7698fb5946ab2bc708dc0f59a6 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 20:15:39 +0000 Subject: [PATCH 54/74] chore(ui): regenerate dashboard api types for ModelInfo pricing fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8e950874a10..a75c23da1cf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35294,6 +35294,10 @@ export interface components { base_model?: string | null; /** Blocked */ blocked?: boolean | null; + /** Cache Creation Input Token Cost */ + cache_creation_input_token_cost?: number | null; + /** Cache Read Input Token Cost */ + cache_read_input_token_cost?: number | null; /** Created At */ created_at?: string | null; /** Created By */ @@ -35305,6 +35309,14 @@ export interface components { db_model: boolean; /** Id */ id: string | null; + /** Input Cost Per Character */ + input_cost_per_character?: number | null; + /** Input Cost Per Token */ + input_cost_per_token?: number | null; + /** Output Cost Per Character */ + output_cost_per_character?: number | null; + /** Output Cost Per Token */ + output_cost_per_token?: number | null; /** Team Id */ team_id?: string | null; /** Team Public Model Name */ From ede84eee15ab0703598f87a8fce2612e406d49e9 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 03:25:38 +0000 Subject: [PATCH 55/74] ci: give the remaining pull_request workflows a concurrency group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/check-schema-sync.yml | 4 ++++ .github/workflows/conventional-commits.yml | 4 ++++ .github/workflows/guard-fork-dependencies.yml | 4 ++++ .github/workflows/helm_unit_test.yml | 4 ++++ .github/workflows/test-linting.yml | 4 ++++ .github/workflows/test-litellm-ui-build.yml | 4 ++++ .github/workflows/test-litellm-ui-lint.yml | 4 ++++ .github/workflows/test-mcp.yml | 4 ++++ .github/workflows/test-model-map.yaml | 4 ++++ 9 files changed, 36 insertions(+) diff --git a/.github/workflows/check-schema-sync.yml b/.github/workflows/check-schema-sync.yml index 0e5e2804e60..a4e78d2c44c 100644 --- a/.github/workflows/check-schema-sync.yml +++ b/.github/workflows/check-schema-sync.yml @@ -10,6 +10,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: check-sync: name: Verify schema.prisma copies match root diff --git a/.github/workflows/conventional-commits.yml b/.github/workflows/conventional-commits.yml index 69ade24d028..eb9eb69f8b6 100644 --- a/.github/workflows/conventional-commits.yml +++ b/.github/workflows/conventional-commits.yml @@ -14,6 +14,10 @@ on: permissions: pull-requests: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: lint-pr-title: name: Validate PR title diff --git a/.github/workflows/guard-fork-dependencies.yml b/.github/workflows/guard-fork-dependencies.yml index f4cbdd63cdf..6b366da78d4 100644 --- a/.github/workflows/guard-fork-dependencies.yml +++ b/.github/workflows/guard-fork-dependencies.yml @@ -15,6 +15,10 @@ on: permissions: {} +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: guard: name: Block fork dependency changes diff --git a/.github/workflows/helm_unit_test.yml b/.github/workflows/helm_unit_test.yml index a44d412c781..f95848945a0 100644 --- a/.github/workflows/helm_unit_test.yml +++ b/.github/workflows/helm_unit_test.yml @@ -9,6 +9,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: unit-test: runs-on: ubuntu-latest diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5e333f2a3ca..3db3fb07a94 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: lint: runs-on: ubuntu-latest diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index 39f4bc1428a..618b0195b5a 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -10,6 +10,10 @@ on: - litellm_oss_staging - "litellm_**" +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: build-ui: runs-on: ubuntu-latest diff --git a/.github/workflows/test-litellm-ui-lint.yml b/.github/workflows/test-litellm-ui-lint.yml index ecc739a87e2..e03d89ee26a 100644 --- a/.github/workflows/test-litellm-ui-lint.yml +++ b/.github/workflows/test-litellm-ui-lint.yml @@ -10,6 +10,10 @@ on: - litellm_oss_staging - "litellm_**" +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: frontend-lint: runs-on: ubuntu-latest diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index a5a4e722133..05cc13d0af2 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: test: runs-on: ubuntu-latest diff --git a/.github/workflows/test-model-map.yaml b/.github/workflows/test-model-map.yaml index cf4b0eb21a1..c2770e5da4c 100644 --- a/.github/workflows/test-model-map.yaml +++ b/.github/workflows/test-model-map.yaml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: validate-model-prices-json: runs-on: ubuntu-latest From 557d14cc71f7a0c1b44fa36f20b0da4cf9330e47 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:27:00 -0700 Subject: [PATCH 56/74] fix(lint): make strict-gate noqas survive base ruff and flag stale ones --- litellm/utils.py | 2 +- ruff-strict-budget.json | 2 +- ruff-strict.toml | 8 ++++++++ ruff.toml | 4 +++- 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index cdf1cc3cf23..911de83b785 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5769,7 +5769,7 @@ def json_schema_type(python_type_name: str): return python_to_json_schema_types.get(python_type_name, "string") -def function_to_dict(input_function) -> dict: # noqa: C901 +def function_to_dict(input_function) -> dict: """Using type hints and numpy-styled docstring, produce a dictionary usable for OpenAI function calling diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 8ff4bfb36c0..a8af4eabb3f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -252,7 +252,7 @@ "limit": 67 }, "RUF100": { - "limit": 100 + "limit": 0 }, "S110": { "limit": 218 diff --git a/ruff-strict.toml b/ruff-strict.toml index 01faf04805f..d20fb2e4d7d 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -4,6 +4,14 @@ extend = "ruff.toml" preview = true select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] extend-select = [] +# Overrides the inherited list: rules this gate enforces itself must NOT be external here, +# so this config's RUF100 flags their stale `# noqa` directives. What remains external is +# only what other tooling enforces: upstream litellm's ruff config, plus the base ruff.toml +# rules (T20, E7xx/F5xx/F8xx) this select list doesn't re-enable. +external = [ + "T20", "E731", "F541", "F841", + "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", +] [lint.mccabe] max-complexity = 15 diff --git a/ruff.toml b/ruff.toml index 095e3e24c52..00743e0f38a 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,9 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", "TID251", + "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "FURB", "I001", "LOG015", "N999", "PERF", + "PIE", "PL", "PYI", "RET", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", + "RUF046", "RUF051", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] From c3536c29a0ebb8e5ad63663d6e0798f80eb5ac9c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:38:53 -0700 Subject: [PATCH 57/74] fix(lint): cover every base-owned ruff rule in the strict gate's external list --- ruff-strict.toml | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/ruff-strict.toml b/ruff-strict.toml index d20fb2e4d7d..974c49c787b 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -6,10 +6,13 @@ select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B0 extend-select = [] # Overrides the inherited list: rules this gate enforces itself must NOT be external here, # so this config's RUF100 flags their stale `# noqa` directives. What remains external is -# only what other tooling enforces: upstream litellm's ruff config, plus the base ruff.toml -# rules (T20, E7xx/F5xx/F8xx) this select list doesn't re-enable. +# only what other tooling enforces: every base ruff.toml rule this select list doesn't +# re-enable (all of the default E/F families plus T20/PGH004/RUF008/RUF009, minus the +# strict-selected F401 and RUF100; F4 is split out so stale F401 noqas stay detectable), +# plus upstream litellm's ruff config. external = [ - "T20", "E731", "F541", "F841", + "T20", "PGH004", "RUF008", "RUF009", "E4", "E7", "E9", + "F402", "F404", "F406", "F407", "F5", "F6", "F7", "F8", "F9", "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] From 7bffbbd1f2132a00206b6864ddd102cfade9a94f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:00:50 -0700 Subject: [PATCH 58/74] refactor(vertex_ai): drop unreachable post-path status checks in batches handler HTTPHandler.post and AsyncHTTPHandler.post call raise_for_status before returning, so the status_code != 200 branches after the create and cancel POSTs could never run. Non-2xx already surfaces as httpx.HTTPStatusError from inside the client. The checks after GETs stay: the get helpers return without raising. Tests that faked a non-raising POST response are replaced by HTTPStatusError propagation coverage. --- litellm/llms/vertex_ai/batches/handler.py | 22 +---- .../llms/vertex_ai/batches/test_handler.py | 92 +++---------------- 2 files changed, 16 insertions(+), 98 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 8e36c2a0faa..6481b67fad7 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -98,11 +98,6 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) - _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( response=_json_response @@ -132,10 +127,6 @@ class VertexAIBatchPrediction(VertexLLM): error_body[:1000], ) raise - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -473,7 +464,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler: Final = _get_httpx_client() try: - response: Final = sync_handler.post( + sync_handler.post( url=api_base, headers=headers, data=json.dumps({}), @@ -487,11 +478,6 @@ class VertexAIBatchPrediction(VertexLLM): ) raise - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) - # HTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = sync_handler.get( url=retrieve_api_base, @@ -525,7 +511,7 @@ class VertexAIBatchPrediction(VertexLLM): llm_provider=litellm.LlmProviders.VERTEX_AI, ) try: - response: Final = await client.post( + await client.post( url=api_base, headers=headers, data=json.dumps({}), @@ -538,10 +524,6 @@ class VertexAIBatchPrediction(VertexLLM): e.response.text[:1000], ) raise - if response.status_code != 200: - raise VertexAIError( - status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" - ) # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = await client.get( diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index b9fb5dfe3c5..9535bf17411 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -5,8 +5,10 @@ The handler is HTTP/auth glue around the (separately-tested) pure ``VertexAIBatchTransformation``. Each public method (create / retrieve / list / cancel) resolves a Vertex access token + URL, branches on ``_is_async`` (returning the coroutine in the async case, doing the sync HTTP call otherwise), -checks the HTTP status, and parses the JSON into ``LiteLLMBatch`` (or the OpenAI -list shape). +and parses the JSON into ``LiteLLMBatch`` (or the OpenAI list shape). POST-backed +calls rely on the client's ``raise_for_status`` (non-2xx surfaces as +``httpx.HTTPStatusError``); GET-backed calls return without raising, so the +handler checks their status codes itself. We mock only true I/O / auth seams: * ``_ensure_access_token`` - the Vertex credential seam. Returns a fixed @@ -20,7 +22,7 @@ We mock only true I/O / auth seams: what URL/headers/body, and that the response is parsed into the litellm type. Sibling seams are asserted NOT called where relevant. -The ``_is_async`` branch, status-code error paths, and the cancel +The ``_is_async`` branch, the error paths, and the cancel retrieve-after-cancel sequencing run for real. """ @@ -179,13 +181,19 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() -def test_create_batch_sync_non_200_raises(): +def test_create_batch_sync_httpstatuserror_propagates(): + """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the + sync create path must surface that error, not swallow it.""" h = _make_handler() client = MagicMock() - client.post.return_value = _http_response(status_code=500) + request = httpx.Request("POST", "https://x/batchPredictionJobs") + err_response = httpx.Response(status_code=500, request=request, text="boom") + client.post.side_effect = httpx.HTTPStatusError( + "boom", request=request, response=err_response + ) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(VertexAIError, match="Error: 500") as exc_info: + with pytest.raises(httpx.HTTPStatusError): h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -197,9 +205,6 @@ def test_create_batch_sync_non_200_raises(): max_retries=None, ) - assert exc_info.value.status_code == 500 - assert "error text" in str(exc_info.value) - def test_create_batch_input_file_id_without_model_raises_400_before_post(): """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" @@ -224,32 +229,6 @@ def test_create_batch_input_file_id_without_model_raises_400_before_post(): client.post.assert_not_called() -def test_create_batch_async_non_200_raises(): - h = _make_handler() - async_client = MagicMock() - async_client.post = AsyncMock(return_value=_http_response(status_code=403)) - - with ( - patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), - patch(f"{HMOD}.get_async_httpx_client", return_value=async_client), - ): - coro = h.create_batch( - _is_async=True, - create_batch_data=CREATE_DATA, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - with pytest.raises(VertexAIError, match="Error: 403") as exc_info: - _run(coro) - - assert exc_info.value.status_code == 403 - assert "error text" in str(exc_info.value) - - # =========================================================================== # # retrieve_batch # =========================================================================== # @@ -554,27 +533,6 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves(): assert post_kwargs["url"].endswith(":cancel") -def test_cancel_batch_sync_cancel_post_non_200_raises(): - h = _make_handler() - client = MagicMock() - client.post.return_value = _http_response(status_code=500) - - with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(VertexAIError, match="Error: 500"): - h.cancel_batch( - _is_async=False, - batch_id=BATCH_ID, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - # cancel POST failed -> retrieve GET must never fire - client.get.assert_not_called() - - def test_cancel_batch_sync_retrieve_non_200_raises(): h = _make_handler() client = MagicMock() @@ -791,28 +749,6 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): _run(coro) async_client.get.assert_not_awaited() - # (a2) cancel POST returns a plain non-200 (no exception) -> raises - async_client_post500 = MagicMock() - async_client_post500.post = AsyncMock(return_value=_http_response(status_code=500)) - async_client_post500.get = AsyncMock() - with ( - patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), - patch(f"{HMOD}.get_async_httpx_client", return_value=async_client_post500), - ): - coro = h.cancel_batch( - _is_async=True, - batch_id=BATCH_ID, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - with pytest.raises(VertexAIError, match="Error: 500"): - _run(coro) - async_client_post500.get.assert_not_awaited() - # (b) retrieve-after-cancel returns non-200 async_client2 = MagicMock() async_client2.post = AsyncMock(return_value=_http_response(json_body={})) From f304b7b19faa9743636965c79a9709fd5bd2b2d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:10:33 -0700 Subject: [PATCH 59/74] refactor(lint): graduate the 35 zero-violation strict rules into ruff.toml Every strict-gate rule whose budget ceiling was already 0 moves into the base config's lint.extend-select, so editors and ruff check --fix surface the diagnostics directly and the budget file shrinks to rules with real debt. Graduates stay in ruff-strict.toml's select so the strict RUF100 pass keeps policing their stale noqa directives, and base external entries they made redundant (FURB, I001, RUF010, RUF022, RUF023, RUF051) are dropped so base RUF100 polices those directly. UP037 had two violations hidden behind a star import; importing Literal explicitly fixes them so UP037 can graduate too. New drift tests pin the invariants: every strict-selected rule is budgeted or hard-failed by base, every base-owned rule stays visible to exactly one RUF100 pass, and graduated rules fail the normal ruff run. --- .../internal_user_endpoints.py | 2 +- ruff-strict-budget.json | 105 --------- ruff.toml | 22 +- tests/test_litellm/test_ruff_strict_gate.py | 220 +++++++++++++++++- 4 files changed, 237 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 640a735c916..abc5d3e53ff 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, cast +from typing import Any, Final, Literal, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..7e350c184af 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -56,9 +56,6 @@ "B026": { "limit": 3 }, - "B033": { - "limit": 0 - }, "BLE001": { "limit": 2924 }, @@ -113,18 +110,6 @@ "F401": { "limit": 17 }, - "FURB136": { - "limit": 0 - }, - "FURB168": { - "limit": 0 - }, - "FURB188": { - "limit": 0 - }, - "I001": { - "limit": 0 - }, "LOG015": { "limit": 5 }, @@ -137,18 +122,9 @@ "PERF401": { "limit": 12 }, - "PERF402": { - "limit": 0 - }, "PERF403": { "limit": 34 }, - "PIE790": { - "limit": 0 - }, - "PIE800": { - "limit": 0 - }, "PIE804": { "limit": 18 }, @@ -158,9 +134,6 @@ "PLC0206": { "limit": 26 }, - "PLC0208": { - "limit": 0 - }, "PLC0414": { "limit": 46 }, @@ -170,24 +143,12 @@ "PLR0206": { "limit": 1 }, - "PLR0402": { - "limit": 0 - }, "PLR1704": { "limit": 3 }, - "PLR1711": { - "limit": 0 - }, "PLR1714": { "limit": 257 }, - "PLR1730": { - "limit": 0 - }, - "PLR2044": { - "limit": 0 - }, "PLW0127": { "limit": 57 }, @@ -206,27 +167,12 @@ "PLW1510": { "limit": 2 }, - "PYI030": { - "limit": 0 - }, "PYI036": { "limit": 3 }, - "PYI041": { - "limit": 0 - }, - "PYI064": { - "limit": 0 - }, - "RET501": { - "limit": 0 - }, "RET504": { "limit": 177 }, - "RUF010": { - "limit": 0 - }, "RUF012": { "limit": 241 }, @@ -236,18 +182,9 @@ "RUF019": { "limit": 38 }, - "RUF022": { - "limit": 0 - }, - "RUF023": { - "limit": 0 - }, "RUF046": { "limit": 4 }, - "RUF051": { - "limit": 0 - }, "RUF059": { "limit": 67 }, @@ -272,18 +209,12 @@ "SIM113": { "limit": 3 }, - "SIM114": { - "limit": 0 - }, "SIM115": { "limit": 2 }, "SIM117": { "limit": 7 }, - "SIM118": { - "limit": 0 - }, "SIM201": { "limit": 1 }, @@ -302,9 +233,6 @@ "TC004": { "limit": 5 }, - "TC005": { - "limit": 0 - }, "TID251": { "limit": 1240 }, @@ -323,46 +251,13 @@ "TRY300": { "limit": 860 }, - "UP006": { - "limit": 0 - }, - "UP007": { - "limit": 0 - }, - "UP008": { - "limit": 0 - }, - "UP012": { - "limit": 0 - }, - "UP018": { - "limit": 0 - }, - "UP024": { - "limit": 0 - }, "UP028": { "limit": 2 }, "UP031": { "limit": 2 }, - "UP032": { - "limit": 0 - }, - "UP034": { - "limit": 0 - }, - "UP035": { - "limit": 0 - }, "UP036": { "limit": 1 - }, - "UP037": { - "limit": 0 - }, - "UP045": { - "limit": 0 } } diff --git a/ruff.toml b/ruff.toml index 00743e0f38a..9b90910b355 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,14 +1,26 @@ lint.ignore = ["F405", "E402", "F403"] -lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] +# The second group is the strict gate's graduates: rules the codebase already has zero +# violations of, so they hard-fail here instead of being ratcheted in ruff-strict-budget.json. +# That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot. +lint.extend-select = [ + "T20", "PGH004", "RUF008", "RUF009", "RUF100", + "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", + "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", + "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", "UP012", + "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", +] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external # so RUF100 leaves their directives alone: the strict gate (ruff-strict.toml) and upstream # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ - # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "FURB", "I001", "LOG015", "N999", "PERF", - "PIE", "PL", "PYI", "RET", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", - "RUF046", "RUF051", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", + # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml). + # Family entries whose every strict rule graduated into extend-select above (FURB), and + # standalone graduated codes (I001, RUF010, RUF022, RUF023, RUF051), are dropped so this + # config's RUF100 polices their directives itself. + "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "LOG015", "N999", "PERF", + "PIE", "PL", "PYI", "RET", "RUF012", "RUF015", "RUF019", + "RUF046", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index abdeb6feecc..206207acb09 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -1,16 +1,24 @@ import importlib.util +import json +import re +import shutil import subprocess +import sys +import tomllib from pathlib import Path import pytest -_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "ruff_strict_gate.py" +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "scripts" / "ruff_strict_gate.py" _spec = importlib.util.spec_from_file_location("ruff_strict_gate", _MODULE_PATH) gate = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(gate) Violation = gate.Violation +_ENABLED_BY_RUFF_DEFAULTS = frozenset({"F401"}) + def rule(name, limit): return {name: {"limit": limit}} @@ -151,3 +159,213 @@ def test_base_point_mid_merge_advances_to_the_merged_in_base_tip(tmp_path): repo, _, base_tip = _branched_repo(tmp_path) _git(repo, "merge", "--no-commit", "--no-ff", "main") assert gate.resolve_base_point("main", cwd=repo) == base_tip + + +def _lint_section(config_name: str) -> dict: + return tomllib.loads((_REPO_ROOT / config_name).read_text())["lint"] + + +def _base_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff.toml")["external"]) + + +def _strict_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff-strict.toml")["external"]) + + +def _strict_selected() -> frozenset: + return frozenset(_lint_section("ruff-strict.toml")["select"]) + + +def _prefix_covered(code: str, prefixes: tuple[str, ...]) -> bool: + return any(code.startswith(prefix) for prefix in prefixes) + + +def _selected_by_the_normal_config() -> frozenset: + return frozenset(_lint_section("ruff.toml")["extend-select"]) | _ENABLED_BY_RUFF_DEFAULTS + + +def _budgeted_rules() -> frozenset: + return frozenset(json.loads((_REPO_ROOT / "ruff-strict-budget.json").read_text())) + + +def _ruff_binary() -> str | None: + beside_interpreter = Path(sys.executable).with_name("ruff") + return str(beside_interpreter) if beside_interpreter.exists() else shutil.which("ruff") + + +_RUFF = _ruff_binary() +_needs_ruff = pytest.mark.skipif(_RUFF is None, reason="ruff is not installed in this environment") + + +def _ruff_output_for_noqa(code: str, *extra_args: str) -> str: + proc = subprocess.run( + [ + _RUFF, + "check", + "--no-cache", + "--stdin-filename", + "litellm/types/_external_probe.py", + *extra_args, + "-", + ], + cwd=_REPO_ROOT, + input=f"def _probe(x: int): # noqa: {code}\n return x\n", + capture_output=True, + text=True, + ) + return proc.stdout + + +def test_every_strict_gate_rule_is_protected_from_base_ruf100(): + unprotected = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _base_external()) + and selector not in _selected_by_the_normal_config() + ) + assert unprotected == frozenset(), ( + f"`ruff check` deletes any `# noqa` naming {sorted(unprotected)} as unused, so suppressing " + "one of those strict-gate rules breaks lint. Cover them in ruff.toml's lint.external or " + "enable them in its lint.extend-select." + ) + + +def test_every_selected_rule_keeps_stale_noqa_detection_somewhere(): + policed_by_strict = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _strict_external()) + ) + policed_by_base = frozenset( + selector + for selector in _selected_by_the_normal_config() + if not _prefix_covered(selector, _base_external()) + ) + shadowed = ( + _strict_selected() | _selected_by_the_normal_config() + ) - policed_by_strict - policed_by_base + assert shadowed == frozenset(), ( + f"no config's RUF100 can ever report a stale `# noqa` for {sorted(shadowed)}: every config " + "that selects each of them also shadows it with an external entry. Narrow the external " + "entry in ruff.toml or ruff-strict.toml." + ) + + +_BASE_OWNED_FAMILY = re.compile(r"E[479]\d+|F\d+|T20\d+") +_BASE_OWNED_SINGLES = frozenset({"PGH004", "RUF008", "RUF009", "RUF100"}) + + +@pytest.fixture(scope="module") +def all_ruff_rule_codes() -> frozenset: + listing = subprocess.run( + [_RUFF, "rule", "--all", "--output-format", "json"], + capture_output=True, + text=True, + ) + assert listing.returncode == 0, listing.stderr + return frozenset( + entry["code"] for entry in json.loads(listing.stdout) if "Removed" not in entry["status"] + ) + + +@_needs_ruff +def test_every_base_owned_rule_is_external_or_selected_in_the_strict_config(all_ruff_rule_codes): + base_owned = frozenset( + code + for code in all_ruff_rule_codes + if _BASE_OWNED_FAMILY.fullmatch(code) or code in _BASE_OWNED_SINGLES + ) + stranded = frozenset( + code + for code in base_owned + if code not in _strict_selected() and not _prefix_covered(code, _strict_external()) + ) + assert stranded == frozenset(), ( + f"the strict gate's RUF100 reads a valid `# noqa` for {sorted(stranded)} as unused, the " + "spurious-breach trap ruff-strict.toml's external override exists to prevent. Cover them " + "there." + ) + double_booked = frozenset( + code + for code in base_owned + if code in _strict_selected() and _prefix_covered(code, _strict_external()) + ) + assert double_booked == frozenset(), ( + f"{sorted(double_booked)} are selected by the strict config yet shadowed by its external " + "list, so their stale suppressions can never be reported. Narrow the external entry in " + "ruff-strict.toml." + ) + + +def test_every_budgeted_rule_is_one_the_gate_actually_measures(): + selectors = tuple(_lint_section("ruff-strict.toml")["select"]) + unmeasured = frozenset(code for code in _budgeted_rules() if not code.startswith(selectors)) + assert unmeasured == frozenset(), ( + f"the gate never counts {sorted(unmeasured)}, so their ceilings are dead config that reads " + "as coverage. Either select them in ruff-strict.toml or drop them from the budget." + ) + + +@_needs_ruff +def test_every_strict_selected_rule_is_budgeted_or_hard_failed_by_the_base_config(all_ruff_rule_codes): + strict_enabled = frozenset( + code + for code in all_ruff_rule_codes + if code.startswith(tuple(_lint_section("ruff-strict.toml")["select"])) + ) + base_hard_failed = tuple(_lint_section("ruff.toml")["extend-select"]) + unpoliced = frozenset( + code + for code in strict_enabled + if code not in _budgeted_rules() + and not code.startswith(base_hard_failed) + and code not in _ENABLED_BY_RUFF_DEFAULTS + ) + assert unpoliced == frozenset(), ( + f"nothing enforces {sorted(unpoliced)}: the gate skips rules missing from the budget, and " + "the base config does not hard-fail them. Re-add a budget ceiling or graduate them into " + "ruff.toml's lint.extend-select." + ) + + +@_needs_ruff +def test_a_noqa_for_a_strict_gate_rule_survives_the_normal_ruff_run(): + assert "RUF100" not in _ruff_output_for_noqa("ANN202") + + +@_needs_ruff +def test_the_external_list_is_what_saves_that_noqa(): + assert "RUF100" in _ruff_output_for_noqa("ANN202", "--config", "lint.external=[]") + + +@_needs_ruff +def test_a_stale_noqa_for_a_locally_enabled_rule_is_still_reported(): + assert "RUF100" in _ruff_output_for_noqa("F401") + + +def _ruff_output_for_source(source: str) -> str: + proc = subprocess.run( + [_RUFF, "check", "--no-cache", "--stdin-filename", "litellm/types/_graduate_probe.py", "-"], + cwd=_REPO_ROOT, + input=source, + capture_output=True, + text=True, + ) + return proc.stdout + + +_DEPRECATED_TYPING_ALIAS = "from typing import List # noqa: UP035\n\n\ndef _probe(x: List[int]) -> None: ...\n" + + +@_needs_ruff +def test_a_graduated_rule_now_fails_the_normal_ruff_run_instead_of_waiting_for_the_gate(): + assert "UP006" in _ruff_output_for_source(_DEPRECATED_TYPING_ALIAS) + + +@_needs_ruff +def test_a_graduated_rule_can_still_be_suppressed_without_tripping_unused_noqa(): + suppressed = _DEPRECATED_TYPING_ALIAS.replace("...\n", "... # noqa: UP006\n") + output = _ruff_output_for_source(suppressed) + assert "UP006" not in output + assert "RUF100" not in output From 5cd027cbbca7d731238bd870c01917e9dcf3af97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 7 Aug 2026 23:11:23 -0700 Subject: [PATCH 60/74] fix(lint): let the ratchet guard recognise a graduated rule A budget rule that graduates into a config's hard-fail select list rightly leaves the budget file, but the ratchet guard read any disappearance as a silently raised ceiling. Teach it the pairing between ruff-strict-budget.json and ruff.toml: a dropped rule is excused only when the paired config's lint.extend-select (minus lint.ignore) now hard-fails it, so deleting a rule without graduating it still trips the guard. --- scripts/budget_ratchet_check.py | 51 ++++++++++++++++--- .../test_litellm/test_budget_ratchet_check.py | 44 ++++++++++++++++ 2 files changed, 89 insertions(+), 6 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 10a78483643..e97cd1bca00 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -10,7 +10,10 @@ content at the merge-base with the target branch and fails (exits 1, red) if: * a rule was dropped from a budget (its ceiling effectively became infinite), or * an entire budget file was deleted. -New rules and lowered/equal limits are fine. +New rules and lowered/equal limits are fine. So is a rule that graduated: once a +paired config (ruff.toml for the ruff-strict budget) selects the rule outright it +hard-fails at the first violation, which is stricter than any ceiling the budget +could hold, so dropping its entry tightens the guard rather than removing it. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -30,7 +33,9 @@ import argparse import json import subprocess import sys +import tomllib from pathlib import Path +from types import MappingProxyType from typing import NamedTuple REPO_ROOT = Path(__file__).resolve().parent.parent @@ -40,6 +45,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "type-discipline-budget.json", "basedpyright-code-budget.json", ) +GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"}) class Regression(NamedTuple): @@ -106,24 +112,57 @@ def _limits(budget: dict) -> dict[str, int]: } +def selectors_hard_failed_by(lint: dict) -> tuple[str, ...]: + """A ruff `[lint]` table's selected codes, minus anything `ignore` turns back off. + + `lint.ignore` wins over `lint.extend-select` in ruff, so an ignored code is not + actually enforced and must not count as a graduation. + """ + ignored = tuple(lint.get("ignore", ())) + return tuple( + selector + for selector in lint.get("extend-select", ()) + if not (ignored and selector.startswith(ignored)) + ) + + +def graduated_selectors(rel: str) -> tuple[str, ...]: + """Selectors the budget's paired ruff config hard-fails, so its ceiling is moot.""" + config = GRADUATION_CONFIGS.get(rel) + if config is None or not (REPO_ROOT / config).exists(): + return () + return selectors_hard_failed_by( + tomllib.loads((REPO_ROOT / config).read_text()).get("lint", {}) + ) + + def _regression_detail( rule: str, base_limits: dict[str, int], head_limits: dict[str, int], + graduated: tuple[str, ...], ) -> str | None: - """Why `rule` regressed vs base, or None when it held flat or fell. + """Why `rule` regressed vs base, or None when it held flat, fell, or graduated. - A dropped rule is terminal; otherwise the only loosening left is a raised limit. + A dropped rule is terminal unless it graduated; otherwise the only loosening + left is a raised limit. """ base_limit = base_limits[rule] if rule not in head_limits: + if graduated and rule.startswith(graduated): + return None return f"rule dropped (limit {base_limit} -> removed)" if head_limits[rule] > base_limit: return f"limit raised {base_limit} -> {head_limits[rule]}" return None -def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: +def regressions_for( + rel: str, + base: dict | None, + head: dict | None, + graduated: tuple[str, ...] = (), +) -> list[Regression]: if base is None: return [] # new budget file: nothing to ratchet against yet if head is None: @@ -133,7 +172,7 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr return [ Regression(rel, rule, detail) for rule in sorted(base_limits) - if (detail := _regression_detail(rule, base_limits, head_limits)) is not None + if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None ] @@ -164,7 +203,7 @@ def main() -> int: print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") continue checked.append(rel) - regressions.extend(regressions_for(rel, base, head)) + regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) if regressions: print( diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 1972c1b6386..22d05f4d00d 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -68,6 +68,50 @@ def test_new_rule_in_head_is_clean(): assert ratchet.regressions_for("b.json", {}, {"new-rule": _spec_of(5)}) == [] +def test_dropped_rule_that_graduated_to_a_hard_failing_config_is_clean(): + base = {"UP006": _spec_of(0)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("UP006",)) == [] + + +def test_graduation_matches_by_prefix_like_ruff_selectors_do(): + base = {"ANN202": _spec_of(865)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("ANN",)) == [] + + +def test_an_unrelated_graduation_does_not_excuse_a_dropped_rule(): + base = {"C901": _spec_of(3)} + regs = ratchet.regressions_for("b.json", base, {}, graduated=("UP006", "SIM118")) + assert [r.rule for r in regs] == ["C901"] + assert "dropped" in regs[0].detail + + +def test_graduation_never_excuses_a_raised_limit(): + base = {"UP006": _spec_of(0)} + regs = ratchet.regressions_for("b.json", base, {"UP006": _spec_of(7)}, graduated=("UP006",)) + assert [r.rule for r in regs] == ["UP006"] + assert "0 -> 7" in regs[0].detail + + +def test_graduated_selectors_come_from_the_paired_ruff_config(): + selectors = ratchet.graduated_selectors("ruff-strict-budget.json") + assert "UP006" in selectors + assert "ANN" not in selectors + + +def test_budgets_without_a_paired_config_can_never_graduate(): + assert ratchet.graduated_selectors("type-discipline-budget.json") == () + assert ratchet.graduated_selectors("basedpyright-code-budget.json") == () + + +def test_a_selector_the_config_also_ignores_does_not_count_as_graduated(): + lint = {"ignore": ["UP006"], "extend-select": ["UP006", "SIM118"]} + assert ratchet.selectors_hard_failed_by(lint) == ("SIM118",) + + +def test_selectors_hard_failed_by_reads_a_config_with_no_ignore_list(): + assert ratchet.selectors_hard_failed_by({"extend-select": ["UP006"]}) == ("UP006",) + + def test_deleted_budget_file_is_a_regression(): regs = ratchet.regressions_for("b.json", {"LIT006": _spec_of(1)}, None) assert [r.rule for r in regs] == ["*"] From 5f7a663005bf3228f913d34cb372b712d318e7de Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:33:32 -0700 Subject: [PATCH 61/74] fix(proxy): enforce require_managed_files on every raw provider id route require_managed_files was only checked on upload, so raw provider ids still reached the batch, fine-tuning and vector store file routes. Ownership rows exist only for managed ids, so those requests were forwarded under shared credentials with no tenant check: knowing another tenant's id was enough to read, run against, cancel or delete their object. Generalise the file-id guard to validate_managed_id_requirement(resource_id, resource_kind) and call it on batch create/retrieve/cancel, fine-tuning create/retrieve/cancel (training_file and validation_file both) and the shared vector store file id resolver. Behaviour is unchanged when the setting is off. --- litellm/proxy/batches_endpoints/endpoints.py | 5 + .../proxy/fine_tuning_endpoints/endpoints.py | 8 + .../openai_files_endpoints/common_utils.py | 27 +- .../openai_files_endpoints/files_endpoints.py | 8 +- .../vector_store_files_endpoints/endpoints.py | 5 + .../proxy/batches_endpoints/test_endpoints.py | 117 +++++++++ .../proxy/fine_tuning_endpoints/__init__.py | 0 .../fine_tuning_endpoints/test_endpoints.py | 236 ++++++++++++++++++ .../test_files_endpoint.py | 8 +- .../vector_store_files_endpoints/__init__.py | 0 .../test_endpoints.py | 83 ++++++ 11 files changed, 479 insertions(+), 18 deletions(-) create mode 100644 tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py create mode 100644 tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f7c332f2849..c9f66c5a48b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_original_file_id, prepare_data_with_credentials, update_batch_in_database, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -176,6 +177,7 @@ async def create_batch( } input_file_id: Final = _create_batch_data.get("input_file_id", None) + validate_managed_id_requirement(resource_id=input_file_id, resource_kind="file") unified_file_id: str | Literal[False] = False model_from_file_id = None @@ -392,6 +394,7 @@ async def retrieve_batch( data: dict = {} try: + validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") model_from_id: Final = decode_model_from_file_id(batch_id) _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, @@ -840,6 +843,8 @@ async def cancel_batch( data: dict = {} try: + validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") + # Check for encoded batch ID with model info model_from_id: Final = decode_model_from_file_id(batch_id) diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index f8ffb77edb8..87e9895eed0 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.utils import LiteLLMFineTuningJob @@ -134,6 +135,11 @@ async def create_fine_tuning_job( ## CHECK IF MANAGED FILE ID unified_file_id: str | Literal[False] = False training_file: Final = fine_tuning_request.training_file + validate_managed_id_requirement(resource_id=training_file, resource_kind="file") + validate_managed_id_requirement( + resource_id=fine_tuning_request.validation_file, + resource_kind="file", + ) response: LiteLLMFineTuningJob | None = None if training_file: unified_file_id = _is_base64_encoded_unified_file_id(training_file) @@ -246,6 +252,7 @@ async def retrieve_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -513,6 +520,7 @@ async def cancel_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b26010c5597..143bd5bc6b5 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -881,17 +881,20 @@ def validate_managed_files_requirement( ) -def validate_managed_file_id_requirement(file_id: str) -> None: +def validate_managed_id_requirement( + resource_id: str | None, + resource_kind: Literal["file", "batch", "fine-tuning job"], +) -> None: """ - Enforce proxy-level managed files on the file read/delete routes when - ``litellm.require_managed_files`` is enabled. + Enforce proxy-level managed resources on every route that accepts a provider-issued id + when ``litellm.require_managed_files`` is enabled. - Ownership is only recorded for LiteLLM managed files, so a raw provider file id sent to - retrieve/content/delete is forwarded to the provider under shared credentials without any - tenant check; knowing another tenant's provider file id would be enough to read or delete it. + Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the + provider under shared credentials without any tenant check; knowing another tenant's provider + id would be enough to read, reuse, or destroy the object behind it. Raises: - HTTPException: 400 if ``file_id`` is not a LiteLLM managed file id. + HTTPException: 400 if ``resource_id`` is set and is not a LiteLLM managed id. """ from fastapi import HTTPException @@ -900,14 +903,18 @@ def validate_managed_file_id_requirement(file_id: str) -> None: if litellm.require_managed_files is not True: return - if _is_base64_encoded_unified_file_id(file_id): + if not resource_id: + return + + if _is_base64_encoded_unified_file_id(resource_id): return raise HTTPException( status_code=400, detail=( - "Raw provider file ids cannot be used when require_managed_files is enabled in " - "litellm_settings. Use the LiteLLM managed file id returned when the file was created." + f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " + f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " + f"{resource_kind} was created." ), ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 5c1a4441d0d..b5950c4f852 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -49,8 +49,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, - validate_managed_file_id_requirement, validate_managed_files_requirement, + validate_managed_id_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -613,7 +613,7 @@ async def get_file_content( data: dict = {"file_id": file_id} try: - validate_managed_file_id_requirement(file_id=file_id) + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) @@ -911,7 +911,7 @@ async def get_file( data: dict = {"file_id": file_id} try: - validate_managed_file_id_requirement(file_id=file_id) + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") custom_llm_provider: Final = ( provider @@ -1103,7 +1103,7 @@ async def delete_file( data: dict = {"file_id": file_id} try: - validate_managed_file_id_requirement(file_id=file_id) + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") custom_llm_provider: Final = ( provider diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 896b7ca33d7..f8fdf607292 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -65,6 +65,11 @@ def _update_request_data_with_managed_file_id( is_base64_encoded_unified_id, parse_unified_id, ) + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + validate_managed_id_requirement(resource_id=file_id, resource_kind="file") # First, check if this is a unified managed file ID (base64 encoded) decoded_id: Final = is_base64_encoded_unified_id(file_id) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index e758aa5ca7f..18ca604f88d 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2253,3 +2253,120 @@ async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_ha "vertex_location": "us-central1", "vertex_credentials": "/creds/customer-sa.json", } + + +# =========================================================================== # +# require_managed_files - raw provider ids must not reach the provider. # +# # +# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to # +# these routes is forwarded under the shared provider credentials with no # +# tenant check, so any caller who learns another tenant's id can read its # +# batch, reuse its file as batch input, or cancel its job. These lock the # +# guard on every batches route that accepts a caller-supplied id. # +# =========================================================================== # + + +def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness): + """A model-encoded id is client-forgeable and has no ownership row, so it is + not a managed file id and must be rejected like any other raw id.""" + set_body( + harness, + { + "input_file_id": AZURE_FILE_ID, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", False): + await call_create(harness) + + assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123" + + +@pytest.mark.asyncio +async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_retrieve(retrieve_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + retrieve_harness.litellm_aretrieve.assert_not_called() + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness): + with patch.object(litellm, "require_managed_files", True): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_cancel(cancel_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + cancel_harness.litellm_acancel.assert_not_called() + cancel_harness.router_acancel.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness): + with patch.object(litellm, "require_managed_files", True): + await call_cancel(cancel_harness, _unified_batch_id()) + + assert cancel_harness.router_acancel.call_count == 1 diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py b/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py new file mode 100644 index 00000000000..35c202057e7 --- /dev/null +++ b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py @@ -0,0 +1,236 @@ +""" +require_managed_files enforcement for litellm/proxy/fine_tuning_endpoints/endpoints.py + +Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to these +routes is forwarded to the provider under the shared proxy credentials with no tenant +check, so any caller who learns another tenant's file id can train on it, and any +caller who learns another tenant's job id can read or cancel it. + +Each test asserts BOTH that the request is rejected AND that every downstream provider +seam stayed untouched, so a guard that raises after the provider call would still fail. +""" + +import base64 +import os +import sys +from contextlib import ExitStack +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import Response + +import litellm +import litellm.proxy.fine_tuning_endpoints.endpoints as endpoints +import litellm.proxy.proxy_server as proxy_server +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.llms.openai import LiteLLMFineTuningJobCreate +from litellm.types.utils import LiteLLMFineTuningJob, SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" +RAW_JOB_ID = "ftjob-victim-abc123" + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _unified_job_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format("gpt-4o-mini-id", RAW_JOB_ID) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _job() -> LiteLLMFineTuningJob: + job = LiteLLMFineTuningJob( + id=RAW_JOB_ID, + created_at=1234567890, + fine_tuned_model=None, + finished_at=None, + hyperparameters={"n_epochs": 1}, + model="gpt-4o-mini", + object="fine_tuning.job", + organization_id="org-test", + result_files=[], + seed=0, + status="running", + trained_tokens=None, + training_file=RAW_FILE_ID, + validation_file=None, + ) + job._hidden_params = {} + return job + + +class FakeRequest: + def __init__(self): + self.headers = {} + self.query_params = {} + + async def json(self): + return {} + + +class Seams: + def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock]): + self.router = router + self.litellm_calls = litellm_calls + + def assert_no_provider_call(self) -> None: + for name, mock in self.litellm_calls.items(): + assert mock.call_count == 0, f"litellm.{name} was called" + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job"): + assert getattr(self.router, name).call_count == 0, f"router.{name} was called" + + +@pytest.fixture +def seams(): + logging = MagicMock(spec=ProxyLogging) + logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"]) + logging.post_call_failure_hook = AsyncMock() + logging.update_request_status = AsyncMock() + logging.get_proxy_hook = MagicMock(return_value=None) + + router = MagicMock(spec=Router) + router.acreate_fine_tuning_job = AsyncMock(return_value=_job()) + router.aretrieve_fine_tuning_job = AsyncMock(return_value=_job()) + router.acancel_fine_tuning_job = AsyncMock(return_value=_job()) + + litellm_calls = { + name: AsyncMock(return_value=_job()) + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job") + } + + with ExitStack() as stack: + stack.enter_context( + patch.object( + ProxyBaseLLMRequestProcessing, + "common_processing_pre_call_logic", + AsyncMock(side_effect=lambda self=None, **kw: (self.data if self else {}, MagicMock())), + ) + ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", MagicMock(return_value={}))) + for name, mock in litellm_calls.items(): + stack.enter_context(patch.object(litellm, name, mock)) + stack.enter_context(patch.object(proxy_server, "llm_router", router)) + stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) + stack.enter_context(patch.object(proxy_server, "premium_user", True)) + stack.enter_context(patch.object(proxy_server, "general_settings", {})) + stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) + stack.enter_context(patch.object(proxy_server, "version", "test-version")) + stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}])) + yield Seams(router=router, litellm_calls=litellm_calls) + + +async def _create(training_file: str, validation_file: str | None = None): + return await endpoints.create_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_request=LiteLLMFineTuningJobCreate( + model="gpt-4o-mini", + training_file=training_file, + validation_file=validation_file, + custom_llm_provider="openai", + ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _retrieve(job_id: str): + return await endpoints.retrieve_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + custom_llm_provider="openai", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _cancel(job_id: str): + return await endpoints.cancel_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +@pytest.mark.asyncio +async def test_create__raw_training_file_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__raw_validation_file_rejected_when_managed_files_required(seams): + """The validation file is uploaded and readable exactly like the training file, + so a managed training_file must not smuggle a raw validation_file past the guard.""" + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id(), validation_file=RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__unified_training_file_allowed_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + await _create(_unified_file_id()) + + assert seams.router.acreate_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams): + with patch.object(litellm, "require_managed_files", False): + await _create(RAW_FILE_ID) + + assert seams.litellm_calls["acreate_fine_tuning_job"].call_count == 1 + + +@pytest.mark.asyncio +async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _retrieve(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + await _retrieve(_unified_job_id()) + + assert seams.router.aretrieve_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _cancel(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + await _cancel(_unified_job_id()) + + assert seams.router.acancel_fine_tuning_job.call_count == 1 diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 24b814bae1f..2e15da7590a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3109,22 +3109,22 @@ def _unified_managed_file_id() -> str: def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( - validate_managed_file_id_requirement, + validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", True) - validate_managed_file_id_requirement(file_id=_unified_managed_file_id()) + validate_managed_id_requirement(resource_id=_unified_managed_file_id(), resource_kind="file") def test_managed_file_id_requirement_is_opt_in(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( - validate_managed_file_id_requirement, + validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", False) - validate_managed_file_id_requirement(file_id="file-victim-abc123") + validate_managed_id_requirement(resource_id="file-victim-abc123", resource_kind="file") def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py b/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py new file mode 100644 index 00000000000..c2bdb1d80f2 --- /dev/null +++ b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py @@ -0,0 +1,83 @@ +""" +require_managed_files enforcement for litellm/proxy/vector_store_files_endpoints/endpoints.py + +Every vector-store file route (create, retrieve, content, update, delete) resolves its +caller-supplied file id through _update_request_data_with_managed_file_id before the +provider call, so the guard lives there once and covers all five. + +A raw provider file id has no ownership row, so without the guard it is attached to a +vector store or read back under the shared provider credentials with no tenant check. +""" + +import base64 +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import HTTPException + +import litellm +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_managed_file_id, +) +from litellm.types.utils import SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _resolve(file_id: str): + return _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": file_id}, + file_id=file_id, + request=MagicMock(headers={}, query_params={}), + llm_router=None, + ) + + +def test_raw_file_id_rejected_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + _resolve(RAW_FILE_ID) + + assert exc.value.status_code == 400 + + +def test_model_encoded_file_id_rejected_when_managed_files_required(): + """encode_file_id_with_model output is client-forgeable and carries no ownership + row, so it is not a managed file id.""" + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + encoded = encode_file_id_with_model(RAW_FILE_ID, "gpt-4o-mini", id_type="file") + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + _resolve(encoded) + + assert exc.value.status_code == 400 + + +def test_unified_file_id_allowed_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + data, original = _resolve(_unified_file_id()) + + assert original == _unified_file_id() + assert data["file_id"] == RAW_FILE_ID + + +def test_raw_file_id_allowed_when_managed_files_not_required(): + with patch.object(litellm, "require_managed_files", False): + data, original = _resolve(RAW_FILE_ID) + + assert original is None + assert data["file_id"] == RAW_FILE_ID From b01eacd67c3c8828c464f29ff5fbaee32f0735ce Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:36:30 -0700 Subject: [PATCH 62/74] ci: run the new fine-tuning and vector store file test dirs --- .github/workflows/test-unit-proxy-endpoints.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 645996f779d..2ea3c521e8b 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -38,6 +38,8 @@ jobs: tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/batches_endpoints + tests/test_litellm/proxy/fine_tuning_endpoints + tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/image_endpoints From 8c0556abf6965fde0de260da5ce424aa1daa1a56 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 02:42:34 -0700 Subject: [PATCH 63/74] fix(proxy): authenticate managed ids before routing --- litellm/proxy/batches_endpoints/endpoints.py | 21 ++++- .../proxy/fine_tuning_endpoints/endpoints.py | 25 +++++- .../openai_files_endpoints/common_utils.py | 58 ++++++++++--- .../openai_files_endpoints/files_endpoints.py | 21 ++++- .../vector_store_files_endpoints/endpoints.py | 56 +++++++++--- .../proxy/batches_endpoints/test_endpoints.py | 56 ++++++++++++ .../fine_tuning_endpoints/test_endpoints.py | 46 +++++++++- .../test_files_endpoint.py | 36 +++++++- .../test_endpoints.py | 87 ++++++++++++++++--- 9 files changed, 355 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index c9f66c5a48b..aef1c5ac17e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -177,7 +177,12 @@ async def create_batch( } input_file_id: Final = _create_batch_data.get("input_file_id", None) - validate_managed_id_requirement(resource_id=input_file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=input_file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) unified_file_id: str | Literal[False] = False model_from_file_id = None @@ -394,7 +399,12 @@ async def retrieve_batch( data: dict = {} try: - validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) model_from_id: Final = decode_model_from_file_id(batch_id) _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, @@ -843,7 +853,12 @@ async def cancel_batch( data: dict = {} try: - validate_managed_id_requirement(resource_id=batch_id, resource_kind="batch") + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Check for encoded batch ID with model info model_from_id: Final = decode_model_from_file_id(batch_id) diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index 87e9895eed0..a13ad00713d 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -135,10 +135,17 @@ async def create_fine_tuning_job( ## CHECK IF MANAGED FILE ID unified_file_id: str | Literal[False] = False training_file: Final = fine_tuning_request.training_file - validate_managed_id_requirement(resource_id=training_file, resource_kind="file") - validate_managed_id_requirement( + await validate_managed_id_requirement( + resource_id=training_file, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + await validate_managed_id_requirement( resource_id=fine_tuning_request.validation_file, resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), ) response: LiteLLMFineTuningJob | None = None if training_file: @@ -252,7 +259,12 @@ async def retrieve_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") - validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -520,7 +532,12 @@ async def cancel_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") - validate_managed_id_requirement(resource_id=fine_tuning_job_id, resource_kind="fine-tuning job") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 143bd5bc6b5..56e986c89cf 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -4,7 +4,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -22,6 +22,21 @@ if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch +@runtime_checkable +class ManagedResourceAccessChecker(Protocol): + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + def _is_base64_encoded_unified_file_id(b64_uid: str) -> str | Literal[False]: # Ensure b64_uid is a string and not a mock object if not isinstance(b64_uid, str): @@ -881,20 +896,24 @@ def validate_managed_files_requirement( ) -def validate_managed_id_requirement( +async def validate_managed_id_requirement( resource_id: str | None, resource_kind: Literal["file", "batch", "fine-tuning job"], + user_api_key_dict: "UserAPIKeyAuth", + managed_files_obj: object | None, ) -> None: """ Enforce proxy-level managed resources on every route that accepts a provider-issued id - when ``litellm.require_managed_files`` is enabled. + when ``litellm.require_managed_files`` is enabled, and authenticate managed ids against + the caller's stored ownership record. Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the provider under shared credentials without any tenant check; knowing another tenant's provider id would be enough to read, reuse, or destroy the object behind it. Raises: - HTTPException: 400 if ``resource_id`` is set and is not a LiteLLM managed id. + HTTPException: 400 for a raw id, 403 for an inaccessible managed id, or 500 when + ownership validation is unavailable. """ from fastapi import HTTPException @@ -906,16 +925,33 @@ def validate_managed_id_requirement( if not resource_id: return - if _is_base64_encoded_unified_file_id(resource_id): + if not _is_base64_encoded_unified_file_id(resource_id): + raise HTTPException( + status_code=400, + detail=( + f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " + f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " + f"{resource_kind} was created." + ), + ) + + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException( + status_code=500, + detail="Managed resource ownership validation is unavailable.", + ) + + can_access: Final = ( + await managed_files_obj.can_user_call_unified_file_id(resource_id, user_api_key_dict) + if resource_kind == "file" + else await managed_files_obj.can_user_call_unified_object_id(resource_id, user_api_key_dict) + ) + if can_access: return raise HTTPException( - status_code=400, - detail=( - f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " - f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " - f"{resource_kind} was created." - ), + status_code=403, + detail=f"The caller does not have access to this managed {resource_kind} id.", ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b5950c4f852..0acaac3bf5d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -613,7 +613,12 @@ async def get_file_content( data: dict = {"file_id": file_id} try: - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) @@ -911,7 +916,12 @@ async def get_file( data: dict = {"file_id": file_id} try: - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) custom_llm_provider: Final = ( provider @@ -1103,7 +1113,12 @@ async def delete_file( data: dict = {"file_id": file_id} try: - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) custom_llm_provider: Final = ( provider diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index f8fdf607292..c9b89bcd390 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -30,10 +30,12 @@ if TYPE_CHECKING: router: Final = APIRouter() -def _update_request_data_with_managed_file_id( +async def _update_request_data_with_managed_file_id( data: dict, file_id: str, request: Request, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: object | None, llm_router: Optional["Router"] = None, ) -> tuple[dict, str | None]: """ @@ -69,7 +71,12 @@ def _update_request_data_with_managed_file_id( validate_managed_id_requirement, ) - validate_managed_id_requirement(resource_id=file_id, resource_kind="file") + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=managed_files_obj, + ) # First, check if this is a unified managed file ID (base64 encoded) decoded_id: Final = is_base64_encoded_unified_id(file_id) @@ -514,8 +521,13 @@ async def vector_store_file_create( # Handle managed file IDs if present in request body original_managed_file_id = None if "file_id" in data: - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=data["file_id"], request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=data["file_id"], + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -712,8 +724,13 @@ async def vector_store_file_retrieve( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -814,8 +831,13 @@ async def vector_store_file_content( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -916,8 +938,13 @@ async def vector_store_file_update( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -1018,8 +1045,13 @@ async def vector_store_file_delete( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 18ca604f88d..f9193db143e 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2275,6 +2275,37 @@ def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-pro return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") +def _unified_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + @pytest.mark.asyncio async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness): set_body( @@ -2334,6 +2365,27 @@ async def test_create__raw_input_file_id_allowed_when_managed_files_not_required assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123" +@pytest.mark.asyncio +async def test_create__other_teams_unified_input_file_id_rejected(harness): + set_body( + harness, + { + "input_file_id": _unified_file_id(), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "403" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + @pytest.mark.asyncio async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness): with patch.object(litellm, "require_managed_files", True): @@ -2347,6 +2399,8 @@ async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retri @pytest.mark.asyncio async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness): + retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await call_retrieve(retrieve_harness, _unified_batch_id()) @@ -2366,6 +2420,8 @@ async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_ @pytest.mark.asyncio async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness): + cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await call_cancel(cancel_harness, _unified_batch_id()) diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py index 35c202057e7..b54787bf428 100644 --- a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py @@ -14,6 +14,7 @@ import base64 import os import sys from contextlib import ExitStack +from dataclasses import dataclass from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -78,10 +79,31 @@ class FakeRequest: return {} +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + class Seams: - def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock]): + def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock], logging: MagicMock): self.router = router self.litellm_calls = litellm_calls + self.logging = logging def assert_no_provider_call(self) -> None: for name, mock in self.litellm_calls.items(): @@ -126,7 +148,7 @@ def seams(): stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) stack.enter_context(patch.object(proxy_server, "version", "test-version")) stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}])) - yield Seams(router=router, litellm_calls=litellm_calls) + yield Seams(router=router, litellm_calls=litellm_calls, logging=logging) async def _create(training_file: str, validation_file: str | None = None): @@ -176,6 +198,8 @@ async def test_create__raw_training_file_rejected_when_managed_files_required(se async def test_create__raw_validation_file_rejected_when_managed_files_required(seams): """The validation file is uploaded and readable exactly like the training file, so a managed training_file must not smuggle a raw validation_file past the guard.""" + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): with pytest.raises(ProxyException) as exc: await _create(_unified_file_id(), validation_file=RAW_FILE_ID) @@ -186,12 +210,26 @@ async def test_create__raw_validation_file_rejected_when_managed_files_required( @pytest.mark.asyncio async def test_create__unified_training_file_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await _create(_unified_file_id()) assert seams.router.acreate_fine_tuning_job.call_count == 1 +@pytest.mark.asyncio +async def test_create__other_teams_unified_training_file_rejected(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id()) + + assert exc.value.code == "403" + seams.assert_no_provider_call() + + @pytest.mark.asyncio async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams): with patch.object(litellm, "require_managed_files", False): @@ -212,6 +250,8 @@ async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams): @pytest.mark.asyncio async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await _retrieve(_unified_job_id()) @@ -230,6 +270,8 @@ async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams): @pytest.mark.asyncio async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + with patch.object(litellm, "require_managed_files", True): await _cancel(_unified_job_id()) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 2e15da7590a..e68e7102fce 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3107,24 +3107,52 @@ def _unified_managed_file_id() -> str: return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") -def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): +class _ManagedResourceAccessCheckerStub: + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return True + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return True + + +@pytest.mark.asyncio +async def test_require_managed_files_allows_owned_unified_managed_file_id(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", True) - validate_managed_id_requirement(resource_id=_unified_managed_file_id(), resource_kind="file") + await validate_managed_id_requirement( + resource_id=_unified_managed_file_id(), + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="owner-user"), + managed_files_obj=_ManagedResourceAccessCheckerStub(), + ) -def test_managed_file_id_requirement_is_opt_in(monkeypatch): +@pytest.mark.asyncio +async def test_managed_file_id_requirement_is_opt_in(monkeypatch): from litellm.proxy.openai_files_endpoints.common_utils import ( validate_managed_id_requirement, ) monkeypatch.setattr("litellm.require_managed_files", False) - validate_managed_id_requirement(resource_id="file-victim-abc123", resource_kind="file") + await validate_managed_id_requirement( + resource_id="file-victim-abc123", + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + managed_files_obj=None, + ) def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py index c2bdb1d80f2..da5dd1934e4 100644 --- a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py @@ -5,13 +5,15 @@ Every vector-store file route (create, retrieve, content, update, delete) resolv caller-supplied file id through _update_request_data_with_managed_file_id before the provider call, so the guard lives there once and covers all five. -A raw provider file id has no ownership row, so without the guard it is attached to a -vector store or read back under the shared provider credentials with no tenant check. +A raw or forged managed-looking file id has no ownership row, so without the guard it +is attached to a vector store or read back under shared provider credentials. """ import base64 import os import sys +from dataclasses import dataclass +from typing import Literal from unittest.mock import MagicMock, patch import pytest @@ -21,12 +23,35 @@ sys.path.insert(0, os.path.abspath("../../../..")) from fastapi import HTTPException import litellm +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.vector_store_files_endpoints.endpoints import ( _update_request_data_with_managed_file_id, ) from litellm.types.utils import SpecialEnums RAW_FILE_ID = "file-victim-abc123" +CALLER = UserAPIKeyAuth(api_key="sk-test", user_id="attacker-user", team_id="team-b") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: Literal["allow", "deny", "missing"] + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + if self.file_access == "missing": + raise HTTPException(status_code=404, detail=f"File not found: {unified_file_id}") + return self.file_access == "allow" + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return False def _unified_file_id() -> str: @@ -36,24 +61,31 @@ def _unified_file_id() -> str: return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") -def _resolve(file_id: str): - return _update_request_data_with_managed_file_id( +async def _resolve( + file_id: str, + file_access: Literal["allow", "deny", "missing"] = "allow", +): + return await _update_request_data_with_managed_file_id( data={"vector_store_id": "vs-test", "file_id": file_id}, file_id=file_id, request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access=file_access), llm_router=None, ) -def test_raw_file_id_rejected_when_managed_files_required(): +@pytest.mark.asyncio +async def test_raw_file_id_rejected_when_managed_files_required(): with patch.object(litellm, "require_managed_files", True): with pytest.raises(HTTPException) as exc: - _resolve(RAW_FILE_ID) + await _resolve(RAW_FILE_ID) assert exc.value.status_code == 400 -def test_model_encoded_file_id_rejected_when_managed_files_required(): +@pytest.mark.asyncio +async def test_model_encoded_file_id_rejected_when_managed_files_required(): """encode_file_id_with_model output is client-forgeable and carries no ownership row, so it is not a managed file id.""" from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model @@ -62,22 +94,53 @@ def test_model_encoded_file_id_rejected_when_managed_files_required(): with patch.object(litellm, "require_managed_files", True): with pytest.raises(HTTPException) as exc: - _resolve(encoded) + await _resolve(encoded) assert exc.value.status_code == 400 -def test_unified_file_id_allowed_when_managed_files_required(): +@pytest.mark.asyncio +async def test_forged_unified_file_id_rejected_without_ownership_record(): + forged_id = _unified_file_id() + data = {"vector_store_id": "vs-test", "file_id": forged_id} + with patch.object(litellm, "require_managed_files", True): - data, original = _resolve(_unified_file_id()) + with pytest.raises(HTTPException) as exc: + await _update_request_data_with_managed_file_id( + data=data, + file_id=forged_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="missing"), + llm_router=None, + ) + + assert exc.value.status_code == 404 + assert data["file_id"] == forged_id + + +@pytest.mark.asyncio +async def test_other_teams_unified_file_id_rejected(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(_unified_file_id(), file_access="deny") + + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_owned_unified_file_id_allowed_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + data, original = await _resolve(_unified_file_id()) assert original == _unified_file_id() assert data["file_id"] == RAW_FILE_ID -def test_raw_file_id_allowed_when_managed_files_not_required(): +@pytest.mark.asyncio +async def test_raw_file_id_allowed_when_managed_files_not_required(): with patch.object(litellm, "require_managed_files", False): - data, original = _resolve(RAW_FILE_ID) + data, original = await _resolve(RAW_FILE_ID) assert original is None assert data["file_id"] == RAW_FILE_ID From f038be22dbaaf34d4ba695b40c0651d6d576effa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 03:25:35 -0700 Subject: [PATCH 64/74] build(lint): rename make pre-commit to make check with a working-tree fallback --- CLAUDE.md | 4 +- Makefile | 18 ++- scripts/install_git_hooks.sh | 2 +- scripts/pre_commit_lint.sh | 144 ++++++++++++++------- tests/e2e/CONTRIBUTING.md | 2 +- tests/test_litellm/test_pre_commit_lint.py | 90 ++++++++++++- 6 files changed, 201 insertions(+), 59 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 59929143c46..0354e3def53 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -41,7 +41,9 @@ Python max line length is 120, not 88 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing -`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice +Run `make check` (formerly `make pre-commit`, which still works as an alias) before every commit, merge commits included. It runs the CI-gating lint scoped to your staged files, so stage everything you intend to commit first; it warns about changed files you left unstaged and names the checks that were skipped because of them. With nothing staged it instead checks the working tree's diff against the merge base with origin/litellm_internal_staging, which is how you predict the CI lint on an already-committed branch, e.g. right after a merge commit + +`make check` saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in diff --git a/Makefile b/Makefile index 493828571b7..94d8c875af5 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ install-dev install-proxy-dev install-test-deps install-hooks \ - install-helm-unittest check-circular-imports check-import-safety pre-commit \ + install-helm-unittest check-circular-imports check-import-safety check pre-commit \ lint-install lint-fetch-base bootstrap # Default target @@ -22,7 +22,8 @@ help: @echo " make install-test-deps - Install the full local test environment" @echo " make install-helm-unittest - Install helm unittest plugin" @echo " make install-hooks - Install git hooks (Conventional Commits + Branches)" - @echo " make pre-commit - Run CI-equivalent lint on staged files (run before committing)" + @echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged" + @echo " make pre-commit - Legacy alias for make check" @echo " make format - Apply ruff format code formatting" @echo " make format-check - Check ruff format code formatting (matches CI)" @echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)" @@ -236,13 +237,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline # Faster linting for local development (only checks changed code) lint-dev: lint-format-changed check-circular-imports check-import-safety -# Run the gating CI checks against your staged files right before committing. Mirrors +# Run the gating CI checks against your changes. Scopes to staged files when anything +# is staged (warning about changed files left unstaged); with nothing staged it falls +# back to the working tree's diff against the merge base with the base branch, so a +# fresh merge commit or an unstaged working tree still gets checked. Mirrors # test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and -# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage. +# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope. # Not auto-installed as a git hook so it never slows an unrelated human commit. -pre-commit: bootstrap +check: bootstrap ./scripts/pre_commit_lint.sh +pre-commit: + @echo "make pre-commit is a legacy alias; use make check" >&2 + @$(MAKE) check + # Testing targets test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/scripts/install_git_hooks.sh b/scripts/install_git_hooks.sh index 7ea8c3ff2e9..8f4e79e4ddd 100755 --- a/scripts/install_git_hooks.sh +++ b/scripts/install_git_hooks.sh @@ -35,7 +35,7 @@ These hooks enforce Conventional Commits and Conventional Branches. Bypass with --no-verify when you need to (e.g. for emergency hotfixes). The CI-equivalent lint is deliberately not installed as an auto-firing hook -(it can take minutes); run it on demand with 'make pre-commit' before committing. +(it can take minutes); run it on demand with 'make check' before committing. To uninstall: git config --unset core.hooksPath EOF diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index af8335e0e84..1bf7fe17832 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -1,18 +1,25 @@ #!/usr/bin/env bash # -# pre_commit_lint.sh — shift CI lint left. Run it (via `make pre-commit`) right -# before `git commit`; it inspects your staged files and runs only the matching -# gating CI checks, so a clean run means a green CI lint: -# - litellm/ Python staged -> `make lint` (test-linting.yml's lint job) -# - tests/e2e Python staged -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) -# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) -# - dashboard staged -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) -# - proxy/types staged -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) +# pre_commit_lint.sh — shift CI lint left. Run it (via `make check`, formerly +# `make pre-commit`) before `git commit`, or after committing (e.g. a merge +# commit) to predict CI for the branch. It picks the files in scope and runs +# only the matching gating CI checks, so a clean run means a green CI lint: +# - anything staged -> scope is the staged files; changed-but-unstaged files +# whose checks were skipped are called out +# - nothing staged -> scope is the working tree's diff against the merge base +# with origin/litellm_internal_staging, untracked files included +# The per-area checks: +# - litellm/ Python -> `make lint` (test-linting.yml's lint job) +# - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) +# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) +# - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) +# - proxy/types -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) # -# Each block is skipped when no matching files are staged, so unrelated commits stay -# fast. This is intentionally not auto-installed as a git hook (see scripts/install_git_hooks.sh): -# the dashboard and basedpyright passes can take minutes, so it's run on demand rather -# than firing on every human commit. It is hook-compatible if you want that anyway: +# Each block is skipped when no matching files are in scope, so unrelated commits +# stay fast. This is intentionally not auto-installed as a git hook (see +# scripts/install_git_hooks.sh): the dashboard and basedpyright passes can take +# minutes, so it's run on demand rather than firing on every human commit. It is +# hook-compatible if you want that anyway: # `ln -s ../../scripts/pre_commit_lint.sh .git/hooks/pre-commit`. set -eu @@ -20,17 +27,17 @@ set -eu if [ -z "${PRE_COMMIT_LINT_INNER:-}" ]; then log_file=$(git rev-parse --path-format=absolute --git-path pre_commit_lint.log) if : > "$log_file" 2>/dev/null; then - echo "pre-commit: logging full output to $log_file" + echo "check: logging full output to $log_file" PRE_COMMIT_LINT_INNER=1 "$0" "$@" 2>&1 | tee "$log_file" pipe_status=("${PIPESTATUS[@]}") if [ "${pipe_status[1]}" -eq 0 ]; then - echo "pre-commit: full log: $log_file" + echo "check: full log: $log_file" else - echo "pre-commit: WARNING - writing $log_file failed; the log may be incomplete" >&2 + echo "check: WARNING - writing $log_file failed; the log may be incomplete" >&2 fi exit "${pipe_status[0]}" fi - echo "pre-commit: WARNING - cannot write $log_file; output will not be saved" >&2 + echo "check: WARNING - cannot write $log_file; output will not be saved" >&2 PRE_COMMIT_LINT_INNER=1 exec "$0" "$@" fi @@ -38,38 +45,79 @@ repo_root=$(git rev-parse --show-toplevel) cd "$repo_root" staged=$(git diff --cached --name-only --diff-filter=ACMR) -staged_match() { printf '%s\n' "$staged" | grep -E "$1" || true; } +unstaged=$(git diff --name-only) +untracked=$(git ls-files --others --exclude-standard) + +if [ -n "$staged" ]; then + scope=$staged +else + git fetch --quiet origin litellm_internal_staging 2>/dev/null || true + merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || { + echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2 + echo " Fix: git fetch origin litellm_internal_staging" >&2 + exit 1 + } + scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMR "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) + if [ -z "$scope" ]; then + echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" + exit 0 + fi + echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:" + printf '%s\n' "$scope" | sed 's/^/ /' +fi + +scope_match() { printf '%s\n' "$scope" | grep -E "$1" || true; } + +litellm_py_pattern='^litellm/.*\.py$' +e2e_py_pattern='^tests/e2e/.*\.py$' +spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' +ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' +ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' # CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or # scripts-only commit can't turn it red; scope the trigger there to skip the slow # make lint when it couldn't catch anything. -litellm_py_files=$(staged_match '^litellm/.*\.py$') -e2e_py_files=$(staged_match '^tests/e2e/.*\.py$') +litellm_py_files=$(scope_match "$litellm_py_pattern") +e2e_py_files=$(scope_match "$e2e_py_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' || true) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types # (Prisma schema and configs included, not just Python) plus the generator and its # lockfiles, so match that whole trigger set rather than a Python subset. -spec_files=$(staged_match '^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$') +spec_files=$(scope_match "$spec_pattern") # CI's frontend-lint runs prettier over a wider extension set than eslint; keep that # split so this flags exactly what the job would. -ui_prettier_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$') -ui_eslint_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$') +ui_prettier_files=$(scope_match "$ui_prettier_pattern") +ui_eslint_files=$(scope_match "$ui_eslint_pattern") -# CI lints the committed tree, so this script predicts CI for what you have STAGED -# (every trigger above reads `git diff --cached`). The tools it runs, though, read -# the working tree, so unstaged edits to tracked files and untracked files fold -# into the result and a green/red here won't match a commit of just the staged -# changes. There's no safe way to lint the index in place, so surface the gap -# instead of hiding it: stage everything you intend to commit before trusting a -# pass. This only warns; it never blocks or touches your changes. -unstaged=$(git diff --name-only) -untracked=$(git ls-files --others --exclude-standard) -if [ -n "$unstaged" ] || [ -n "$untracked" ]; then - echo "pre-commit: NOTE - unstaged/untracked changes are included in these checks but" >&2 - echo " won't be in a commit of only your staged changes, so this result may differ from" >&2 - echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2 - printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sed 's/^/ /' >&2 +# CI lints the committed tree, so with staged files this script predicts CI for +# what you have STAGED (every trigger above reads `git diff --cached`). The tools +# it runs, though, read the working tree, so unstaged edits to tracked files and +# untracked files fold into the result and a green/red here won't match a commit +# of just the staged changes. There's no safe way to lint the index in place, so +# surface the gap instead of hiding it: stage everything you intend to commit +# before trusting a pass. This only warns; it never blocks or touches your changes. +if [ -n "$staged" ]; then + not_staged=$(printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sort -u) + if [ -n "$not_staged" ]; then + echo "check: NOTE - unstaged/untracked changes are included in these checks but" >&2 + echo " won't be in a commit of only your staged changes, so this result may differ from" >&2 + echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2 + printf '%s\n' "$not_staged" | sed 's/^/ /' >&2 + fi + warn_skipped() { + local check_name=$1 pattern=$2 triggered=$3 + [ -n "$triggered" ] && return 0 + local missed + missed=$(printf '%s\n' "$not_staged" | grep -E "$pattern" || true) + [ -z "$missed" ] && return 0 + echo "check: SKIPPED $check_name because these changed files are not staged:" >&2 + printf '%s\n' "$missed" | sed 's/^/ /' >&2 + } + warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" + warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" + warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_files" + warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi lint_dashboard() { @@ -114,15 +162,15 @@ bootstrap_hint() { python_checks() { local rc=0 - echo "pre-commit: linting Python (make lint)" - make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; rc=1; } + echo "check: linting Python (make lint)" + make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make check." >&2; rc=1; } # `make lint` format-checks files in origin/base...HEAD, which at pre-commit time - # predates the staged change, so format-check the staged litellm files directly to + # predates the staged change, so format-check the scoped litellm files directly to # cover a brand-new commit before it lands. if [ -n "$fmt_files" ]; then - echo "pre-commit: ruff format --check (staged litellm files)" + echo "check: ruff format --check (scoped litellm files)" printf '%s\n' "$fmt_files" | xargs uv run --no-sync ruff format --check --exclude '/enterprise/' \ - || { echo "✗ Unformatted staged files. Fix with: make format, then re-stage." >&2; rc=1; } + || { echo "✗ Unformatted files in scope. Fix with: make format, then re-stage." >&2; rc=1; } fi return $rc } @@ -146,18 +194,18 @@ if [ -n "$litellm_py_files" ]; then fi if [ -n "$e2e_py_files" ] && [ -z "$litellm_py_files" ]; then - echo "pre-commit: type-checking tests/e2e (make lint-e2e-basedpyright)" - make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make pre-commit." >&2; status=1; } + echo "check: type-checking tests/e2e (make lint-e2e-basedpyright)" + make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make check." >&2; status=1; } fi if [ -n "$e2e_py_files" ]; then - echo "pre-commit: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" + echo "check: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" uv run --no-sync python tests/code_coverage_tests/check_e2e_no_raw_requests.py \ - || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make pre-commit." >&2; status=1; } + || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; } fi dashboard_checks() { - echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)" + echo "check: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2 bootstrap_hint @@ -176,7 +224,7 @@ fi genapi_checks() { local status=0 - echo "pre-commit: checking dashboard API types are in sync (npm run gen:api)" + echo "check: checking dashboard API types are in sync (npm run gen:api)" # gen-api-types.mjs imports litellm.proxy.proxy_server, which needs the proxy deps # and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs # prisma generate before gen:api, so mirror that here or a stale client can mask @@ -194,7 +242,7 @@ genapi_checks() { status=1 elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then if ! git diff --quiet -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then - echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make pre-commit only if other checks failed too." >&2 + echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make check only if other checks failed too." >&2 status=1 fi else diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 5999e5772d1..dc69bd42171 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -138,7 +138,7 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover Before you push -1. Run `make lint-e2e-basedpyright` (or `make pre-commit` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py` +1. Run `make lint-e2e-basedpyright` (or `make check` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py` 2. Add the models your test needs to the config your local proxy loads diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 274d3c517f6..12b8d338b49 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -129,6 +129,90 @@ def _run(repo: Path, bin_dir: Path, extra_env: dict[str, str]) -> subprocess.Com ) +def _commit_all(repo: Path, message: str) -> None: + subprocess.run(["git", "add", "."], cwd=repo, check=True) + subprocess.run( + ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", message], + cwd=repo, + check=True, + ) + + +def _set_base_ref(repo: Path) -> None: + subprocess.run( + ["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"], + cwd=repo, + check=True, + ) + + +def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").write_text("x = 2\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_checks_committed_branch_changes(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").write_text("x = 2\n") + _commit_all(repo, "branch change") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_includes_untracked_files_in_scope(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "brand_new.py").write_text("z = 3\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "litellm/brand_new.py" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing to check" in proc.stdout + assert "linting Python" not in proc.stdout + + +def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 1 + assert "cannot resolve the merge base" in proc.stdout + assert "git fetch origin litellm_internal_staging" in proc.stdout + + +def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + (repo / "notes.md").write_text("hi\n") + subprocess.run(["git", "add", "notes.md"], cwd=repo, check=True) + (repo / "litellm" / "foo.py").write_text("x = 4\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "SKIPPED Python lint (make lint)" in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" not in proc.stdout + + def test_python_dashboard_and_gen_api_blocks_run_concurrently_with_grouped_output(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) barrier_dir = tmp_path / "barrier" @@ -159,8 +243,8 @@ def test_full_output_is_saved_to_a_log_file_in_the_git_dir(tmp_path: Path) -> No assert "linting dashboard" in log assert "API types" in log assert "unstaged/untracked changes" in log - assert f"pre-commit: full log: {log_file}" in proc.stdout - assert "pre-commit: full log:" not in log + assert f"check: full log: {log_file}" in proc.stdout + assert "check: full log:" not in log def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Path) -> None: @@ -170,7 +254,7 @@ def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Pa assert proc.returncode == 0, proc.stdout + proc.stderr assert "linting Python" in proc.stdout assert "output will not be saved" in proc.stderr - assert "pre-commit: full log:" not in proc.stdout + assert "check: full log:" not in proc.stdout failing = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"}) assert failing.returncode == 1 From 20eb7bb43718958bfac0e06225ead0b5ddb1d5b7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 08:14:29 -0700 Subject: [PATCH 65/74] chore(typing): clear 1.4k basedpyright Any errors across 21 hotspot files Typing-only pass over the 21 files with the highest reportAny and reportExplicitAny density among self-contained modules: management endpoints, guardrails, streaming internals, response transformations, MCP server, enterprise managed files, and vector store management. Whole-tree basedpyright drops from 148,648 to 146,984 errors (-1,664), with reportAny -1,111 and reportExplicitAny -296. No rule increased repo-wide and no file regressed on any rule. No cast(), type: ignore, noqa, suppression comments, or new Any annotations anywhere in the diff, and no runtime behavior changes. Budgets ratcheted by make lint-budget-update: basedpyright -1,663 across 48 rules, ruff-strict -86, type-discipline -110. --- basedpyright-code-budget.json | 28 +- .../proxy/hooks/managed_files.py | 524 +++++++----------- .../pydantic_ai_agents/transformation.py | 265 +++++---- .../websearch_interception/handler.py | 12 +- .../litellm_core_utils/realtime_streaming.py | 55 +- .../streaming_chunk_builder_utils.py | 97 +++- .../adapters/handler.py | 84 +-- litellm/llms/azure/assistants.py | 36 +- .../mcp_server/rest_endpoints.py | 2 +- .../proxy/_experimental/mcp_server/server.py | 91 +-- .../proxy/guardrails/guardrail_endpoints.py | 173 ++++-- .../cisco_ai_defense/cisco_ai_defense.py | 100 ++-- .../unified_guardrail/unified_guardrail.py | 136 +++-- litellm/proxy/hooks/litellm_skills/main.py | 60 +- .../key_management_endpoints.py | 77 +-- .../model_management_endpoints.py | 160 ++++-- .../management_endpoints/team_endpoints.py | 224 +++++--- litellm/proxy/management_endpoints/ui_sso.py | 187 +++++-- .../proxy_setting_endpoints.py | 123 +++- .../management_endpoints.py | 72 ++- .../transformation.py | 184 +++--- litellm/responses/streaming_iterator.py | 51 +- ruff-strict-budget.json | 10 +- type-discipline-budget.json | 8 +- 24 files changed, 1644 insertions(+), 1115 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 32b8eb3d4d0..0385f7a96e7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 28842 + "limit": 27731 }, "reportArgumentType": { - "limit": 2634 + "limit": 2626 }, "reportAssignmentType": { "limit": 329 @@ -12,7 +12,7 @@ "limit": 514 }, "reportCallIssue": { - "limit": 117 + "limit": 116 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9103 + "limit": 8807 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5843 + "limit": 5835 }, "reportMissingTypeArgument": { - "limit": 15816 + "limit": 15790 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1078 + "limit": 1077 }, "reportOptionalOperand": { "limit": 0 @@ -90,28 +90,28 @@ "limit": 8 }, "reportReturnType": { - "limit": 218 + "limit": 217 }, "reportTypedDictNotRequiredAccess": { - "limit": 27 + "limit": 26 }, "reportUndefinedVariable": { "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45098 + "limit": 45063 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39826 + "limit": 39773 }, "reportUnknownParameterType": { - "limit": 20237 + "limit": 20207 }, "reportUnknownVariableType": { - "limit": 31371 + "limit": 31281 }, "reportUnnecessaryCast": { "limit": 122 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 864 + "limit": 862 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d8318962633..f0914240f79 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,9 +3,21 @@ import base64 import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Final, + List, + Literal, + Optional, + Protocol, + TypedDict, + Union, + cast, +) from uuid import NAMESPACE_URL, uuid5 from fastapi import HTTPException @@ -98,33 +110,76 @@ def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMB try: batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}" - ) + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}") return None batch_obj.id = row.unified_object_id return batch_obj -def _parse_managed_file_object( - raw_file_object: object, unified_file_id: str -) -> Optional[OpenAIFileObject]: +def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]: if raw_file_object is None: return None try: return OpenAIFileObject.model_validate(raw_file_object) except Exception as e: - verbose_logger.warning( - f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}" - ) + verbose_logger.warning(f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}") return None +class _ManagedFileRow(Protocol): + unified_file_id: str + file_object: OpenAIFileObject + storage_backend: Optional[str] + storage_url: Optional[str] + created_by: Optional[str] + team_id: Optional[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _ManagedFileTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[_ManagedFileRow]: ... + + async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ... + + async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ... + + +class _ManagedObjectTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ... + + async def find_many( + self, + where: Mapping[str, object], + take: int, + order: Union[Mapping[str, str], Sequence[Mapping[str, str]]], + cursor: Mapping[str, str] = ..., + skip: int = ..., + ) -> "Sequence[PrismaManagedObjectRow]": ... + + async def upsert( + self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> "PrismaManagedObjectRow": ... + + +class _CursorPageArgs(TypedDict, total=False): + cursor: Mapping[str, str] + skip: int + + +def _managed_file_table(prisma_client: PrismaClient) -> _ManagedFileTableActions: + return prisma_client.db.litellm_managedfiletable + + +def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableActions: + return prisma_client.db.litellm_managedobjecttable + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes - def __init__( - self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient - ): + def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client @@ -143,9 +198,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed File object with id={file_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache") if file_object is not None: litellm_managed_file_object = LiteLLM_ManagedFileTable( unified_file_id=file_id, @@ -196,13 +249,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"storage_url={db_data.get('storage_url')}" ) - result = await self.prisma_client.db.litellm_managedfiletable.upsert( + result = await _managed_file_table(self.prisma_client).upsert( where={"unified_file_id": file_id}, data={"create": db_data, "update": update_data}, ) - verbose_logger.debug( - f"LiteLLM Managed File object with id={file_id} stored in db: {result}" - ) + verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") async def store_unified_object_id( self, @@ -213,9 +264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -228,7 +277,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedobjecttable.upsert( + await _managed_object_table(self.prisma_client).upsert( where={"unified_object_id": unified_object_id}, data={ "create": { @@ -265,9 +314,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return LiteLLM_ManagedFileTable.model_validate(result) ## CHECK DB - db_object = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_object = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if db_object: return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump()) @@ -277,9 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, file_id: str, litellm_parent_otel_span: Optional[Span] = None ) -> OpenAIFileObject: ## get old value - initial_value = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if initial_value is None: raise Exception(f"LiteLLM Managed File object with id={file_id} not found") ## delete old value @@ -288,15 +333,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): value=None, litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedfiletable.delete( - where={"unified_file_id": file_id} - ) + await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id}) return initial_value.file_object - async def can_user_call_unified_file_id( - self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_file = await self.prisma_client.db.litellm_managedfiletable.find_first( + async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_file = await _managed_file_table(self.prisma_client).find_first( where={"unified_file_id": unified_file_id} ) @@ -311,13 +352,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"File not found: {unified_file_id}", ) - async def can_user_call_unified_object_id( - self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_object = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={"unified_object_id": unified_object_id} - ) + async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_object = await _managed_object_table(self.prisma_client).find_first( + where={"unified_object_id": unified_object_id} ) if managed_object: @@ -339,34 +376,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): provider: Optional[str] = None, target_model_names: Optional[str] = None, llm_router: Optional[Router] = None, - ) -> Dict[str, Any]: + ) -> Dict[str, object]: # Provider filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the provider information # To support provider filtering, we would need to store the provider information in the encoded object ids if provider: - raise Exception( - "Filtering by 'provider' is not supported when using managed batches." - ) + raise Exception("Filtering by 'provider' is not supported when using managed batches.") # Model name filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the model name # A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids. if target_model_names: - raise Exception( - "Filtering by 'target_model_names' is not supported when using managed batches." - ) + raise Exception("Filtering by 'target_model_names' is not supported when using managed batches.") owner_filter = build_owner_filter(user_api_key_dict) if owner_filter is None: return build_list_page([]) - where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} + where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter} if after: - cursor_row = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={**where_clause, "unified_object_id": after} - ) + cursor_row = await _managed_object_table(self.prisma_client).find_first( + where={**where_clause, "unified_object_id": after} ) if cursor_row is None: raise HTTPException( @@ -375,11 +406,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) page_size: Final = min(limit or 20, 100) - cursor_args: Dict[str, Any] = ( - {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - ) + cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + batches = await _managed_object_table(self.prisma_client).find_many( where=where_clause, take=page_size + 1, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], @@ -389,9 +418,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size parsed_rows: Final = tuple( - (row, batch_obj) - for row in batches[:page_size] - if (batch_obj := _parse_managed_batch_row(row)) is not None + (row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None ) unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( raw_file_ids=frozenset( @@ -432,14 +459,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): verbose_proxy_logger=verbose_logger, user_api_key_dict=user_api_key_dict, db_batch_object=row, - unified_batch_id=_is_base64_encoded_unified_file_id( - row.unified_object_id - ), + unified_batch_id=_is_base64_encoded_unified_file_id(row.unified_object_id), ) except Exception as e: - verbose_logger.warning( - f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}") return None return batch_obj @@ -458,7 +481,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return [] - file_ids = await self.prisma_client.db.litellm_managedfiletable.find_many( + file_ids = await _managed_file_table(self.prisma_client).find_many( where={ **owner_filter, "flat_model_file_ids": {"hasSome": model_object_ids}, @@ -467,27 +490,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return [ parsed_file_object.model_copy(update={"id": row.unified_file_id}) for row in file_ids - if ( - parsed_file_object := _parse_managed_file_object( - row.file_object, row.unified_file_id - ) - ) - is not None + if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None ] - async def check_managed_file_id_access( - self, data: Dict, user_api_key_dict: UserAPIKeyAuth - ) -> bool: + async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and retrieve_file_id: - if await self.can_user_call_unified_file_id( - retrieve_file_id, user_api_key_dict - ): + if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict): return True else: raise HTTPException( @@ -496,9 +506,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return False - async def check_file_ids_access( - self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None: """ Check if the user has access to a list of file IDs. Only checks managed (unified) file IDs. @@ -513,9 +521,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for file_id in file_ids: is_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_unified_file_id: - if not await self.can_user_call_unified_file_id( - file_id, user_api_key_dict - ): + if not await self.can_user_call_unified_file_id(file_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", @@ -543,10 +549,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ### HANDLE TRANSFORMATIONS ### # Check both completion and acompletion call types - is_completion_call = ( - call_type == CallTypes.completion.value - or call_type == CallTypes.acompletion.value - ) + is_completion_call = call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value if is_completion_call: messages = data.get("messages") @@ -559,9 +562,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check if any files are stored in storage backends and need base64 conversion # This is needed for Vertex AI/Gemini which requires base64 content - is_vertex_ai = model and ( - "vertex_ai" in model or "gemini" in model.lower() - ) + is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower()) if is_vertex_ai: await self._convert_storage_files_to_base64( messages=messages, @@ -573,10 +574,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids, user_api_key_dict.parent_otel_span ) data["model_file_id_mapping"] = model_file_id_mapping - elif ( - call_type == CallTypes.aresponses.value - or call_type == CallTypes.responses.value - ): + elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value: # Handle managed files in responses API input and tools file_ids = [] @@ -603,23 +601,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if tools: unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools) if unified_vs_ids: - await self.check_vector_store_ids_access( - unified_vs_ids, user_api_key_dict - ) + await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict) elif call_type == CallTypes.afile_content.value: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: data["model"] = model_id - data["file_id"] = self.get_output_file_id_from_unified_file_id( - potential_file_id - ) + data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -636,10 +626,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ): accessor_key: Optional[str] = None retrieve_object_id: Optional[str] = None - if ( - call_type == CallTypes.aretrieve_batch.value - or call_type == CallTypes.acancel_batch.value - ): + if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value: accessor_key = "batch_id" elif ( call_type == CallTypes.acancel_fine_tuning_job.value @@ -651,32 +638,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): retrieve_object_id = cast(Optional[str], data.get(accessor_key)) potential_llm_object_id = ( - _is_base64_encoded_unified_file_id(retrieve_object_id) - if retrieve_object_id - else False + _is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False ) if potential_llm_object_id and retrieve_object_id: ## VALIDATE USER HAS ACCESS TO THE OBJECT ## - if not await self.can_user_call_unified_object_id( - retrieve_object_id, user_api_key_dict - ): + if not await self.can_user_call_unified_object_id(retrieve_object_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the object {retrieve_object_id}", ) ## for managed batch id - get the model id - potential_model_id = get_model_id_from_unified_batch_id( - potential_llm_object_id - ) + potential_model_id = get_model_id_from_unified_batch_id(potential_llm_object_id) if potential_model_id is None: raise Exception( f"LiteLLM Managed {accessor_key} with id={retrieve_object_id} is invalid - does not contain encoded model_id." ) data["model"] = potential_model_id - data[accessor_key] = get_batch_id_from_unified_batch_id( - potential_llm_object_id - ) + data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id) elif call_type == CallTypes.acreate_fine_tuning_job.value: input_file_id = cast(Optional[str], data.get("training_file")) if input_file_id: @@ -732,24 +711,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if accessor_key: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) - model_file_id_mapping = cast( - Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping") - ) + model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) if model_id is None: model_id = cast( Optional[str], - kwargs.get("litellm_metadata", {}) - .get("model_info", {}) - .get("id", None), + kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), ) mapped_file_id: Optional[str] = None if input_file_id and model_file_id_mapping and model_id: - mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get( - model_id, None - ) + mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None) if mapped_file_id: kwargs[accessor_key] = mapped_file_id @@ -775,9 +748,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input( - self, input: Union[str, List[Dict[str, Any]]] - ) -> List[str]: + def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]: """ Gets file ids from responses API input. @@ -809,19 +780,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): content = item.get("content") if isinstance(content, list): for content_item in content: - if ( - isinstance(content_item, dict) - and content_item.get("type") == "input_file" - ): + if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") if file_id: file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Gets file ids from responses API tools parameter. @@ -854,9 +820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_ids - def get_vector_store_ids_from_file_search_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Extract unified vector_store_ids from file_search tools. @@ -949,9 +913,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping( - self, file_ids: List[str], litellm_parent_otel_span: Span - ) -> dict: + async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id @@ -981,9 +943,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get all cache keys matching the pattern file_id:* for file_id in litellm_managed_file_ids: # Search for any cache key starting with this file_id - unified_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + unified_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) if unified_file_object: file_id_mapping[file_id] = unified_file_object.model_mappings @@ -1001,9 +961,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): raise Exception("LLM Router not initialized. Ensure models added to proxy.") responses = [] for model in target_model_names_list: - individual_response = await llm_router.acreate_file( - model=model, **_create_file_request - ) + individual_response = await llm_router.acreate_file(model=model, **_create_file_request) responses.append(individual_response) return responses @@ -1034,9 +992,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str] = {} for file_object in responses: - model_file_id_mapping = file_object._hidden_params.get( - "model_file_id_mapping" - ) + model_file_id_mapping = file_object._hidden_params.get("model_file_id_mapping") if model_file_id_mapping and isinstance(model_file_id_mapping, dict): model_mappings.update(model_file_id_mapping) @@ -1051,17 +1007,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Emit Prometheus metrics for managed file creation prom_logger = self._get_prometheus_logger() if prom_logger: - first_model = ( - target_model_names_list[0] if target_model_names_list else None - ) + first_model = target_model_names_list[0] if target_model_names_list else None first_provider = "" if responses: - first_provider = ( - getattr(responses[0], "_hidden_params", {}).get( - "custom_llm_provider" - ) - or "" - ) + first_provider = getattr(responses[0], "_hidden_params", {}).get("custom_llm_provider") or "" prom_logger.record_managed_file_created( model=first_model or "", api_provider=first_provider, @@ -1104,9 +1053,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # Convert to URL-safe base64 and strip padding - base64_unified_file_id = ( - base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") - ) + base64_unified_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") ## CREATE RESPONSE OBJECT @@ -1123,46 +1070,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response - def get_unified_generic_response_id( - self, model_id: str, generic_response_id: str - ) -> str: - unified_generic_response_id = ( - SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( - model_id, generic_response_id - ) - ) - return ( - base64.urlsafe_b64encode(unified_generic_response_id.encode()) - .decode() - .rstrip("=") + def get_unified_generic_response_id(self, model_id: str, generic_response_id: str) -> str: + unified_generic_response_id = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( + model_id, generic_response_id ) + return base64.urlsafe_b64encode(unified_generic_response_id.encode()).decode().rstrip("=") def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: - unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( - model_id, batch_id - ) + unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id( - self, output_file_id: str, model_id: str, model_name: Optional[str] - ) -> str: - deterministic_uuid: Final = uuid5( - uuid5(NAMESPACE_URL, model_id), output_file_id - ) - unified_output_file_id = ( - SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", - str(deterministic_uuid), - model_name or "", - output_file_id, - model_id, - ) - ) - return ( - base64.urlsafe_b64encode(unified_output_file_id.encode()) - .decode() - .rstrip("=") + def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) + unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", + str(deterministic_uuid), + model_name or "", + output_file_id, + model_id, ) + return base64.urlsafe_b64encode(unified_output_file_id.encode()).decode().rstrip("=") def get_model_id_from_unified_file_id(self, file_id: str) -> str: return file_id.split("llm_output_file_model_id,")[1].split(";")[0] @@ -1170,59 +1097,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): def get_output_file_id_from_unified_file_id(self, file_id: str) -> str: marker = "llm_output_file_id," if marker not in file_id: - raise ValueError( - f"Unified id does not contain {marker!r}: {file_id[:80]!r}" - ) + raise ValueError(f"Unified id does not contain {marker!r}: {file_id[:80]!r}") return file_id.split(marker, 1)[1].split(";")[0] async def async_post_call_success_hook( self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes - ) -> Any: + ) -> LLMResponseTypes: if isinstance(response, LiteLLMBatch): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id - unified_batch_id = response._hidden_params.get( - "unified_batch_id" - ) # managed batch id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id + unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) resolved_model_name = resolve_managed_output_file_model_name( - unified_input_file_id=unified_file_id - if isinstance(unified_file_id, str) - else response.input_file_id, + unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id, fallback_model_name=model_name, ) original_response_id = response.id if (unified_batch_id or unified_file_id) and model_id: - response.id = self.get_unified_batch_id( - batch_id=response.id, model_id=model_id - ) + response.id = self.get_unified_batch_id(batch_id=response.id, model_id=model_id) # Handle both output_file_id and error_file_id for file_attr in ["output_file_id", "error_file_id"]: file_id_value = getattr(response, file_attr, None) if file_id_value and model_id: - decoded_output_file_id = _is_base64_encoded_unified_file_id( - file_id_value - ) - if ( - decoded_output_file_id - and "llm_output_file_id," in decoded_output_file_id - ): - provider_file_id = ( - self.get_output_file_id_from_unified_file_id( - decoded_output_file_id - ) - ) + decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value) + if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id: + provider_file_id = self.get_output_file_id_from_unified_file_id(decoded_output_file_id) unified_file_id = file_id_value elif decoded_output_file_id: verbose_logger.warning( - f"Skipping {file_attr}={file_id_value!r}: " - "unified id is not a managed file output id" + f"Skipping {file_attr}={file_id_value!r}: unified id is not a managed file output id" ) continue else: @@ -1241,23 +1148,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Import module and use getattr for better testability with mocks import litellm.proxy.proxy_server as proxy_server_module - _llm_router = getattr( - proxy_server_module, "llm_router", None - ) + _llm_router = getattr(proxy_server_module, "llm_router", None) if _llm_router is not None and model_id: - _creds = ( - _llm_router.get_deployment_credentials_with_provider( - model_id - ) - or {} - ) + _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {} file_object = await litellm.afile_retrieve( file_id=provider_file_id, **_creds, ) else: file_object = await litellm.afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type] + custom_llm_provider=model_name.split("/")[0] + if model_name and "/" in model_name + else "openai", # type: ignore[arg-type] file_id=provider_file_id, ) verbose_logger.debug( @@ -1311,9 +1213,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_finetuning_job_id = response._hidden_params.get( "unified_finetuning_job_id" ) # managed finetuning job id @@ -1321,9 +1221,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name = cast(Optional[str], response._hidden_params.get("model_name")) original_response_id = response.id if (unified_file_id or unified_finetuning_job_id) and model_id: - response.id = self.get_unified_generic_response_id( - model_id=model_id, generic_response_id=response.id - ) + response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id) await self.store_unified_object_id( unified_object_id=response.id, file_object=response, @@ -1338,9 +1236,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ ## check if file object if hasattr(response, "data") and isinstance(response.data, list): - if all( - isinstance(file_object, FileObject) for file_object in response.data - ): + if all(isinstance(file_object, FileObject) for file_object in response.data): ## Get all file id's ## Check which file id's were created by the user ## Filter the response to only include the files created by the user @@ -1349,9 +1245,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object.id for file_object in cast(List[FileObject], response.data) # type: ignore ] - user_created_file_ids = await self.get_user_created_file_ids( - user_api_key_dict, file_ids - ) + user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore return response @@ -1359,11 +1253,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response async def afile_retrieve( - self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None + self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None ) -> OpenAIFileObject: - stored_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) # Case 1 : This is not a managed file if not stored_file_object: @@ -1386,21 +1278,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) try: - model_id, model_file_id = next( - iter(stored_file_object.model_mappings.items()) - ) - credentials = ( - llm_router.get_deployment_credentials_with_provider(model_id) or {} - ) - response = await litellm.afile_retrieve( - file_id=model_file_id, **credentials - ) + model_id, model_file_id = next(iter(stored_file_object.model_mappings.items())) + credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {} + response = await litellm.afile_retrieve(file_id=model_file_id, **credentials) response.id = file_id # Replace with unified ID return response except Exception as e: - raise Exception( - f"Failed to retrieve file {file_id} from provider: {str(e)}" - ) from e + raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e async def afile_list( self, @@ -1437,12 +1321,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return False except Exception as e: - verbose_logger.warning( - f"Error checking batch polling configuration: {e}. Assuming disabled." - ) + verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.") return False - async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]: + async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]: """ Find batches that reference this file and still need cost tracking. Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost. @@ -1458,9 +1340,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get model-specific file IDs for this unified file ID if it's a managed file try: - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span=None - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span=None) if model_file_id_mapping and file_id in model_file_id_mapping: # Add all provider file IDs for this unified file @@ -1468,8 +1348,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids_to_check.extend(provider_file_ids) except Exception as e: verbose_logger.debug( - f"Could not get model file ID mapping for {file_id}: {e}. " - f"Will only check unified file ID." + f"Could not get model file ID mapping for {file_id}: {e}. Will only check unified file ID." ) MAX_MATCHES_TO_RETURN = 10 @@ -1487,11 +1366,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for batch in batches: try: # Parse the batch file_object to check for file references - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) + batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id @@ -1500,9 +1375,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): output_file_id = batch_data.get("output_file_id") error_file_id = batch_data.get("error_file_id") - referenced_file_ids = [ - fid for fid in [input_file_id, output_file_id, error_file_id] if fid - ] + referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid] # Check if any referenced file ID matches the file we're trying to delete if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids): @@ -1514,9 +1387,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) except Exception as e: - verbose_logger.warning( - f"Error parsing batch object {batch.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Error parsing batch object {batch.unified_object_id}: {e}") continue return referencing_batches @@ -1545,21 +1416,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if referencing_batches: # File is referenced by non-terminal batches and polling is enabled - MAX_BATCHES_IN_ERROR = ( - 5 # Limit batches shown in error message for readability - ) + MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability # Show up to MAX_BATCHES_IN_ERROR in the error message batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR] - batch_statuses = [ - f"{b['batch_id']}: {b['status']}" for b in batches_to_show - ] + batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show] # Determine the count message count_message = f"{len(referencing_batches)}" - if ( - len(referencing_batches) >= 10 - ): # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file + if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file count_message = "10+" error_message = ( @@ -1600,23 +1465,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await self._check_file_deletion_allowed(file_id) # file_id = convert_b64_uid_to_unified_uid(file_id) - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) delete_response = None specific_model_file_id_mapping = model_file_id_mapping.get(file_id) if specific_model_file_id_mapping: # Remove conflicting keys from data to avoid duplicate keyword arguments - filtered_data = { - k: v for k, v in data.items() if k not in ("model", "file_id") - } + filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore - stored_file_object = await self.delete_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) # Record successful deletion metric only on actual success if stored_file_object or delete_response: @@ -1643,9 +1502,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): Get the content of a file from first model that has it """ model_file_id_mapping = data.pop("model_file_id_mapping", None) - model_file_id_mapping = ( - model_file_id_mapping - or await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) + model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping( + [file_id], litellm_parent_otel_span ) specific_model_file_id_mapping = model_file_id_mapping.get(file_id) @@ -1658,13 +1516,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # against the deployment's configured bucket, which they only # trust from this immutable server-side snapshot, never from # request params. - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_id - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) if credentials is not None: - data["_litellm_internal_model_credentials"] = cast( - Dict, MappingProxyType(dict(credentials)) - ) + data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) else: data.pop("_litellm_internal_model_credentials", None) return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore @@ -1699,9 +1553,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check database for storage backend info # IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version) # So we query with the original file_id (which is base64 encoded) - db_file = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_file = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if not db_file or not db_file.storage_backend or not db_file.storage_url: continue @@ -1727,22 +1579,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_content = await storage_backend.download_file(storage_url) # Determine content type from file object - content_type = self._get_content_type_from_file_object( - db_file.file_object - ) + content_type = self._get_content_type_from_file_object(db_file.file_object) # Convert to base64 base64_data = base64.b64encode(file_content).decode("utf-8") base64_data_uri = f"data:{content_type};base64,{base64_data}" # Update messages to use base64 instead of file_id - self._update_messages_with_base64_data( - messages, file_id, base64_data_uri, content_type - ) + self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type) except Exception as e: - verbose_logger.exception( - f"Error converting file {file_id} from storage backend to base64: {str(e)}" - ) + verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}") # Continue with other files even if one fails continue diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..024e8c179c2 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any, Final, Protocol, cast, runtime_checkable from uuid import uuid4 +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ... + + +@runtime_checkable +class _SupportsPydanticDict(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + class PydanticAITransformation: """ @@ -28,7 +45,7 @@ class PydanticAITransformation: """ @staticmethod - def _remove_none_values(obj: Any) -> Any: + def _remove_none_values(obj: object) -> object: """ Recursively remove None values from a dict/list structure. @@ -42,14 +59,18 @@ class PydanticAITransformation: Cleaned object with None values removed """ if isinstance(obj, dict): - return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None} + typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj) + return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None} elif isinstance(obj, list): - return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None] + typed_list: Final = _LIST_ADAPTER.validate_python(obj) + return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None] else: return obj @staticmethod - def _params_to_dict(params: Any) -> dict[str, Any]: + def _params_to_dict( + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", + ) -> Mapping[str, object]: """ Convert params to a dict, handling Pydantic models. @@ -59,10 +80,10 @@ class PydanticAITransformation: Returns: Dict representation of params """ - if hasattr(params, "model_dump"): + if isinstance(params, _SupportsModelDump): # Pydantic v2 model return params.model_dump(mode="python", exclude_none=True) - elif hasattr(params, "dict"): + elif isinstance(params, _SupportsPydanticDict): # Pydantic v1 model return params.dict(exclude_none=True) elif isinstance(params, dict): @@ -75,12 +96,12 @@ class PydanticAITransformation: async def _poll_for_completion( client: AsyncHTTPHandler, endpoint: str, - task_id: str, + task_id: object, request_id: str, max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,10 +133,10 @@ class PydanticAITransformation: }, ) response.raise_for_status() - poll_data = response.json() + poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) - result = poll_data.get("result", {}) - status = result.get("status", {}) + result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {})) + status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state = status.get("state", "") verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state) @@ -133,10 +154,10 @@ class PydanticAITransformation: async def _send_and_poll_raw( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -153,14 +174,16 @@ class PydanticAITransformation: Raw Pydantic AI task response (with history/artifacts) """ # Convert params to dict if it's a Pydantic model - params_dict = PydanticAITransformation._params_to_dict(params) - # Remove None values - FastA2A doesn't accept null for optional fields - params_dict = PydanticAITransformation._remove_none_values(params_dict) + params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python( + PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params)) + ) # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI if "message" in params_dict: - params_dict["message"]["kind"] = "message" + message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"]) + message_value["kind"] = "message" + params_dict["message"] = message_value # Build A2A JSON-RPC request using message/send method for FastA2A compatibility a2a_request: Final = { @@ -189,11 +212,11 @@ class PydanticAITransformation: }, ) response.raise_for_status() - response_data = response.json() + response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) # Check if task is already completed - result: Final = response_data.get("result", {}) - status: Final = result.get("status", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state: Final = status.get("state", "") if state != "completed": @@ -217,10 +240,10 @@ class PydanticAITransformation: async def send_non_streaming_request( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -253,10 +276,10 @@ class PydanticAITransformation: async def send_and_get_raw_response( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -282,9 +305,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -328,7 +351,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]: """ Extract response text from completed task response. @@ -342,52 +365,53 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result: Final = response_data.get("result", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) # Try to extract from artifacts first (preferred for results) artifacts: Final = result.get("artifacts", []) if artifacts: - for artifact in artifacts: - parts = artifact.get("parts", []) + for artifact in _LIST_ADAPTER.validate_python(artifacts): + parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", [])) for part in parts: - if part.get("kind") == "text": - text = part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + text = part_dict.get("text", "") if text: return text, str(uuid4()), parts # Fall back to history - get the last agent message - history: Final = result.get("history", []) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) for msg in reversed(history): - if msg.get("role") == "agent": - parts = msg.get("parts", []) - message_id = msg.get("messageId", str(uuid4())) + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent": + parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", [])) + message_id = msg_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) if full_text: return full_text, message_id, parts # Fall back to message field (original format) message: Final = result.get("message", {}) if message: - parts = message.get("parts", []) - message_id = message.get("messageId", str(uuid4())) + message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", [])) + message_id = message_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) return full_text, message_id, parts return "", str(uuid4()), [] @staticmethod async def fake_streaming_from_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -410,12 +434,12 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result: Final = response_data.get("result", {}) - history: Final = result.get("history", []) - input_message = {} + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) + input_message = _STR_KEY_DICT_ADAPTER.validate_python({}) for msg in history: - if msg.get("role") == "user": - input_message = msg + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user": + input_message = msg_dict break # Generate IDs for streaming events @@ -426,45 +450,49 @@ class PydanticAITransformation: # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "history": [ - { - "contextId": context_id, - "kind": "message", - "messageId": input_message_id, - "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), - "role": "user", - "taskId": task_id, - } - ], - "id": task_id, - "kind": "task", - "status": { - "state": "submitted", + task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "history": [ + { + "contextId": context_id, + "kind": "message", + "messageId": input_message_id, + "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), + "role": "user", + "taskId": task_id, + } + ], + "id": task_id, + "kind": "task", + "status": { + "state": "submitted", + }, }, - }, - } + } + ) yield task_event # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": False, - "kind": "status-update", - "status": { - "state": "working", + working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": False, + "kind": "status-update", + "status": { + "state": "working", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield working_event # Small delay to simulate processing @@ -473,29 +501,32 @@ class PydanticAITransformation: # 3. Emit artifact update chunks (kind: "artifact-update") # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event if full_text: + full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text) # Split text into chunks - for i in range(0, len(full_text), chunk_size): - chunk_text = full_text[i : i + chunk_size] - is_last_chunk = (i + chunk_size) >= len(full_text) + for i in range(0, len(full_text_str), chunk_size): + chunk_text = full_text_str[i : i + chunk_size] + is_last_chunk = (i + chunk_size) >= len(full_text_str) - artifact_event = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "kind": "artifact-update", - "taskId": task_id, - "artifact": { - "artifactId": artifact_id, - "parts": [ - { - "kind": "text", - "text": chunk_text, - } - ], + artifact_event = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "kind": "artifact-update", + "taskId": task_id, + "artifact": { + "artifactId": artifact_id, + "parts": [ + { + "kind": "text", + "text": chunk_text, + } + ], + }, }, - }, - } + } + ) yield artifact_event # Add delay between chunks (except for last chunk) @@ -503,19 +534,21 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": True, - "kind": "status-update", - "status": { - "state": "completed", + completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": True, + "kind": "status-update", + "status": { + "state": "completed", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield completed_event verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 9748db2dcd2..7abbf0c96e5 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -42,7 +42,7 @@ from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import AgenticLoopParams, CallTypes, LlmProviders from litellm.utils import ProviderConfigManager if TYPE_CHECKING: @@ -265,7 +265,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None # Check if request has tools with native web_search - tools: Final = kwargs.get("tools") + tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools") if not tools: return None @@ -314,7 +314,9 @@ class WebSearchInterceptionLogger(CustomLogger): return kwargs - def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None: + def _convert_responses_tools( + self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]] + ) -> dict[str, object] | None: """Convert Responses API web search tools to the LiteLLM standard function tool.""" if not any(is_web_search_tool_responses(tool) for tool in tools): return None @@ -379,7 +381,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _tool_name(tool: dict[str, Any]) -> str | None: + def _tool_name(tool: Mapping[str, object]) -> object: """Effective tool name, handling OpenAI ``function`` wrapper shape.""" fn: Final = tool.get("function") if tool.get("type") == "function" and isinstance(fn, dict): @@ -1271,7 +1273,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = agentic_params.get("model", model) verbose_logger.debug( "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]", diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 858d10df53b..d68bdc4a250 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm @@ -20,11 +21,24 @@ from .litellm_logging import Logging as LiteLLMLogging if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from litellm.types.guardrails import GuardrailEventHooks + CLIENT_CONNECTION_CLASS = ClientConnection else: CLIENT_CONNECTION_CLASS = Any +class _ClientWebSocketExceptions(Protocol): + ConnectionClosed: type[Exception] + + +class _ClientWebSocket(Protocol): + exceptions: _ClientWebSocketExceptions + + async def send_text(self, data: str) -> None: ... + async def receive_text(self) -> str: ... + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... @@ -48,13 +62,13 @@ class RealTimeStreaming: logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, request_data: dict | None = None, backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, ): - self.websocket = websocket + self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.messages: list[OpenAIRealtimeEvents] = [] @@ -127,7 +141,7 @@ class RealTimeStreaming: ] ) _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) - _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = { + _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000}, @@ -281,6 +295,7 @@ class RealTimeStreaming: if event_obj.get("type") != "response.done": return response: Final = cast(dict[str, Any], event_obj.get("response", {})) + item: Mapping[str, object] for item in response.get("output", []): if item.get("type") == "function_call": self.tool_calls.append( @@ -384,7 +399,7 @@ class RealTimeStreaming: return message try: - message_obj: Final = json.loads(message) + message_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return message @@ -487,7 +502,7 @@ class RealTimeStreaming: if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: return False try: - msg_obj: Final = json.loads(message) + msg_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return False return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES @@ -555,7 +570,7 @@ class RealTimeStreaming: def _event_to_client_json(self, event: dict) -> str: return json.dumps(self._normalize_event_for_ga_client(event)) - async def _send_event_to_client(self, event: Any, event_str: str) -> bool: + async def _send_event_to_client(self, event: object, event_str: str) -> bool: if self._should_drop_event_from_client(event): return False if isinstance(event, dict): @@ -595,12 +610,12 @@ class RealTimeStreaming: def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, Any]] = { + turn_detection: Final[dict[str, str | bool]] = { "type": "server_vad", "create_response": False, } if self._backend_uses_beta_protocol: - session: dict[str, Any] = {"turn_detection": turn_detection} + session: dict[str, object] = {"turn_detection": turn_detection} else: session = { "type": "realtime", @@ -654,7 +669,7 @@ class RealTimeStreaming: def _has_realtime_guardrails_for_event_hooks( self, - event_hooks: list[Any], + event_hooks: Sequence["GuardrailEventHooks"], ) -> bool: """Return True if any callback would run for one of ``event_hooks``.""" from litellm.integrations.custom_guardrail import CustomGuardrail @@ -699,7 +714,7 @@ class RealTimeStreaming: transcript: str, item_id: str | None = None, pre_block_backend_message: str | None = None, - event_hooks: list[Any] | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ Run registered guardrails on realtime text (transcript, user message, tool output). @@ -753,7 +768,7 @@ class RealTimeStreaming: raise # Extract the human-readable error from the detail dict (HTTPException) # or fall back to str(e) for plain ValueError. - detail = getattr(e, "detail", None) + detail: object | None = getattr(e, "detail", None) if isinstance(detail, dict): safe_msg = detail.get("error") or str(e) elif detail is not None: @@ -826,7 +841,7 @@ class RealTimeStreaming: return True return False - async def _handle_provider_config_message(self, raw_response) -> None: + async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" returned_object: Final = self.provider_config.transform_realtime_response( raw_response, @@ -910,7 +925,7 @@ class RealTimeStreaming: await self._send_event_to_client(event, event_str) @staticmethod - def _parse_backend_event(raw_response: str) -> dict | None: + def _parse_backend_event(raw_response: str) -> dict[str, object] | None: """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" try: event: Final = json.loads(raw_response) @@ -1020,7 +1035,7 @@ class RealTimeStreaming: objects and any test doubles that expose a .scope dict. """ try: - headers: Final = websocket.scope.get("headers", []) + headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", []) for name, value in headers: if isinstance(name, bytes): name = name.decode("latin-1") @@ -1071,9 +1086,9 @@ class RealTimeStreaming: session["output_modalities"] = ["text"] # 3-7. Lift flat audio fields into the nested audio object - audio: Final[dict[str, Any]] = {} - inp: Final[dict[str, Any]] = {} - out: Final[dict[str, Any]] = {} + audio: Final[dict[str, object]] = {} + inp: Final[dict[str, object]] = {} + out: Final[dict[str, object]] = {} # voice → audio.output.voice if "voice" in session: @@ -1190,7 +1205,7 @@ class RealTimeStreaming: # model; check them with the same guardrail used for # user text so an attacker cannot smuggle blocked # content into a function_call_output. - output = item.get("output", "") + output: object = item.get("output", "") output_text = output if isinstance(output, str) else json.dumps(output) if output_text: # Build the sanitized function_call_output up @@ -1241,7 +1256,7 @@ class RealTimeStreaming: # interaction turn. continue elif item.get("role") == "user": - content_list = item.get("content", []) + content_list: Sequence[object] = item.get("content", []) texts = [ c.get("text", "") for c in content_list @@ -1280,7 +1295,7 @@ class RealTimeStreaming: and not self._guardrail_turn_detection_update_sent and self._has_audio_transcription_guardrails() ): - session = msg_obj.setdefault("session", {}) + session: object = msg_obj.setdefault("session", {}) if isinstance(session, dict): existing_td = session.get("turn_detection") if not isinstance(existing_td, dict): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..886ba6a3a18 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,7 +3,7 @@ import time from collections.abc import Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast from litellm._logging import verbose_logger from litellm.types.llms.openai import ( @@ -30,6 +30,7 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) @@ -39,6 +40,60 @@ if TYPE_CHECKING: ) +class _ThinkingBlockFragment(TypedDict, total=False): + type: str | None + data: str | None + thinking: str | None + signature: str | None + + +class _ThinkingDelta(TypedDict, total=False): + thinking_blocks: Sequence[_ThinkingBlockFragment] + + +class _ThinkingChoice(TypedDict, total=False): + delta: _ThinkingDelta + + +class _ThinkingChunk(TypedDict): + choices: Sequence[_ThinkingChoice] + + +class _ContentChoice(TypedDict, total=False): + delta: Mapping[str, str | None] + + +class _ContentChunk(TypedDict): + choices: Sequence[_ContentChoice] + + +class _AudioDelta(TypedDict, total=False): + audio: ChatCompletionAudioDelta | None + + +class _AudioChoice(TypedDict, total=False): + delta: _AudioDelta + + +class _AudioChunk(TypedDict): + choices: Sequence[_AudioChoice] + + +class _UsageBearingChunk(TypedDict, total=False): + usage: Usage | None + _hidden_params: Mapping[str, str] + + +class _UsageSummary(TypedDict): + prompt_tokens: int | None + completion_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + completion_tokens_details: CompletionTokensDetails | None + prompt_tokens_details: PromptTokensDetailsWrapper | None + cost: float | None + + def capture_cache_creation_token_details( prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, @@ -78,7 +133,7 @@ class ChunkProcessor: return [] first_chunk: Final = chunks[0] - first_hidden_params: dict[str, Any] = {} + first_hidden_params: dict[str, object] = {} if isinstance(first_chunk, dict): candidate = first_chunk.get("_hidden_params", {}) if isinstance(candidate, dict): @@ -115,8 +170,8 @@ class ChunkProcessor: @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[object], + logging_obj: "Logging | None" = None, ) -> None: if not chunks: return @@ -456,7 +511,7 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: Sequence["_ContentChunk"], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: Final[list[str]] = [] for chunk in chunks: @@ -475,7 +530,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: Sequence["_ThinkingChunk"] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -532,10 +587,10 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse: base64_data_list: Final[list[str]] = [] transcript_list: Final[list[str]] = [] expires_at: int | None = None @@ -544,7 +599,7 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta") or {} + delta: _AudioDelta = choice.get("delta") or {} audio: ChatCompletionAudioDelta | None = delta.get("audio") if audio is not None: for k, v in audio.items(): @@ -565,7 +620,7 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: + def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary": prompt_tokens = 0 completion_tokens = 0 ## anthropic prompt caching information ## @@ -623,8 +678,8 @@ class ChunkProcessor: return reasoning_tokens @staticmethod - def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None: - usage_chunk: Usage | dict[str, Any] | None = None + def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None: + usage_chunk: Usage | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -640,7 +695,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -721,13 +776,7 @@ class ChunkProcessor: "web_search_requests", ) - prompt_tokens_details = ( - cast( - PromptTokensDetailsWrapper | None, - usage_chunk_dict["prompt_tokens_details"], - ) - or prompt_tokens_details - ) + prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details cache_creation_token_details = capture_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details @@ -758,7 +807,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -797,7 +846,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, messages: list | None = None, @@ -851,8 +900,8 @@ class ChunkProcessor: setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate( + completion_tokens_details.model_dump() ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a9751489473..36f3e875a7e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,8 +1,9 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import ( TYPE_CHECKING, Any, Final, + TypeAlias, cast, ) @@ -33,8 +34,12 @@ if TYPE_CHECKING: # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"}) +_AnthropicMessages: TypeAlias = "list[dict[str, object]]" +_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None" +_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None" -def _messages_have_compaction_block(messages: list[dict]) -> bool: + +def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -54,8 +59,10 @@ def _proxy_router_fallback() -> "Router | None": return _proxy_router -def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: - """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise. +def _extract_proxy_litellm_metadata( + kwargs: Mapping[str, object], +) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]": + """Return ``(kwargs["litellm_metadata"], its user_api_key_auth)`` when it's a dict; ``(None, None)`` otherwise. The proxy attaches its auth/spend-attribution fields (``user_api_key``, ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth`` @@ -68,18 +75,19 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | """ litellm_metadata: Final = kwargs.get("litellm_metadata") if not isinstance(litellm_metadata, dict): - return None - return litellm_metadata + return None, None + user_api_key_auth: Final[UserAPIKeyAuth | None] = litellm_metadata.get("user_api_key_auth") + return litellm_metadata, user_api_key_auth async def _prepare_context_managed_request( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -102,11 +110,11 @@ async def _prepare_context_managed_request( if polyfill_will_run: history_result: PolyfillResult | None = None - working_messages: list[dict] = messages - working_system: Any | None = system + working_messages: _AnthropicMessages = messages + working_system: _AnthropicSystem = system else: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) working_messages = history_result.messages if history_result is not None else messages @@ -136,7 +144,7 @@ async def _prepare_context_managed_request( # to non-Anthropic backends that would reject them. if polyfill_will_run and history_result is None: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) return history_result @@ -144,7 +152,7 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -171,7 +179,7 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -209,9 +217,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run. Delegates spec-shape normalization to the dispatcher's ``_normalize_spec`` @@ -236,11 +244,11 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -306,7 +314,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _route_openai_thinking_to_responses_api_if_needed( completion_kwargs: dict[str, Any], *, - thinking: dict[str, Any] | None, + thinking: Mapping[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -407,12 +415,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _prepare_completion_kwargs( *, max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: _AnthropicSystem = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, @@ -420,7 +428,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, + extra_kwargs: Mapping[str, object] | None = None, ) -> tuple[dict[str, Any], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. @@ -433,7 +441,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -528,7 +536,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -537,7 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -551,10 +559,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: requested_router if requested_router is not None else _proxy_router_fallback() ) - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result: Final = await _prepare_context_managed_request( model=model, @@ -618,7 +623,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -627,7 +632,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -688,10 +693,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if context_management is None and not _messages_have_compaction_block(messages): polyfill_result: PolyfillResult | None = None else: - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result = run_async_function( _prepare_context_managed_request, model=model, diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 671e4633af4..f7b419405ac 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -1,8 +1,9 @@ from collections.abc import Coroutine, Iterable -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypedDict import httpx from openai import AsyncAzureOpenAI, AzureOpenAI +from openai.types.shared_params.metadata import Metadata from typing_extensions import overload from ...types.llms.openai import ( @@ -22,6 +23,16 @@ from ...types.llms.openai import ( from .common_utils import BaseAzureLLM +class _RunThreadStreamData(TypedDict): + thread_id: str + assistant_id: str + additional_instructions: str | None + instructions: str | None + metadata: Metadata | None + model: str | None + tools: Iterable[AssistantToolParam] | None + + class AzureAssistantsAPI(BaseAzureLLM): def __init__(self) -> None: super().__init__() @@ -212,9 +223,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj # fmt: off @@ -301,9 +312,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj async def async_get_messages( @@ -443,7 +454,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) # fmt: off @@ -539,7 +550,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = azure_openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) async def async_get_thread( self, @@ -566,7 +577,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # fmt: off @@ -642,7 +653,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # def delete_thread(self): # pass @@ -730,7 +741,8 @@ class AzureAssistantsAPI(BaseAzureLLM): event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AssistantStreamManager[AssistantEventHandler]: - data: Final[dict[str, Any]] = { + stream_fn: Final = client.beta.threads.runs.stream + base_data: Final[_RunThreadStreamData] = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, @@ -740,8 +752,8 @@ class AzureAssistantsAPI(BaseAzureLLM): "tools": tools, } if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return stream_fn(**base_data, event_handler=event_handler) + return stream_fn(**base_data) # fmt: off diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 76618e0f742..e285feb77ee 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -109,7 +109,7 @@ if MCP_AVAILABLE: ############ MCP Server REST API Routes ################# async def _safe_fire_mcp_tool_call_logging( logging_obj: Any | None, - result: Any, + result: "CallToolResult", start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1c6ad84ddb4..49a1f1314f0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,9 +13,9 @@ import time import traceback import types import uuid -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -145,7 +145,7 @@ try: ) # Robust auth lookup keyed by session_object. - _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() + _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() except ImportError as e: verbose_logger.debug("MCP module not found: %s", e) MCP_AVAILABLE = False @@ -493,14 +493,14 @@ if MCP_AVAILABLE: def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, - experimental_capabilities: dict[str, dict[str, Any]] | None = None, + experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: opts: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) - updates: Final[dict[str, Any]] = {} + updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: updates["instructions"] = merged @@ -549,6 +549,17 @@ if MCP_AVAILABLE: _stateful_session_locks: Final[dict[str, asyncio.Lock]] = {} _stateful_session_active_request_counts: Final[dict[str, int]] = {} + class _TerminableTransport(Protocol): + async def terminate(self) -> None: ... + + class _TransportRegistry(Protocol): + def __contains__(self, session_id: object, /) -> bool: ... + + def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ... + + def _stateful_server_instances() -> _TransportRegistry: + return getattr(session_manager_stateful, "_server_instances", {}) + def _remove_stateful_session_tracking(session_id: str) -> None: _stateful_session_auth_contexts.pop(session_id, None) _stateful_session_auth_context_last_seen.pop(session_id, None) @@ -578,8 +589,8 @@ if MCP_AVAILABLE: ) -> None: """Terminate expired stateful sessions and drop their auth contexts.""" now = time.monotonic() if now is None else now - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) - expired_session_ids: Final = [] + server_instances: Final = _stateful_server_instances() + expired_session_ids: Final[list[str]] = [] for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): if _stateful_session_active_request_counts.get(session_id, 0) > 0: continue @@ -619,7 +630,7 @@ if MCP_AVAILABLE: session may proceed, or ``False`` when the caller is already at the cap with every session in flight (the new ``initialize`` should be rejected). """ - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) + server_instances: Final = _stateful_server_instances() def _owned_live_session_ids() -> list[str]: return [ @@ -778,7 +789,7 @@ if MCP_AVAILABLE: get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -847,7 +858,7 @@ if MCP_AVAILABLE: async def _build_virtual_call_logging_obj( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual @@ -885,7 +896,7 @@ if MCP_AVAILABLE: async def _dispatch_virtual_mcp_tool( name: str, - arguments: dict[str, Any] | None, + arguments: dict[str, object] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, mcp_servers: list[str] | None = None, @@ -957,7 +968,7 @@ if MCP_AVAILABLE: ) @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -1621,7 +1632,7 @@ if MCP_AVAILABLE: async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, - prefetched_creds: dict[str, dict[str, Any]] | None = None, + prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. @@ -1646,7 +1657,7 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not user_id: return {} try: @@ -1871,7 +1882,7 @@ if MCP_AVAILABLE: list_tools_start_time: Final = datetime.now() litellm_logging_obj: LiteLLMLoggingObj | None = None - list_tools_request_data: dict[str, Any] = {} + list_tools_request_data: dict[str, object] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook @@ -1879,7 +1890,7 @@ if MCP_AVAILABLE: list_tools_call_id: Final = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Final[dict[str, Any]] = { + spend_logs_metadata: Final[dict[str, object]] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): @@ -2615,7 +2626,7 @@ if MCP_AVAILABLE: async def execute_mcp_tool( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -2882,7 +2893,7 @@ if MCP_AVAILABLE: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) @@ -2956,7 +2967,7 @@ if MCP_AVAILABLE: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args local_content = await _handle_local_mcp_tool(original_tool_name, arguments) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) return await _run_post_mcp_call_guardrails( result=response, @@ -3003,7 +3014,7 @@ if MCP_AVAILABLE: async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, - result: Any, + result: CallToolResult, start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -3070,7 +3081,7 @@ if MCP_AVAILABLE: @client async def call_mcp_tool( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3161,7 +3172,7 @@ if MCP_AVAILABLE: async def mcp_get_prompt( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3262,7 +3273,7 @@ if MCP_AVAILABLE: def _get_standard_logging_mcp_tool_call( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: @@ -3291,13 +3302,13 @@ if MCP_AVAILABLE: async def _handle_managed_mcp_tool( server_name: str, name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, - litellm_logging_obj: Any | None = None, + litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" @@ -3320,7 +3331,7 @@ if MCP_AVAILABLE: return call_tool_result async def _handle_local_mcp_tool( - name: str, arguments: dict[str, Any] + name: str, arguments: dict[str, object] ) -> list[TextContent | ImageContent | EmbeddedResource]: """ Handle tool execution for local registry tools @@ -3426,7 +3437,8 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - for header_name, header_value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + for header_name, header_value in scope_headers: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": return header_value.decode() if isinstance(header_value, bytes) else str(header_value) @@ -3528,7 +3540,7 @@ if MCP_AVAILABLE: if message.get("type") != "http.request": break - body = message.get("body", b"") or b"" + body: bytes = message.get("body", b"") or b"" if body: # Only retain up to the remaining peek budget for sniffing. # The full ``message`` is already in memory (delivered by @@ -3571,9 +3583,9 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final = scope.get("headers", []) + _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) - def _normalize_header_name(header_name: Any) -> bytes | None: + def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): return header_name.lower() if isinstance(header_name, str): @@ -3902,7 +3914,8 @@ if MCP_AVAILABLE: def _get_authorization_header_from_scope(scope: Scope) -> str | None: """First ``Authorization`` header value in the ASGI scope, or None.""" - for key, value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + for key, value in scope_headers: if key.lower() == b"authorization": return value.decode("latin-1") return None @@ -3921,7 +3934,8 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers) if not has_litellm_key_header: return None return _get_authorization_header_from_scope(scope) @@ -4115,7 +4129,7 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4135,7 +4149,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). @@ -4436,7 +4451,7 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4456,7 +4471,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar so the # downstream probe list matches the fully-authorized server set @@ -4680,7 +4696,8 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - for key, value in message.get("headers", []): + response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) + for key, value in response_headers: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": session_id = value.decode() if isinstance(value, bytes) else str(value) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..b68d4a68b79 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -6,10 +6,10 @@ import concurrent.futures import inspect import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone from types import UnionType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request @@ -54,8 +54,8 @@ from litellm.types.guardrails import ( if TYPE_CHECKING: from types import CodeType - from prisma.actions import LiteLLM_GuardrailsTableActions from prisma.models import LiteLLM_GuardrailsTable + from pydantic.fields import FieldInfo from litellm.proxy.utils import PrismaClient @@ -65,24 +65,44 @@ router: Final = APIRouter() GUARDRAIL_REGISTRY: Final = GuardrailRegistry() -def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]": - table: Final[LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]] = GuardrailsRepository(prisma_client).table +class _GuardrailsTableActions(Protocol): + async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ... + + async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_many( + self, where: Mapping[str, object], order: Mapping[str, str] + ) -> "Sequence[LiteLLM_GuardrailsTable]": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_GuardrailsTable | None": ... + + +def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]: + return mapping + + +def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions: + table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table return table async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": - row: Final[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.create(data=data) + row: Final = await _guardrails_table(prisma_client).create(data=data) return row async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None: - await GuardrailsRepository(prisma_client).table.delete(where=where) + await _guardrails_table(prisma_client).delete(where=where) async def _find_team_guardrail_rows( prisma_client: "PrismaClient", where: Mapping[str, object] ) -> "Sequence[LiteLLM_GuardrailsTable]": - rows: Final[Sequence[LiteLLM_GuardrailsTable]] = await GuardrailsRepository(prisma_client).table.find_many( + rows: Final = await _guardrails_table(prisma_client).find_many( where=where, order={"created_at": "desc"}, ) @@ -499,10 +519,12 @@ async def update_guardrail( if existing_guardrail is None: raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=request.guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=request.guardrail, + prisma_client=prisma_client, + ) ) guardrail_name: Final = result.get("guardrail_name", "Unknown") @@ -613,7 +635,7 @@ class RegisterGuardrailRequest(BaseModel): """Request body for POST /guardrails/register. Follows Generic Guardrail API config.""" guardrail_name: str - litellm_params: dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional + litellm_params: dict[str, object] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: dict[str, object] | None = None team_id: str | None = None @@ -1172,12 +1194,14 @@ async def patch_guardrail( ) # Update litellm_params if default_on is provided or pii_entities_config is provided - litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {}))) + existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {}))) + litellm_params = LitellmParams(**existing_litellm_params) if request.litellm_params is not None: requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True) litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) - litellm_params = LitellmParams(**litellm_params_dict) + merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) + litellm_params = LitellmParams(**merged_litellm_params) # Update guardrail_info if provided guardrail_info: Final = ( @@ -1193,10 +1217,12 @@ async def patch_guardrail( litellm_params=litellm_params, guardrail_info=guardrail_info, ) - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=guardrail, + prisma_client=prisma_client, + ) ) guardrail_name = result.get("guardrail_name", "Unknown") @@ -1552,31 +1578,46 @@ async def validate_blocked_words_file(request: dict[str, str]): return {"valid": False, "error": f"Validation error: {e}"} -def _get_field_type_from_annotation(field_annotation: Any) -> str: +def _dunder_origin(annotation: object) -> object: + origin: Final[object] = getattr(annotation, "__origin__", None) + return origin + + +def _dunder_name(annotation: object) -> object: + name: Final[object] = getattr(annotation, "__name__", None) + return name + + +def _dunder_args(annotation: object) -> tuple[object, ...]: + args: Final[tuple[object, ...]] = getattr(annotation, "__args__", ()) + return args + + +def _get_field_type_from_annotation(field_annotation: object) -> str: """ Convert a Python type annotation to a UI-friendly type string """ # Handle Union types (like Optional[T]) if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[T], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: field_annotation = non_none_args[0] # Handle List types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is list: return "array" # Handle Dict types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is dict: return "dict" # Handle Literal types if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"): # Check for Literal types (Python 3.8+) - origin: Final = field_annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": + origin: Final = _dunder_origin(field_annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": return "select" # For dropdown/select inputs # Handle basic types @@ -1595,66 +1636,66 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str: return "string" -def _extract_literal_values(annotation: Any) -> list[str]: +def _extract_literal_values(annotation: object) -> Sequence[object]: """ Extract literal values from a Literal type annotation """ if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"): - origin: Final = annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": - return list(annotation.__args__) + origin: Final = _dunder_origin(annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": + return list(_dunder_args(annotation)) return [] -def _get_dict_key_options(field_annotation: Any) -> list[str] | None: +def _get_dict_key_options(field_annotation: object) -> Sequence[object] | None: """ Extract key options from Dict[Literal[...], T] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: key_type: Final = args[0] return _extract_literal_values(key_type) return None -def _get_dict_value_type(field_annotation: Any) -> str: +def _get_dict_value_type(field_annotation: object) -> str: """ Get the value type from Dict[K, V] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: value_type: Final = args[1] return _get_field_type_from_annotation(value_type) return "string" -def _get_list_element_options(field_annotation: Any) -> list[str] | None: +def _get_list_element_options(field_annotation: object) -> Sequence[object] | None: """ Extract element options from List[Literal[...]] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is list + and _dunder_origin(field_annotation) is list and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 1: element_type: Final = args[0] return _extract_literal_values(element_type) return None -def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool: +def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool: """Check if optional_params field should be skipped (not meaningfully overridden).""" if field_name != "optional_params": return False @@ -1664,12 +1705,12 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Check if the annotation is still a generic TypeVar (not specialized) if isinstance(field_annotation, TypeVar) or ( - hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar + hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is TypeVar ): return True # Also skip if it's a generic type that wasn't specialized - if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( + if hasattr(field_annotation, "__name__") and _dunder_name(field_annotation) in ( "T", "TypeVar", ): @@ -1677,18 +1718,18 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Handle Optional[T] where T is still a TypeVar if hasattr(field_annotation, "__args__"): - non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)] + non_none_args: Final = [arg for arg in _dunder_args(field_annotation) if arg is not type(None)] if non_none_args and isinstance(non_none_args[0], TypeVar): return True return False -def _unwrap_optional_type(field_annotation: Any) -> Any: +def _unwrap_optional_type(field_annotation: object) -> object: """Unwrap Optional types to get the actual type.""" if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[BaseModel], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: return non_none_args[0] @@ -1696,20 +1737,20 @@ def _unwrap_optional_type(field_annotation: Any) -> Any: def _build_field_dict( - field: Any, - field_annotation: Any, + field: "FieldInfo", + field_annotation: object, description: str, required: bool, -) -> dict[str, Any]: +) -> dict[str, object]: """Build field dictionary for non-nested fields.""" # Determine the field type from annotation field_type = _get_field_type_from_annotation(field_annotation) # Check for custom UI type override - field_json_schema_extra: Final = getattr(field, "json_schema_extra", {}) + field_json_schema_extra: Final[Mapping[str, object]] = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: ui_type: Final = field_json_schema_extra["ui_type"] - field_type = ui_type.value if hasattr(ui_type, "value") else ui_type + field_type = getattr(ui_type, "value", ui_type) elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] @@ -1748,8 +1789,9 @@ def _build_field_dict( field_dict["options"] = literal_options # Add default value if it exists - if field.default is not None and field.default is not ...: - field_dict["default_value"] = field.default + field_default: Final[object] = getattr(field, "default", None) + if field_default is not None and field_default is not ...: + field_dict["default_value"] = field_default # Copy min, max, step from json_schema_extra for number/percentage inputs if field_json_schema_extra: @@ -1763,7 +1805,7 @@ def _build_field_dict( def _extract_fields_recursive( model: type[BaseModel], depth: int = 0, -) -> dict[str, Any]: +) -> dict[str, object]: # Check if we've exceeded the maximum recursion depth if depth > DEFAULT_MAX_RECURSE_DEPTH: raise HTTPException( @@ -1817,7 +1859,7 @@ def _extract_fields_recursive( return fields -def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]: +def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]: """ Get the fields from a Pydantic model as a nested dictionary structure """ @@ -2141,7 +2183,26 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: return "response" if input_type == "response" else "request" -def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None: +class _GuardrailLoggingObj(Protocol): + call_type: str + model_call_details: dict[str, object] + + @property + def update_messages(self) -> "Callable[..., object]": ... + + @property + def async_success_handler(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def success_handler(self) -> "Callable[..., object]": ... + + +class _GuardrailProxyLogging(Protocol): + @property + def post_call_success_hook(self) -> "Callable[..., Awaitable[object]]": ... + + +def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj, request: ApplyGuardrailRequest) -> None: """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" litellm_logging_obj.call_type = "pass_through_endpoint" litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" @@ -2151,8 +2212,8 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGua async def _emit_guardrail_success_logs( - proxy_logging_obj: Any, - litellm_logging_obj: Any, + proxy_logging_obj: _GuardrailProxyLogging, + litellm_logging_obj: _GuardrailLoggingObj | None, data: dict, user_api_key_dict: UserAPIKeyAuth, response: ApplyGuardrailResponse, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 068a3ecf31b..facb822d00d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -19,7 +19,7 @@ request is sent with the ``X-Cisco-AI-Defense-API-Key`` header. import json import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -94,13 +94,13 @@ class _CiscoVerdict: is_safe: bool | None classifications: list[str] severity: str | None - rules: list[dict[str, Any]] + rules: list[dict[str, object]] explanation: str | None event_id: str | None action: str | None = None sanitized_text: str | None = None - sanitized_messages: list[dict[str, Any]] | None = None - sanitized_mcp_arguments: dict[str, Any] | None = None + sanitized_messages: list[dict[str, object]] | None = None + sanitized_mcp_arguments: dict[str, object] | None = None class CiscoAIDefenseGuardrailMissingSecrets(Exception): @@ -136,7 +136,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): api_base: str | None = None, inspection_type: str | None = None, inspect_path: str | None = None, - enabled_rules: list[dict[str, Any]] | None = None, + enabled_rules: Sequence[object] | None = None, integration_profile_id: str | None = None, integration_profile_version: str | None = None, integration_tenant_id: str | None = None, @@ -415,7 +415,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: AsyncIterator[Any], + response: AsyncIterator[object], request_data: dict, ): """Buffer and inspect streaming chat output before delivery.""" @@ -437,7 +437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self.guardrail_name, ) - all_chunks: Final[list[Any]] = [] + all_chunks: Final[list[object]] = [] try: async for chunk in response: all_chunks.append(chunk) @@ -497,7 +497,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): response_obj=assembled, ) except HTTPException as exc: - error_obj: dict[str, Any] = self._http_exception_to_error_obj(exc) + error_obj: dict[str, object] = self._http_exception_to_error_obj(exc) verbose_proxy_logger.warning( "Cisco AI Defense guardrail (%s): streaming response " "blocked — emitting SSE error event instead of " @@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): for chunk in all_chunks: yield chunk - def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, Any]: + def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, object]: """Canonical block payload used across all four block paths. Same dict is the ``HTTPException.detail`` for chat / MCP request @@ -555,34 +555,34 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): "event_id": verdict.event_id, } - def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, Any]: + def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, object]: """Wrap an ``HTTPException`` detail into the SSE ``error`` payload. For Cisco's own blocks the detail is already the canonical block payload, so this is a near-passthrough that just adds ``code`` / ``guardrail`` defaults for non-Cisco / unstructured details. """ - error_obj: dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)} + error_obj: dict[str, object] = {**exc.detail} if isinstance(exc.detail, dict) else {"message": str(exc.detail)} error_obj.setdefault("message", error_obj.get("error", "Guardrail block")) error_obj.setdefault("code", exc.status_code) error_obj.setdefault("guardrail", self.guardrail_name) return error_obj @classmethod - def _streaming_content_was_modified(cls, original_chunks: list[Any], assembled: ModelResponse) -> bool: + def _streaming_content_was_modified(cls, original_chunks: Sequence[object], assembled: ModelResponse) -> bool: """Decide whether redact changed content or tool/function arguments.""" original_text: Final = cls._extract_streaming_chunk_scan_text(original_chunks) assembled_text: Final = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled)) return original_text != assembled_text @classmethod - def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str: + def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str: original_text = "" argument_text = "" for chunk in chunks: choices = getattr(chunk, "choices", None) or [] for c in choices: - delta = getattr(c, "delta", None) + delta: object | None = getattr(c, "delta", None) if delta is None: continue text = getattr(delta, "content", None) @@ -595,7 +595,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): args = cls._extract_tool_call_arguments(tc) if args: argument_text += args - fc = getattr(delta, "function_call", None) + fc: object | None = getattr(delta, "function_call", None) if fc is not None: args = cls._extract_function_call_arguments(fc) if args: @@ -673,7 +673,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): allow, WARNING for intervened/redacted, ERROR is left for upstream API failures. """ - fields: Final[dict[str, Any]] = { + fields: Final[dict[str, object]] = { "guardrail": self.guardrail_name, "surface": context.surface, "direction": context.direction, @@ -752,7 +752,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, direction: str = "input", response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_chat_payload(messages, request_data, user_api_key_dict) start_time: Final = datetime.now() @@ -784,7 +784,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): messages: list[dict[str, str]], request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: + ) -> dict[str, object]: return { "messages": messages, "metadata": self._build_metadata(request_data, user_api_key_dict), @@ -798,9 +798,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def _post_inspection( self, url: str, - payload: dict[str, Any], + payload: dict[str, object], surface: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: headers: Final = self._build_headers() verbose_proxy_logger.debug( "Cisco AI Defense guardrail: posting %s inspection to %s", @@ -856,8 +856,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: - metadata: Final[dict[str, Any]] = {} + ) -> dict[str, object]: + metadata: Final[dict[str, object]] = {} user: Final = request_data.get("user") or getattr(user_api_key_dict, "user_id", None) if user: @@ -884,8 +884,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return metadata - def _build_config(self) -> dict[str, Any]: - config: Final[dict[str, Any]] = {} + def _build_config(self) -> dict[str, object]: + config: Final[dict[str, object]] = {} if self.enabled_rules: config["enabled_rules"] = self.enabled_rules if self.integration_profile_id: @@ -899,7 +899,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return config @staticmethod - def _normalize_rule(rule: object) -> dict[str, Any]: + def _normalize_rule(rule: object) -> dict[str, object]: """Coerce a user-supplied rule into the wire-shape dict Cisco expects. Accepts ``str``, ``dict``, and Pydantic model inputs. @@ -922,7 +922,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): rule = dumped if isinstance(rule, dict): - normalized: Final[dict[str, Any]] = {} + normalized: Final[dict[str, object]] = {} rule_name: Final = rule.get("rule_name") if rule_name: normalized["rule_name"] = rule_name @@ -950,7 +950,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): context: _ScanContext, start_time: datetime, response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Parse, log, and (optionally) raise/redact on the Cisco verdict. ``context.direction`` is ``"input"`` for request scans and ``"output"`` @@ -1119,10 +1119,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _sanitize_response_for_logging( cls, - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], surface: str, action: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Drop bulky / privacy-sensitive fields, recursing into nested dicts. MCP verdicts are commonly nested under ``result``, so a @@ -1138,9 +1138,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return sanitized @classmethod - def _strip_sensitive_keys(cls, d: dict[str, Any]) -> dict[str, Any]: + def _strip_sensitive_keys(cls, d: Mapping[str, object]) -> dict[str, object]: """Recursively strip privacy-sensitive keys from a verdict dict.""" - out: Final[dict[str, Any]] = {} + out: Final[dict[str, object]] = {} for key, value in d.items(): if key.startswith("_") or key in cls._REDACTED_LOG_KEYS: continue @@ -1222,8 +1222,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_jsonrpc_error( - inspect_response: dict[str, Any], - ) -> dict[str, Any] | None: + inspect_response: Mapping[str, object], + ) -> dict[str, object] | None: """Detect a JSON-RPC error envelope inside an HTTP 200 response. The Cisco Inspect API can return ``{"error": {...}}`` (or nest one @@ -1270,7 +1270,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_text( - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], ) -> str | None: """Pull ``sanitized_text`` (or camelCase variant) off the verdict.""" for key in ("sanitized_text", "sanitizedText"): @@ -1287,8 +1287,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_messages( - inspect_response: dict[str, Any], - ) -> list[dict[str, Any]] | None: + inspect_response: Mapping[str, object], + ) -> list[dict[str, object]] | None: """Pull a sanitized OpenAI-format messages array off the verdict. Cisco can return the rewrite under several keys; we accept any of @@ -1354,7 +1354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_mcp_input( request_data: dict, sanitized_text: str | None, - sanitized_mcp_arguments: dict[str, Any] | None, + sanitized_mcp_arguments: dict[str, object] | None, ) -> bool: """Rewrite MCP request arguments in all locations the proxy reads.""" if sanitized_mcp_arguments is not None: @@ -1388,7 +1388,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat request input (``messages`` or ``input``).""" if sanitized_messages and self._extract_tool_definition_text(request_data): @@ -1444,7 +1444,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): cls, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: if sanitized_messages: instruction_text: Final = cls._instruction_text_from_messages(sanitized_messages) @@ -1457,7 +1457,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return False @classmethod - def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str | None: + def _instruction_text_from_messages(cls, messages: list[dict[str, object]]) -> str | None: for message in messages: if not isinstance(message, dict): continue @@ -1468,7 +1468,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return None @classmethod - def _non_instruction_messages(cls, messages: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None: + def _non_instruction_messages(cls, messages: list[dict[str, object]] | None) -> list[dict[str, object]] | None: if messages is None: return None return [ @@ -1499,7 +1499,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, response_obj: object, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat response (``ModelResponse`` or ``ResponsesAPIResponse``).""" if response_obj is None: @@ -1526,7 +1526,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_model_response_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Redact every returned choice, including tool-call/reasoning fields.""" if sanitized_messages: @@ -1570,7 +1570,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_text_completion_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite ``/v1/completions`` text choices after Cisco redaction.""" replacement = sanitized_text @@ -1638,7 +1638,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, output_items: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: replacement_text: str | None = sanitized_text if not replacement_text and sanitized_messages: @@ -1672,14 +1672,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _sanitized_messages_to_responses_input( - sanitized_messages: list[dict[str, Any]], - ) -> list[dict[str, Any]] | None: + sanitized_messages: list[dict[str, object]], + ) -> list[dict[str, object]] | None: """Convert chat-shape sanitized_messages to Responses API ``input``. Returns ``None`` if nothing usable could be converted, so the caller falls back to ``on_flagged_action``. """ - out: Final[list[dict[str, Any]]] = [] + out: Final[list[dict[str, object]]] = [] for m in sanitized_messages: if not isinstance(m, dict): continue @@ -1764,7 +1764,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): start_time: datetime | None = None, surface: str = "chat", direction: str = "input", - ) -> dict[str, Any]: + ) -> dict[str, object]: verbose_proxy_logger.error( "Cisco AI Defense guardrail (%s): API communication failed: %s", surface, @@ -2060,7 +2060,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return getattr(obj, key, None) @classmethod - def _field_list(cls, obj: object, key: str) -> list[Any]: + def _field_list(cls, obj: object, key: str) -> list[object]: value: Final = cls._field(obj, key) return value if isinstance(value, list) else [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c29da89b15f..5bbb01c6c8e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,8 +8,8 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import HTTPException @@ -34,6 +34,9 @@ if TYPE_CHECKING: # Imported lazily at runtime (inside the streaming hook) to avoid a # module-level cyclic import with litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) @@ -41,12 +44,35 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +class _EndpointTranslation(Protocol): + @property + def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ... + + @property + def process_output_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + + +def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: + return translation + + +def _chunk_choices(item: object) -> Sequence[object]: + choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] + return choices + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" -def _get_a2a_request_id(responses_so_far: list[Any], request_data: dict) -> str | None: +def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None: """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" for item in responses_so_far: if isinstance(item, dict) and "id" in item: @@ -138,7 +164,9 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -156,7 +184,7 @@ class UnifiedLLMGuardrails(CustomLogger): async def async_moderation_hook( self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral - ) -> Any: + ) -> object: """ Runs in parallel to LLM API call Runs on only Input @@ -187,7 +215,9 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: return data - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -202,7 +232,7 @@ class UnifiedLLMGuardrails(CustomLogger): data: dict, user_api_key_dict: UserAPIKeyAuth, response, - ) -> Any: + ) -> object: """ Runs on response from LLM API call @@ -271,7 +301,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) try: response = await endpoint_translation.process_output_response( @@ -299,10 +331,10 @@ class UnifiedLLMGuardrails(CustomLogger): async def _handle_streaming_block( self, exc: "ModifyResponseException", - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, stream_started: bool, - responses_so_far: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + ) -> AsyncGenerator[object, None]: """ Terminate a streamed response cleanly when a guardrail blocks it. @@ -323,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _resolve_transform_call_type( user_api_key_dict: UserAPIKeyAuth, - mappings: dict, + mappings: Mapping[CallTypes, type["BaseTranslation"]], ) -> str | None: """Resolve the call type for the incremental_diff path, or None if the route is unresolvable / unsupported. @@ -356,9 +388,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, exc: HTTPException, call_type: str | None, - responses_so_far: list[Any], + responses_so_far: Sequence[object], request_data: dict, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the response has already started, so emit an in-stream JSON-RPC error chunk; otherwise re-raise so the proxy can report it. @@ -387,7 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger): def _build_transform_chunk( self, *, - reference_chunk: Any, + reference_chunk: object, mutated_text_per_choice: dict[int, str], emitted_text_per_choice: dict[int, str], holdback_per_choice: dict[int, int], @@ -500,18 +532,18 @@ class UnifiedLLMGuardrails(CustomLogger): async def _emit_transform_round( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, - reference_chunk: Any, - responses_so_far: list[Any], - responses_yielded: list[Any], + reference_chunk: object, + responses_so_far: Sequence[object], + responses_yielded: list[object], emitted_text_per_choice: dict[int, str], finish_reason_per_choice: dict[int, str | None], is_final: bool, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Run one guardrail processing round and emit the resulting diff chunk. Raises ``_StreamTerminated`` (after emitting the terminal block message or @@ -564,14 +596,14 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, guardrail_to_apply: CustomGuardrail, - response: Any, + response: AsyncIterable[object], request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, sampling_rate: int, end_of_stream_only: bool, - mappings: dict, - ) -> AsyncGenerator[Any, None]: + mappings: Mapping[CallTypes, type["BaseTranslation"]], + ) -> AsyncGenerator[object, None]: """Emit guardrail text transformations as new deltas on the stream. Raw chunks are withheld and accumulated; on each sampled processing round @@ -580,15 +612,15 @@ class UnifiedLLMGuardrails(CustomLogger): synthetic chunk. A BLOCK terminates the stream via the shared block handler; an underflow surfaces as an HTTPException. """ - endpoint_translation: Final = mappings[CallTypes(call_type)]() - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] emitted_text_per_choice: Final[dict[int, str]] = {} finish_reason_per_choice: Final[dict[int, str | None]] = {} chunk_counter = 0 - last_chunk: Any | None = None + last_chunk: object | None = None - def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]: + def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]: return self._emit_transform_round( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, @@ -694,13 +726,13 @@ class UnifiedLLMGuardrails(CustomLogger): async def _inspect_full_response_for_block( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - responses_so_far: list[Any], - responses_yielded: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + responses_yielded: Sequence[object], + ) -> AsyncGenerator[object, None]: """Run the block-only guardrail inspection over the full assembled response (text + tool calls) so nothing bypasses the block decision. @@ -734,17 +766,17 @@ class UnifiedLLMGuardrails(CustomLogger): raise _StreamTerminated() @staticmethod - def _chunk_has_tool_calls(item: Any) -> bool: - for choice in getattr(item, "choices", None) or []: + def _chunk_has_tool_calls(item: object) -> bool: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) if getattr(delta, "tool_calls", None): return True return False @staticmethod - def _chunk_carries_text(item: Any) -> bool: + def _chunk_carries_text(item: object) -> bool: """True if any choice in this chunk has non-empty string ``delta.content``.""" - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) content = getattr(delta, "content", None) if isinstance(content, str) and content != "": @@ -753,7 +785,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _tool_call_passthrough_chunk( - item: Any, + item: object, finish_reason_per_choice: "dict[int, str | None] | None" = None, ) -> ModelResponseStream: """Copy of a chunk carrying tool calls with all text content stripped. @@ -772,7 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger): redaction purpose. """ synthetic_choices: Final[list[StreamingChoices]] = [] - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) idx = getattr(choice, "index", 0) or 0 original_finish = getattr(choice, "finish_reason", None) @@ -801,15 +833,15 @@ class UnifiedLLMGuardrails(CustomLogger): ) @staticmethod - def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None: - for choice in getattr(item, "choices", None) or []: + def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None: + for choice in _chunk_choices(item): finish_reason = getattr(choice, "finish_reason", None) if finish_reason is not None: finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason @staticmethod - def _chunk_has_finish_reason(item: Any) -> bool: - choices: Final = getattr(item, "choices", None) or [] + def _chunk_has_finish_reason(item: object) -> bool: + choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) async def async_post_call_streaming_iterator_hook( @@ -845,22 +877,22 @@ class UnifiedLLMGuardrails(CustomLogger): # Get streaming configuration. Resolution order (later wins): default # < guardrail attribute < guardrail_config dict < this callback's # optional_params. - def _streaming_flag(name: str, default: Any) -> Any: + def _streaming_flag(name: str, default: object) -> Any: value = default if guardrail_to_apply is not None: value = getattr(guardrail_to_apply, name, value) - config: Final = getattr(guardrail_to_apply, "guardrail_config", {}) + config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) if isinstance(config, dict): value = config.get(name, value) return self.optional_params.get(name, value) - sampling_rate: Final = _streaming_flag("streaming_sampling_rate", 5) + sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). - end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False) + end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False) # "block_only" (default) drops guardrail text rewrites on the streaming # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). - streaming_transform_mode: Final = _streaming_flag("streaming_transform_mode", "block_only") + streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -868,7 +900,9 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default) + buffer_until_moderated: bool = _streaming_flag( + "streaming_buffer_until_moderated", buffer_until_moderated_default + ) if ( buffer_until_moderated @@ -939,9 +973,9 @@ class UnifiedLLMGuardrails(CustomLogger): # Infer call type from first chunk call_type = None chunk_counter = 0 - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] - pending_end_of_stream_items: Final[list[Any]] = [] + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] + pending_end_of_stream_items: Final[list[object]] = [] # Whether any real response chunk has been forwarded to the client. # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index dd61cad15a1..a64ed764a67 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -26,7 +26,8 @@ Usage: import base64 import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -36,7 +37,10 @@ from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, ) from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes + +if TYPE_CHECKING: + from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor class SkillsInjectionHook(CustomLogger): @@ -99,7 +103,7 @@ class SkillsInjectionHook(CustomLogger): verbose_proxy_logger.debug("SkillsInjectionHook: Processing %s skills", len(skills)) litellm_skills: Final[list[LiteLLM_SkillsTable]] = [] - anthropic_skills: Final[list[dict[str, Any]]] = [] + anthropic_skills: Final[list[dict[str, object]]] = [] # Separate skills by prefix for skill in skills: @@ -324,9 +328,9 @@ class SkillsInjectionHook(CustomLogger): async def async_post_call_success_deployment_hook( self, request_data: dict, - response: Any, + response: LLMResponseTypes, call_type: CallTypes | None, - ) -> Any | None: + ) -> LLMResponseTypes | None: """ Post-call hook to handle automatic code execution. @@ -372,7 +376,7 @@ class SkillsInjectionHook(CustomLogger): # Check if any tool call needs execution (litellm_code_execution or skill tool) has_executable_tool = False for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX): has_executable_tool = True @@ -441,7 +445,7 @@ class SkillsInjectionHook(CustomLogger): data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes | None: """ Execute the code execution loop for messages API (Anthropic format). @@ -466,7 +470,7 @@ class SkillsInjectionHook(CustomLogger): max_tokens: Final = data.get("max_tokens", 4096) executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response = response for iteration in range(self.max_iterations): @@ -511,9 +515,9 @@ class SkillsInjectionHook(CustomLogger): # Process tool calls tool_results = [] for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") tool_id = tc.get("id", "") - tool_input = tc.get("input", {}) + tool_input: Mapping[str, str] = tc.get("input", {}) # Execute if it's litellm_code_execution OR a skill tool if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: @@ -561,8 +565,8 @@ class SkillsInjectionHook(CustomLogger): self, code: str, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute code in sandbox and return result string.""" try: @@ -574,7 +578,8 @@ class SkillsInjectionHook(CustomLogger): # Collect generated files if exec_result.get("files"): - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: generated_files.append( { "name": f["name"], @@ -595,10 +600,10 @@ class SkillsInjectionHook(CustomLogger): async def _execute_skill_tool( self, tool_name: str, - tool_input: dict[str, Any], + tool_input: Mapping[str, str], skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a skill tool by generating and running code based on skill content.""" # Generate code based on available skill modules @@ -670,7 +675,7 @@ print('No executable skill module found') data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes: """ Execute the code execution loop until model gives final response. @@ -704,7 +709,7 @@ print('No executable skill module found') kwargs: Final = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS} executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response: Any = response for iteration in range(self.max_iterations): @@ -713,7 +718,7 @@ print('No executable skill module found') stop_reason = current_response.choices[0].finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -781,13 +786,13 @@ print('No executable skill module found') self, tool_call: Any, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a litellm_code_execution tool call and return result string.""" try: args: Final = json.loads(tool_call.function.arguments) - code: Final = args.get("code", "") + code: Final[str] = args.get("code", "") verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code)) @@ -802,7 +807,8 @@ print('No executable skill module found') # Collect generated files if exec_result.get("files"): tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: file_content = base64.b64decode(f["content_base64"]) generated_files.append( { @@ -830,8 +836,8 @@ print('No executable skill module found') def _attach_files_to_response( self, response: Any, - generated_files: list[dict[str, Any]], - ) -> Any: + generated_files: list[dict[str, object]], + ) -> LLMResponseTypes: """ Attach generated files to the response object. @@ -841,11 +847,13 @@ print('No executable skill module found') if not generated_files: return response + raw_response: Final = response + # Handle dict response (Anthropic/messages API format) if isinstance(response, dict): response["_litellm_generated_files"] = generated_files verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to dict response", len(generated_files)) - return response + return raw_response # Handle object response (OpenAI format) try: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f2cb1124fa0..2de1d177b33 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,7 +18,7 @@ import os import re import secrets import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast @@ -171,8 +171,12 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ... + async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ... + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def update( self, *, @@ -181,6 +185,10 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): ) -> _PrismaRowT | None: ... +class _TxTables(Protocol): + litellm_proxymodeltable: _PrismaTableActions[object] + + def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: @@ -1650,9 +1658,12 @@ async def generate_key_fn( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( + user_custom_key_generate + ) + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1847,9 +1858,10 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1918,7 +1930,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) try: for k, v in data_json.items(): @@ -2179,7 +2191,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2722,9 +2734,10 @@ async def update_key_fn( ) # Custom key update hook - if user_custom_key_update is not None: - if inspect.iscoroutinefunction(user_custom_key_update): - result: Final = await user_custom_key_update(data) + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + if custom_key_update_hook is not None: + if inspect.iscoroutinefunction(custom_key_update_hook): + result: Final = await custom_key_update_hook(data) else: raise ValueError("user_custom_key_update must be a coroutine") decision: Final = result.get("decision", True) @@ -4089,10 +4102,11 @@ async def delete_verification_tokens( failed_tokens: list = [] try: if prisma_client: - tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"token": {"in": tokens}}) + hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens] + tokens = hashed_tokens + _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table( + VerificationTokenRepository(prisma_client) + ).find_many(where={"token": {"in": hashed_tokens}}) if len(_keys_being_deleted) == 0: raise HTTPException( @@ -4291,7 +4305,7 @@ async def _rotate_master_key( if models: decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models) verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models)) - new_models: Final = [] + new_models: Final[list[dict[str, object]]] = [] for model in decrypted_models: new_model = await _add_model_to_db( model_params=Deployment(**model), @@ -4306,7 +4320,8 @@ async def _rotate_master_key( _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxTables] = tx_ctx await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4630,7 +4645,7 @@ async def _execute_virtual_key_regeneration( _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) + jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( @@ -4642,9 +4657,9 @@ async def _execute_virtual_key_regeneration( updated_token: Final = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, - data=update_data, + data=jsonified_update_data, ) - updated_token_dict: Final = dict(updated_token) if updated_token is not None else {} + updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5589,7 +5604,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5602,7 +5617,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5695,7 +5710,7 @@ def _build_key_filter_conditions( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, -) -> dict[str, str | dict[str, Any] | list[dict[str, Any]]]: +) -> Mapping[str, object]: """Build filter conditions for key listing. Visibility rules: @@ -5707,14 +5722,14 @@ def _build_key_filter_conditions( so former members cannot see service accounts they created after leaving. """ # Prepare filter conditions - where: dict[str, str | dict[str, Any] | list[dict[str, Any]]] = {} + where: dict[str, object] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) # Build the OR conditions for user's keys and admin team keys - or_conditions: Final[list[dict[str, Any]]] = [] + or_conditions: Final[list[dict[str, object]]] = [] # Base conditions for user's own keys - user_condition: Final[dict[str, Any]] = {} + user_condition: Final[dict[str, object]] = {} if user_id and isinstance(user_id, str): if use_substring_matching: user_condition["user_id"] = { @@ -5784,7 +5799,7 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) - global_filters: tuple[dict[str, Any], ...] = ( + global_filters: Final[tuple[dict[str, object], ...]] = ( *( ( {"key_alias": {"contains": key_alias, "mode": "insensitive"}} @@ -5805,7 +5820,7 @@ def _build_key_filter_conditions( else () ), ) - combined_where = {"AND": [where, *global_filters]} if global_filters else where + combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where verbose_proxy_logger.debug("Filter conditions: %s", combined_where) return combined_where @@ -5986,7 +6001,7 @@ async def _list_key_helper( ) -def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, Any]: +def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]: """ Condition to filter out UI session tokens """ @@ -6395,7 +6410,7 @@ async def _can_user_query_key_info( async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, - key_logging: list[dict[str, Any]], + key_logging: Sequence[Mapping[str, str]], ) -> LoggingCallbackStatus: """ Test the key-based logging diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 71407c89813..c2087005863 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -78,9 +78,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, - DeploymentTypedDict, GenericLiteLLMParams, - LiteLLMParamsTypedDict, updateDeployment, ) from litellm.utils import get_utc_datetime @@ -104,10 +102,80 @@ class UpdatePublicModelGroupsRequest(BaseModel): model_config = ConfigDict(extra="forbid") +class _ProxyModelRow(Protocol): + model_id: str + model_name: str + model_info: Mapping[str, object] | None + + def model_dump_json(self, *, exclude_none: bool = False) -> str: ... + + +class _ProxyModelTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + + def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ... + + +class _TxModelTables(Protocol): + litellm_proxymodeltable: _ProxyModelTable + + +class _TeamRow(Protocol): + models: Sequence[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _TeamTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ... + + def update( + self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool] + ) -> Awaitable[LiteLLM_TeamTable]: ... + + +class _TeamIdRef(Protocol): + team_id: str + + +class _ModelAliasRow(Protocol): + id: int + model_aliases: dict[str, str] + team: _TeamIdRef | None + + +class _ModelAliasTable(Protocol): + def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ... + + +def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable: + return ModelRepository(prisma_client).table + + +def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable: + return TeamRepository(prisma_client).table + + +def _db_team_table(prisma_client: PrismaClient) -> _TeamTable: + return prisma_client.db.litellm_teamtable + + +def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable: + return ModelTableRepository(prisma_client).table + + async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None: db_model: Final = cast( BaseModel | None, - await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), + await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}), ) if not db_model: @@ -166,14 +234,9 @@ def _raise_on_strategy_router_write_violation( def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: - merged_deployment_dict: Final = DeploymentTypedDict( - model_name=db_model.model_name, - litellm_params=LiteLLMParamsTypedDict(**db_model.litellm_params.model_dump(exclude_none=True)), - model_info=db_model.model_info.model_dump(exclude_none=True), - ) - # update model name - if updated_patch.model_name: - merged_deployment_dict["model_name"] = updated_patch.model_name + merged_model_name: Final = updated_patch.model_name or db_model.model_name + merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) + merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) # update litellm params if updated_patch.litellm_params: @@ -182,13 +245,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } - merged_deployment_dict["litellm_params"].update(encrypted_params) + merged_litellm_params.update(encrypted_params) # update model info if updated_patch.model_info: - if "model_info" not in merged_deployment_dict: - merged_deployment_dict["model_info"] = {} - merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True)) + merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI # passes through (which today re-sends the OLD pricing on every save) cannot @@ -202,29 +263,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.litellm_params: for field in updated_patch.litellm_params.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_deployment_dict["litellm_params"].pop(field, None) - merged_deployment_dict.get("model_info", {}).pop(field, None) + merged_litellm_params.pop(field, None) + merged_model_info.pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_deployment_dict["model_info"].pop(field, None) - merged_deployment_dict.get("litellm_params", {}).pop(field, None) + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) # convert to prisma compatible format - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel() - if "model_name" in merged_deployment_dict: - prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"] + for key, value in merged_model_info.items(): + if isinstance(value, datetime.datetime): + merged_model_info[key] = value.isoformat() - if "litellm_params" in merged_deployment_dict: - prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"]) - - if "model_info" in merged_deployment_dict: - model_info: Final = merged_deployment_dict["model_info"] - for key, value in model_info.items(): - if isinstance(value, datetime.datetime): - model_info[key] = value.isoformat() - prisma_compatible_model_dict["model_info"] = json.dumps(model_info) + prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( + model_name=merged_model_name, + litellm_params=json.dumps(merged_litellm_params), + model_info=json.dumps(merged_model_info), + ) if updated_patch.blocked is not None: prisma_compatible_model_dict["blocked"] = updated_patch.blocked @@ -338,7 +395,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model: Final = await ModelRepository(prisma_client).table.update( + updated_model: Final = await _proxy_model_table(prisma_client).update( where={"model_id": model_id}, data=update_data, ) @@ -769,8 +826,8 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None -) -> list[LiteLLM_ProxyModelTable]: + team_id: str, prisma_client: PrismaClient, table: _ProxyModelTable | None = None +) -> Sequence[_ProxyModelRow]: """ Fetch all deployments for a given team_id from the database. @@ -785,7 +842,7 @@ async def _get_team_deployments( existing transaction. """ prefix: Final = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table + table = table or _proxy_model_table(prisma_client) response: Final = await table.find_many( where={ "model_name": {"startswith": prefix}, @@ -806,7 +863,7 @@ async def _get_team_deployments( async def delete_team_models( team_ids: list[str], prisma_client: PrismaClient, - llm_router: Any | None, + llm_router: Router | None, ) -> list[str]: """ Delete every BYOK model owned by the given teams, from the DB and the router. @@ -820,7 +877,8 @@ async def delete_team_models( Returns the model_ids that were deleted. """ deleted_model_ids: Final[list[str]] = [] - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxModelTables] = tx_ctx for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] @@ -920,11 +978,11 @@ async def _remove_unbacked_team_models( if not names_to_remove: return - existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + existing_team_row: Final = await _db_team_table(prisma_client).find_unique(where={"team_id": team_id}) if existing_team_row is None: return - updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update( + updated_team_row: Final[LiteLLM_TeamTable] = await _db_team_table(prisma_client).update( where={"team_id": team_id}, data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, @@ -953,7 +1011,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: dict | str | None, + model_info: object, ) -> str | None: parsed: Final = model_info_as_mapping(model_info) if parsed is None: @@ -1062,7 +1120,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1091,7 +1149,7 @@ class ModelManagementAuthChecks: ) -> Literal[True]: ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: - team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique( + team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1192,7 +1250,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) + result: Final = await _proxy_model_table(prisma_client).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1265,9 +1323,9 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) + team_model_aliases: Final = await _model_alias_table(prisma_client).find_many(include={"team": True}) tasks: Final = [] - removed_model_aliases: Final = [] + removed_model_aliases: Final[list[tuple[str, str]]] = [] for team_model_alias in team_model_aliases: model_aliases = team_model_alias.model_aliases # {"alias": "public model name"} id = team_model_alias.id @@ -1278,7 +1336,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - ModelTableRepository(prisma_client).table.update( + _model_alias_table(prisma_client).update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1492,7 +1550,7 @@ async def update_model( }, ) - _model_id = None + _model_id: str | None = None _model_info: Final = getattr(model_params, "model_info", None) if _model_info is None: raise Exception("model_info not provided") @@ -1551,11 +1609,11 @@ async def update_model( else: pass - _data: Final[dict] = { + _data: Final[dict[str, str]] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response: Final = await ModelRepository(prisma_client).table.update( + model_response: Final = await _proxy_model_table(prisma_client).update( where={"model_id": _model_id}, data=_data, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0b99879f9fe..97f494c51de 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,11 +15,11 @@ import math import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Annotated, Final, Protocol, TypeVar, cast +from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm._logging import verbose_proxy_logger @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetNewRequest, CommonProxyErrors, DeleteTeamRequest, LiteLLM_AccessGroupTable, @@ -156,6 +157,15 @@ router: Final = APIRouter() _DbRecordT = TypeVar("_DbRecordT") +class _TeamIdKeyCount(TypedDict): + team_id: int + + +class _TeamIdGroupRow(TypedDict): + team_id: str + _count: _TeamIdKeyCount + + class _PrismaTableActions(Protocol[_DbRecordT]): async def find_unique( self, @@ -220,59 +230,127 @@ class _PrismaTableActions(Protocol[_DbRecordT]): where: Mapping[str, object] | None = None, ) -> int: ... + async def group_by( + self, + by: Sequence[str], + where: Mapping[str, object] | None = None, + count: Mapping[str, bool] | None = None, + ) -> Sequence[_TeamIdGroupRow]: ... + + +class _HasTableActions(Protocol[_DbRecordT]): + @property + def table(self) -> "_PrismaTableActions[_DbRecordT]": ... + + +def _typed_table( + repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT] +) -> "_PrismaTableActions[_DbRecordT]": + return repo.table + + +def _as_object(value: object) -> object: + return value + + +def _nullable(value: _DbRecordT | None) -> _DbRecordT | None: + return value + + +class _UserIdRow(Protocol): + @property + def user_id(self) -> str | None: ... + + +class _HasUserIdTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserIdRow]": ... + + +def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]": + return repo.table + + +class _RawTeamRow(Protocol): + @property + def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ... + + +class _HasRawTeamTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_RawTeamRow]": ... + + +def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]": + return repo.table + + +class _BudgetWriteCall(Protocol): + async def __call__( + self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth + ) -> LiteLLM_BudgetTableFull: ... + + +def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall": + return fn + + +class _TeamFindManyArgs(TypedDict, total=False): + take: int + skip: int + order: Mapping[str, str] + cursor: Mapping[str, object] + + +class _TeamUiViewFilters(TypedDict, total=False): + team_id: Mapping[str, str] + team_alias: Mapping[str, str] + + +class _TeamIdInFilter(TypedDict, total=False): + team_id: Mapping[str, Sequence[str]] + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": - team_table: Final[_PrismaTableActions[LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table + return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]": - membership_table: Final[_PrismaTableActions[LiteLLM_TeamMembership]] = TeamMembershipRepository(prisma_client).table - return membership_table + return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership) def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]": - user_table: Final[_PrismaTableActions[LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table + return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable) def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]": - model_table: Final[_PrismaTableActions[LiteLLM_ModelTable]] = ModelTableRepository(prisma_client).table - return model_table + return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable) def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]": - org_table: Final[_PrismaTableActions[LiteLLM_OrganizationTable]] = OrganizationRepository(prisma_client).table - return org_table + return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable) def _org_membership_db( prisma_client: PrismaClient | None, ) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]": - org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository( - prisma_client - ).table - return org_membership_table + return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable) def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]": - budget_table: Final[_PrismaTableActions[LiteLLM_BudgetTableFull]] = BudgetRepository(prisma_client).table - return budget_table + return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull) def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]": - deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table - return deleted_team_table + return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable) def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]": - access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table - return access_group_table + return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable) def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]": - tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table - return tokens_table + return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken) def _sanitize_for_log(value: object) -> str: @@ -408,7 +486,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table: Final = await new_budget( + team_member_budget_table: Final = await _as_budget_write(new_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -456,7 +534,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row: Final = await update_budget( + budget_row: Final = await _as_budget_write(update_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -571,7 +649,7 @@ class TeamMemberBudgetHandler: ) if missing: - await TeamMembershipRepository(prisma_client).table.create_many( + await _team_membership_db(prisma_client).create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -1407,9 +1485,10 @@ async def new_team( complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"]) complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) + team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict - team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create( - data=complete_team_data_dict, + team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( + data=team_creation_data, include={"litellm_model_table": True}, ) @@ -1856,7 +1935,7 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -1884,7 +1963,7 @@ async def update_team( ) if data.max_budget is not None: - existing_soft_budget: Final = getattr(existing_team_row, "soft_budget", None) + existing_soft_budget: Final[object] = _as_object(getattr(existing_team_row, "soft_budget", None)) soft_budget_to_check: Final = data.soft_budget if data.soft_budget is not None else existing_soft_budget if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): if data.max_budget <= soft_budget_to_check: @@ -1943,7 +2022,7 @@ async def update_team( data.organization_id = None # check org team limits - if updating team that belongs to an org - org_id_to_check: Final = ( + org_id_to_check: Final[object] = _as_object( data.organization_id if data.organization_id is not None else existing_team_row.organization_id ) if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: @@ -1976,7 +2055,7 @@ async def update_team( TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) if "metadata" in updated_kv: - stored_metadata: Final = ( + stored_metadata: Final[Mapping[str, JsonValue] | None] = ( { # mutable-ok: the validator payload's isinstance guard requires a plain dict key: value for key, value in existing_team_row.metadata.items() @@ -2079,16 +2158,19 @@ async def update_team( updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Final[LiteLLM_TeamTable | None] = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data=updated_kv, - # `object_permission` is included so `_refresh_cached_team` - # doesn't write a cached team with the relation nulled out — - # see team_model_add for the full rationale. - include={ - "litellm_model_table": True, - "object_permission": True, - }, + team_update_data: Final[Mapping[str, object]] = updated_kv + team_row: Final[LiteLLM_TeamTable | None] = _nullable( + await _team_db(prisma_client).update( + where={"team_id": data.team_id}, + data=team_update_data, + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, + ) ) if team_row is None or team_row.team_id is None: @@ -2603,7 +2685,7 @@ async def _resolve_existing_member_user_ids( if not requested_user_ids: return frozenset() - found: Final = await UserRepository(prisma_client).table.find_many( + found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many( where={ # mutable-ok: Prisma query filters are dict-shaped "user_id": { # mutable-ok: Prisma query filters are dict-shaped "in": sorted(requested_user_ids) @@ -3098,7 +3180,9 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val) + existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many( + where=key_val + ) if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): for existing_user in existing_user_rows: @@ -3106,7 +3190,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await UserRepository(prisma_client).table.update( + await _user_db(prisma_client).update( where={ "user_id": existing_user.user_id, }, @@ -3114,7 +3198,7 @@ async def team_member_delete( ) # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = set() + user_ids_to_delete: Final = set[str]() if data.user_id is not None: user_ids_to_delete.add(data.user_id) if existing_user_rows is not None and isinstance(existing_user_rows, list): @@ -3123,9 +3207,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"team_id": data.team_id, "user_id": _uid} - ) + await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: @@ -3134,9 +3216,7 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many( + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3151,7 +3231,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await VerificationTokenRepository(prisma_client).table.delete_many( + await _tokens_db(prisma_client).delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3311,7 +3391,7 @@ async def team_member_update( ### upsert new budget budget_patch: Final = _build_member_budget_patch(data) - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3654,7 +3734,7 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many( where={"team_id": {"in": data.team_ids}} ) @@ -4469,7 +4549,7 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by( + grouped: Final = await _tokens_db(prisma_client).group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, @@ -4786,7 +4866,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams: Final = await TeamRepository(prisma_client).table.find_many( + org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4800,7 +4880,9 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) + response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( + include={"litellm_model_table": True} + ) return [ team for team in response @@ -4808,7 +4890,7 @@ async def _authorize_and_filter_teams( ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})) @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -4860,7 +4942,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) + keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( @@ -4911,7 +4993,7 @@ async def get_paginated_teams( total_count: Final = await _team_db(prisma_client).count() # Get paginated teams - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4961,7 +5043,7 @@ async def ui_view_teams( skip: Final = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Final = {} + where_conditions: Final[_TeamUiViewFilters] = {} if team_id: where_conditions["team_id"] = { @@ -4976,7 +5058,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -5166,13 +5248,13 @@ async def team_model_delete( ) # Get current models list - current_models: Final = team_obj.models or [] + current_models: Final[Sequence[str]] = team_obj.models or [] # Remove specified models updated_models: Final = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team: Final = await TeamRepository(prisma_client).table.update( + updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, @@ -5425,7 +5507,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi BATCH_SIZE: Final = 500 while True: - find_args: dict = { + find_args: _TeamFindManyArgs = { "take": BATCH_SIZE, "order": {"team_id": "asc"}, } @@ -5433,7 +5515,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams = await _team_db(prisma_client).find_many(**find_args) if not teams: break @@ -5528,11 +5610,11 @@ async def get_team_daily_activity( ) ## Fetch team aliases and check team admin status - where_condition: Final = {} + where_condition: Final[_TeamIdInFilter] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 44abc56713f..a2c50590dd5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,9 +16,22 @@ import json import os import re import secrets +from collections.abc import Mapping, Sequence from copy import deepcopy from html import escape -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: @@ -155,6 +168,102 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset( } ) +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique( + self, + where: Mapping[str, object], + ) -> _DbRecordT | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + ) -> _DbRecordT | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + ) -> Sequence[_DbRecordT]: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _DbRecordT: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + +class _UserMetadataRow(Protocol): + @property + def metadata(self) -> Mapping[str, object] | None: ... + + +class _HasUserMetadataTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ... + + +def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]": + return repo.table + + +class _SsoConfigRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoConfigTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ... + + +def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]": + return repo.table + + +class _TeamDetailRow(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _HasTeamDetailTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ... + + +def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]": + return repo.table + + +class _CustomSsoCall(Protocol): + async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... + + +class _ServicePrincipalAssignment(Protocol): + def get(self, key: str) -> str: ... + + +class _ServicePrincipalPage(Protocol): + @overload + def get( + self, + key: Literal["value"], + default: Sequence["_ServicePrincipalAssignment"], + ) -> Sequence["_ServicePrincipalAssignment"]: ... + + @overload + def get(self, key: Literal["@odata.nextLink"]) -> str | None: ... + + +def _as_object(value: object) -> object: + return value + def _hash_cli_sso_secret(secret: str) -> str: return hashlib.sha256(secret.encode("utf-8")).hexdigest() @@ -256,7 +365,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: flow = cache.get_cache(key=cache_key) if isinstance(flow, str): try: - flow = json.loads(flow) + flow = _as_object(json.loads(flow)) except ValueError: flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: @@ -421,7 +530,7 @@ def _flatten_cli_sso_metadata_for_poll( def build_cli_sso_attribution_metadata( result: CustomOpenID | OpenID | dict, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build allowlisted, non-secret scalar attribution metadata from an SSO result. @@ -432,7 +541,7 @@ def build_cli_sso_attribution_metadata( if not claim_map: return {} - metadata: Final[dict[str, Any]] = {} + metadata: Final[dict[str, object]] = {} for source_claim, dest_key in claim_map: if not _is_safe_cli_sso_metadata_dest_key(dest_key): verbose_proxy_logger.debug("Skipping unsafe CLI SSO metadata destination key: %s", dest_key) @@ -474,14 +583,14 @@ def _merge_cli_sso_attribution_metadata( async def _persist_cli_sso_user_metadata( prisma_client: PrismaClient, user_id: str, - attribution_metadata: dict[str, Any], + attribution_metadata: dict[str, object], ) -> None: if not attribution_metadata: return try: - user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - existing_metadata: dict[str, Any] = {} + user_row: Final = await _user_meta_db(UserRepository(prisma_client)).find_unique(where={"user_id": user_id}) + existing_metadata: dict[str, object] = {} if user_row is not None: row_metadata: Final = user_row.metadata if isinstance(row_metadata, dict): @@ -491,7 +600,7 @@ async def _persist_cli_sso_user_metadata( existing_metadata=existing_metadata, attribution_metadata=attribution_metadata, ) - await UserRepository(prisma_client).table.update_many( + await _user_meta_db(UserRepository(prisma_client)).update_many( where={"user_id": user_id}, data={"metadata": merged_metadata}, ) @@ -1104,7 +1213,7 @@ def generic_response_convertor( ) # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified - extra_fields: dict[str, Any] | None = None + extra_fields: dict[str, object] | None = None if generic_user_extra_attributes: extra_fields = {} for attr_name in generic_user_extra_attributes.split(","): @@ -1193,7 +1302,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1225,7 +1336,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1273,7 +1386,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: return role_mappings -def _parse_generic_sso_headers() -> dict: +def _parse_generic_sso_headers() -> dict[str, str]: """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict.""" raw: Final = os.getenv("GENERIC_SSO_HEADERS", None) if raw is None: @@ -1677,7 +1790,7 @@ def _build_sso_user_update_data( result: Union["CustomOpenID", OpenID, dict] | None, user_email: str | None, user_id: str | None, -) -> dict: +) -> dict[str, object]: """ Build the update data dictionary for SSO user upsert. @@ -1689,7 +1802,7 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid """ - update_data: Final[dict] = {"user_email": normalize_email(user_email)} + update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid sso_role: Final = getattr(result, "user_role", None) @@ -1740,7 +1853,7 @@ async def _sync_user_role_from_jwt_role_map( # Update existing DB record if role differs if user_info is not None and user_info.user_role != mapped_role.value: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_info.user_id}, data={"user_role": mapped_role.value}, ) @@ -1796,7 +1909,7 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism return user_role if prisma_client: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_id}, data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) @@ -2016,10 +2129,11 @@ async def _build_cli_sso_user_defined_values( ) -> SSOUserDefinedValues | None: from litellm.proxy.proxy_server import user_custom_sso + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso user_id: Final = parsed_openid_result.get("user_id") - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - return await user_custom_sso(result) + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + return await custom_sso_handler(result) raise ValueError("user_custom_sso must be a coroutine function") if user_id is None: return None @@ -2035,12 +2149,14 @@ async def _build_cli_sso_user_defined_values( async def _fetch_cli_sso_team_details( prisma_client: PrismaClient, - teams: list[str], -) -> list[dict[str, Any]]: - team_details: Final[list[dict[str, Any]]] = [] + teams: Sequence[str], +) -> list[dict[str, object]]: + team_details: Final[list[dict[str, object]]] = [] try: if teams: - prisma_teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}}) + prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( + where={"team_id": {"in": teams}} + ) for team_row in prisma_teams: team_dict = team_row.model_dump() team_details.append( @@ -2257,12 +2373,12 @@ async def cli_poll_key( verbose_proxy_logger.info("Returning teams list for user %s to select from: %s", user_id, user_teams) # Best-effort construction of team_details if it wasn't # already cached for some reason. - team_details_response: list[dict[str, Any]] | None = None + team_details_response: list[dict[str, object]] | None = None if isinstance(user_team_details, list) and user_team_details: team_details_response = user_team_details elif user_teams: team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams] - poll_response: dict[str, Any] = { + poll_response: dict[str, object] = { "status": "ready", "user_id": user_id, "teams": user_teams, @@ -2997,7 +3113,9 @@ class SSOAuthenticationHandler: user_id=user_id, ) - await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data) + await _user_meta_db(UserRepository(prisma_client)).update_many( + where={"user_id": user_id}, data=update_data + ) else: verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") # user not in DB, insert User into LiteLLM DB @@ -3089,7 +3207,9 @@ class SSOAuthenticationHandler: code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj: Final = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id}) + team_obj: Final = await _team_detail_db(TeamRepository(prisma_client)).find_first( + where={"team_id": litellm_team_id} + ) verbose_proxy_logger.debug("Team object: %s", team_obj) # only create a new team if it doesn't exist @@ -3278,9 +3398,10 @@ class SSOAuthenticationHandler: # But if it is, we want their models preferences user_defined_values: SSOUserDefinedValues | None = None - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - user_defined_values = await user_custom_sso(result) + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + user_defined_values = await custom_sso_handler(result) else: raise ValueError("user_custom_sso must be a coroutine function") elif user_id is not None: @@ -3448,7 +3569,7 @@ class SSOAuthenticationHandler: dict: Token exchange parameters """ # Prepare token exchange parameters (may add code_verifier: str later) - token_params: Final[dict[str, Any]] = {"include_client_id": generic_include_client_id} + token_params: Final[dict[str, object]] = {"include_client_id": generic_include_client_id} # Retrieve PKCE code_verifier if PKCE was used in authorization. # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip @@ -3663,7 +3784,7 @@ class SSOAuthenticationHandler: access_token string. Raises ProxyException on any validation failure. """ try: - token_response_raw: Final = response.json() + token_response_raw: Final[object] = _as_object(response.json()) except Exception as json_err: verbose_proxy_logger.error( "Failed to parse token response as JSON: %s. Body: %s", @@ -4253,7 +4374,7 @@ class MicrosoftSSOHandler: while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: response = await async_client.get(next_link, headers=headers) - response_json = response.json() + response_json: _ServicePrincipalPage = response.json() verbose_proxy_logger.debug("Response from service principal app role assigned to: %s", response_json) for _object in response_json.get("value", []): diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 382df608a0c..08bb8698cac 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -4,11 +4,11 @@ import json import os from collections import Counter from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol, TypeVar from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile -from pydantic import ConfigDict, ValidationError, create_model +from pydantic import ConfigDict, JsonValue, ValidationError, create_model from pydantic.fields import FieldInfo import litellm @@ -36,6 +36,73 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router: Final = APIRouter() +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + +class _SsoSettingsMappingRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoSettingsMappingTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ... + + +def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]: + return repo.table + + +class _StoredSsoSettingsRow(Protocol): + @property + def sso_settings(self) -> object: ... + + +class _HasStoredSsoSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ... + + +def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]: + return repo.table + + +class _UiSettingsRow(Protocol): + @property + def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _HasUiSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_UiSettingsRow]: ... + + +def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]: + return repo.table + + +class _ConfigParamRow(Protocol): + @property + def param_value(self) -> str | Mapping[str, object] | None: ... + + +class _HasConfigParamTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_ConfigParamRow]: ... + + +def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]: + return repo.table + + # Maps each UIThemeConfig field to the env var the UI branding path reads it # from. /update/ui_theme_settings writes both the stored ui_theme_config and # these env vars, so /get/ui_theme_settings resolves the same env vars to @@ -54,7 +121,7 @@ def _is_public_http_url(value: str | None) -> bool: return parsed.scheme in ("http", "https") and bool(parsed.netloc) -def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None: +def _resolve_ui_theme_field(stored_values: Mapping[str, object], field_name: str) -> str | None: """Resolve one UI theme field to the value the branding path actually uses. The stored ui_theme_config wins; a field absent or blank there falls back to @@ -263,7 +330,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ # include generics like ``Optional[int]`` / ``List[str]`` that are not # instances of ``type`` — so tightening this to ``type`` would reject # valid inputs. -_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[Any, FieldInfo]]] = {} +_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[object, FieldInfo]]] = {} # Settings OSS knows about as enterprise-gated. If a caller sends one of # these keys and no extension package has registered it, the PATCH @@ -275,7 +342,7 @@ _ENTERPRISE_ONLY_UI_SETTINGS: Final[set[str]] = {"enable_projects_ui"} _EFFECTIVE_UI_SETTINGS_CLASS: type[UISettings] | None = None -def register_extra_ui_setting(name: str, annotation: Any, field: FieldInfo) -> None: +def register_extra_ui_setting(name: str, annotation: object, field: FieldInfo) -> None: """Register an additional UI settings field contributed by an extension package. ``field`` must be a ``FieldInfo`` instance — construct it directly @@ -470,7 +537,7 @@ async def delete_allowed_ip( async def _get_settings_with_schema( settings_key: str, - settings_class: Any, + settings_class: type[BaseModel], config: dict, ) -> dict: """ @@ -842,7 +909,9 @@ async def get_sso_settings(): # Resolve the effective SSO config: the stored row wins, else the process # environment, else each field's default. Unlike the legacy read path this # does not write os.environ; a GET has no business mutating the environment. - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_settings_mapping_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) sso_db_settings: Final = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None resolved: Final = resolve_sso_config(sso_db_settings, os.environ) @@ -914,8 +983,10 @@ async def update_sso_settings( # before-snapshot has the same shape as after_value, and rely on # create_config_audit_log's secret-name redaction to mask the # *_client_secret fields before the audit row is written. - existing_sso_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - before_sso_data: dict[str, Any] | None = None + existing_sso_record: Final = await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) + before_sso_data: dict[str, JsonValue] | None = None if existing_sso_record and existing_sso_record.sso_settings: stored = existing_sso_record.sso_settings if isinstance(stored, str): @@ -948,7 +1019,7 @@ async def update_sso_settings( encrypted_sso_data: Final = proxy_config._encrypt_env_variables(environment_variables=sso_data) # Save to dedicated SSO table - await SSOConfigRepository(prisma_client).table.upsert( + await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).upsert( where={"id": "sso_config"}, data={ "create": { @@ -974,7 +1045,7 @@ async def update_sso_settings( # Remove SSO-related env vars from config.environment_variables try: - env_var_entry: Final = await ConfigRepository(prisma_client).table.find_unique( + env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique( where={"param_name": "environment_variables"} ) @@ -982,7 +1053,7 @@ async def update_sso_settings( if env_var_entry is not None: if env_var_entry.param_value is not None: if isinstance(env_var_entry.param_value, str): - environment_variables = json.loads(env_var_entry.param_value) + environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) else: environment_variables = dict(env_var_entry.param_value) else: @@ -993,7 +1064,7 @@ async def update_sso_settings( key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } - await ConfigRepository(prisma_client).table.update( + await _config_param_db(ConfigRepository(prisma_client)).update( where={"param_name": "environment_variables"}, data={ "param_value": json.dumps(filtered_env_vars, default=str), @@ -1239,8 +1310,10 @@ async def get_ui_settings_cached() -> dict[str, Any]: if prisma_client is None: return {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) - ui_settings: dict[str, Any] = {} + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) + ui_settings: dict[str, JsonValue] = {} if db_record and db_record.ui_settings: raw: Final = db_record.ui_settings ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) @@ -1272,9 +1345,11 @@ async def get_ui_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - ui_settings: dict[str, Any] = {} + ui_settings: Mapping[str, JsonValue] = {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_record and db_record.ui_settings: ui_settings_json: Final = db_record.ui_settings @@ -1300,7 +1375,7 @@ async def get_ui_settings(): await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) # Build config-like object for schema helper - config: Final[dict[str, Any]] = {"litellm_settings": {"ui_settings": ui_settings}} + config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} return await _get_settings_with_schema( settings_key="ui_settings", @@ -1315,7 +1390,7 @@ async def get_ui_settings(): dependencies=[Depends(user_api_key_auth)], ) async def update_ui_settings( - settings_body: dict[str, Any] = Body(...), + settings_body: dict[str, object] = Body(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1352,7 +1427,7 @@ async def update_ui_settings( raise HTTPException(status_code=422, detail=e.errors()) # Only include fields the caller actually sent (not Pydantic defaults). - settings_dict: Final = settings.model_dump(exclude_unset=True) + settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True) # Reject enterprise-only settings up front so the caller gets a clear # signal instead of a silent drop. @@ -1373,15 +1448,17 @@ async def update_ui_settings( # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. - existing: dict = {} - db_existing: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + existing: dict[str, JsonValue] = {} + db_existing: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_existing and db_existing.ui_settings: raw: Final = db_existing.ui_settings existing = json.loads(raw) if isinstance(raw, str) else dict(raw) ui_settings: Final = {**existing, **incoming} - await UISettingsRepository(prisma_client).table.upsert( + await _ui_settings_db(UISettingsRepository(prisma_client)).upsert( where={"id": "ui_settings"}, data={ "create": { diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6b40281198c..2b037bef795 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,10 +10,17 @@ All /vector_store management endpoints import copy import json -from typing import Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +if TYPE_CHECKING: + from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow + + from litellm.proxy.utils import PrismaClient + from litellm.router import Router + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -43,6 +50,25 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() + +class _VectorStoreTableActions(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + +def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions: + return ManagedVectorStoresRepository(prisma_client).table + + +def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore(**row.model_dump()) + + _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker() @@ -117,22 +143,20 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An async def _fetch_and_authorize_vector_store( vector_store_id: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: "PrismaClient", ) -> "LiteLLM_ManagedVectorStore": """ Look up a vector store by id and confirm the caller can access it. Raises HTTPException(404) on miss and HTTPException(403) on access denial. """ - row: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( - where={"vector_store_id": vector_store_id} - ) + row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id}) if row is None: raise HTTPException( status_code=404, detail=f"Vector store with ID {vector_store_id} not found", ) - typed: Final = LiteLLM_ManagedVectorStore(**row.model_dump()) + typed: Final = _row_to_vector_store(row) if not await _check_vector_store_access(typed, user_api_key_dict): raise HTTPException( status_code=403, @@ -141,7 +165,7 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, Any] | None: +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: """ Resolve embedding config from router's config-defined models. @@ -177,7 +201,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d litellm_params = deployment.litellm_params # Build embedding config from model params - embedding_config: dict[str, Any] = {} + embedding_config: dict[str, object] = {} # Extract api_key api_key = getattr(litellm_params, "api_key", None) @@ -217,7 +241,9 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d return None -async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> dict[str, Any] | None: +async def _resolve_embedding_config_from_db( + embedding_model: str, prisma_client: "PrismaClient" +) -> dict[str, object] | None: """ Resolve embedding config from database model configuration. @@ -307,7 +333,9 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) return None -async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> dict[str, Any] | None: +async def _resolve_embedding_config( + embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None +) -> dict[str, object] | None: """ Resolve embedding config from either router (config-defined) or database models. @@ -388,7 +416,7 @@ async def _check_vector_store_access( async def create_vector_store_in_db( vector_store_id: str, custom_llm_provider: str, - prisma_client, + prisma_client: "PrismaClient | None", vector_store_name: str | None = None, vector_store_description: str | None = None, vector_store_metadata: dict | None = None, @@ -417,7 +445,7 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": vector_store_id} ) if existing_vector_store is not None: @@ -427,7 +455,7 @@ async def create_vector_store_in_db( ) # Prepare data for database - data_to_create: Final[dict[str, Any]] = { + data_to_create: Final[dict[str, object]] = { "vector_store_id": vector_store_id, "custom_llm_provider": custom_llm_provider, } @@ -463,9 +491,9 @@ async def create_vector_store_in_db( data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create) + _new_vector_store: Final = await _vector_store_table(prisma_client).create(data=data_to_create) - new_vector_store: Final[LiteLLM_ManagedVectorStore] = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump()) + new_vector_store: Final[LiteLLM_ManagedVectorStore] = _row_to_vector_store(_new_vector_store) # Add vector store to registry if litellm.vector_store_registry is not None: @@ -682,12 +710,12 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": data.vector_store_id} ) if existing_vector_store is not None: db_vector_store_exists = True - vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump()) + vector_store_to_check = _row_to_vector_store(existing_vector_store) # Check in-memory registry if litellm.vector_store_registry is not None: @@ -715,9 +743,7 @@ async def delete_vector_store( # Delete from database if exists if db_vector_store_exists: - await ManagedVectorStoresRepository(prisma_client).table.delete( - where={"vector_store_id": data.vector_store_id} - ) + await _vector_store_table(prisma_client).delete(where={"vector_store_id": data.vector_store_id}) # Delete from in-memory registry if exists if memory_vector_store_exists and litellm.vector_store_registry is not None: @@ -829,7 +855,7 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) - vector_store_id: Final = update_data.pop("vector_store_id") + vector_store_id: Final[str] = update_data.pop("vector_store_id") # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — @@ -857,12 +883,12 @@ async def update_vector_store( update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database - updated: Final = await ManagedVectorStoresRepository(prisma_client).table.update( + updated: Final = await _vector_store_table(prisma_client).update( where={"vector_store_id": vector_store_id}, data=update_data, ) - updated_vs: Final = LiteLLM_ManagedVectorStore(**updated.model_dump()) + updated_vs: Final = _row_to_vector_store(updated) # Immediately update in-memory registry to keep it in sync if litellm.vector_store_registry is not None: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..b2d065ea23b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Final, Literal, cast +from collections.abc import Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable from openai.types.chat.chat_completion_named_tool_choice_param import ( ChatCompletionNamedToolChoiceParam, @@ -16,6 +16,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_param import FunctionToolParam +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm._logging import verbose_logger @@ -78,9 +79,35 @@ from .custom_tools import ( unwrap_custom_tool_arguments, ) +if TYPE_CHECKING: + from openai.types.responses.response_apply_patch_tool_call import ( + ResponseApplyPatchToolCall, + ) + ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE: Final = InMemoryCache() +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _HasToolCalls(Protocol): + tool_calls: object + + +@runtime_checkable +class _HasId(Protocol): + id: object + class ChatCompletionSession(TypedDict, total=False): messages: list[ @@ -205,7 +232,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, stream: bool | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, **kwargs, ) -> dict: """ @@ -462,7 +489,9 @@ class LiteLLMCompletionResponsesConfig: if not chat_completion_messages: continue - deduped_in_place: list[Any] = [] + deduped_in_place: list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage + ] = [] for m in chat_completion_messages: role = "" if isinstance(m, dict): @@ -472,7 +501,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": - tool_calls: Any = ( + tool_calls: object = ( m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" @@ -534,7 +563,7 @@ class LiteLLMCompletionResponsesConfig: call_id = "" if role == "assistant": - tool_calls: Any = None + tool_calls: object = None if isinstance(tool_call_message, dict): tool_calls = tool_call_message.get("tool_calls") else: @@ -578,7 +607,16 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None: + def _find_previous_assistant_idx( + messages: Sequence[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message + ], + current_idx: int, + ) -> int | None: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": @@ -586,7 +624,18 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: + def _recover_tool_call_id_from_assistant( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -594,17 +643,23 @@ class LiteLLMCompletionResponsesConfig: else getattr(assistant_message, "tool_calls", None) ) if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: - first_tool_call: Final = tool_calls_raw[0] + first_tool_call: Final = _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)[0] if isinstance(first_tool_call, dict): - tool_call_id_raw = first_tool_call.get("id", "") + tool_call_id_raw = _ANY_KEY_DICT_ADAPTER.validate_python(first_tool_call).get("id", "") return str(tool_call_id_raw) if tool_call_id_raw is not None else "" - elif hasattr(first_tool_call, "id"): - tool_call_id_raw = getattr(first_tool_call, "id", None) + elif isinstance(first_tool_call, _HasId): + tool_call_id_raw = first_tool_call.id return str(tool_call_id_raw) if tool_call_id_raw is not None else "" return "" @staticmethod - def _get_tool_calls_list(assistant_message: Any) -> list[Any]: + def _get_tool_calls_list( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> Sequence[object]: """Extract tool_calls as a list from assistant message.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -614,18 +669,18 @@ class LiteLLMCompletionResponsesConfig: if tool_calls_raw is None: return [] if isinstance(tool_calls_raw, list): - return tool_calls_raw - if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): + return _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw) + if isinstance(tool_calls_raw, _SupportsIter) and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @staticmethod - def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: + def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool: """Check if a tool_call with the given ID exists in the list.""" for tool_call in tool_calls: - tool_call_id_to_check: str | None = None + tool_call_id_to_check: object = None if isinstance(tool_call, dict): - tool_call_id_to_check = tool_call.get("id") + tool_call_id_to_check = _ANY_KEY_DICT_ADAPTER.validate_python(tool_call).get("id") elif hasattr(tool_call, "id"): tool_call_id_to_check = getattr(tool_call, "id", None) if tool_call_id_to_check == tool_call_id: @@ -633,12 +688,13 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): - tool_function = tool.get("function") or {} - tool_name = tool_function.get("name") or tool.get("name") or "" + tool_map = _ANY_KEY_DICT_ADAPTER.validate_python(tool) + tool_function = _ANY_KEY_DICT_ADAPTER.validate_python(tool_map.get("function") or {}) + tool_name = tool_function.get("name") or tool_map.get("name") or "" if tool_name: return { "id": tool_call_id, @@ -651,7 +707,7 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object: """ Safely read a field from dict-like or attribute-based objects. """ @@ -659,7 +715,7 @@ class LiteLLMCompletionResponsesConfig: return default if isinstance(obj, dict): - return obj.get(key, default) + return _ANY_KEY_DICT_ADAPTER.validate_python(obj).get(key, default) getter: Final = getattr(obj, "get", None) if callable(getter): @@ -672,13 +728,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _create_tool_call_chunk( - tool_use_definition: dict[str, Any], tool_call_id: str, index: int + tool_use_definition: Mapping[object, object], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") function_name_raw: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") - function: Final[dict[str, Any]] = { + function: Final[dict[str, object]] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } @@ -697,7 +753,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None: + def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[object, object] | None: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -705,7 +761,7 @@ class LiteLLMCompletionResponsesConfig: return None if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + normalized_definition: dict[object, object] = _ANY_KEY_DICT_ADAPTER.validate_python(tool_use_definition) else: tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") @@ -738,7 +794,7 @@ class LiteLLMCompletionResponsesConfig: return normalized_definition @staticmethod - def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None: + def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None: """Add a tool_call to an assistant message.""" if isinstance(assistant_message, dict): prev_assistant_dict: Final = cast(dict[str, Any], assistant_message) @@ -747,7 +803,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls_list: Final = prev_assistant_dict["tool_calls"] if isinstance(tool_calls_list, list): tool_calls_list.append(tool_call_chunk) - elif hasattr(assistant_message, "tool_calls"): + elif isinstance(assistant_message, _HasToolCalls): if assistant_message.tool_calls is None: assistant_message.tool_calls = [] if isinstance(assistant_message.tool_calls, list): @@ -762,7 +818,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -851,7 +907,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant) if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id): - _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) + _tool_use_definition: object = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) if not _tool_use_definition and tools: _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( @@ -908,7 +964,7 @@ class LiteLLMCompletionResponsesConfig: function_call=input_item ) else: - content: Final = input_item.get("content") + content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: @@ -923,7 +979,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_tool_call_output(input_item: Any) -> bool: + def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a tool call output """ @@ -936,7 +992,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_function_call(input_item: Any) -> bool: + def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a function call or custom tool call. Both need to be reconstructed as assistant tool_calls for Chat @@ -946,7 +1002,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_tool_call_output_to_chat_completion_message( - tool_call_output: dict[str, Any], + tool_call_output: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ ChatCompletionToolMessage is used to indicate the output from a tool call @@ -958,7 +1014,7 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, + output: object, ) -> Any: """ Normalize Responses API function_call_output.output into a shape that downstream @@ -981,7 +1037,7 @@ class LiteLLMCompletionResponsesConfig: # Some adapters represent tool output as a list of "input_*" parts if isinstance(output, list): - normalized_blocks: Final[list[dict[str, Any]]] = [] + normalized_blocks: Final[list[dict[str, object]]] = [] text_acc: Final[list[str]] = [] for part in output: if not isinstance(part, dict): @@ -1082,7 +1138,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_function_call_to_chat_completion_message( - function_call: dict[str, Any], + function_call: Mapping[str, str], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API function_call into a Chat Completion message with tool calls @@ -1127,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: return [chat_completion_response_message] @staticmethod - def _resolve_file_id(item: dict[str, Any]) -> str | None: + def _resolve_file_id(item: Mapping[str, object]) -> object: """ Return the effective file_id for a Responses API input_file item. Explicit file_id takes precedence; file_url is used as fallback so @@ -1136,7 +1192,7 @@ class LiteLLMCompletionResponsesConfig: return item.get("file_id") or item.get("file_url") or None @staticmethod - def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]: + def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]: """ Transform a Responses API input_file item to a Chat Completion file item @@ -1146,21 +1202,21 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary with transformed file structure for Chat Completion """ - file_dict: Final[dict[str, Any]] = {} + file_dict: Final[dict[str, object]] = {} file_id: Final = LiteLLMCompletionResponsesConfig._resolve_file_id(item) if file_id: file_dict["file_id"] = file_id if item.get("file_data"): file_dict["file_data"] = item["file_data"] - new_item: Final[dict[str, Any]] = {"type": "file", "file": file_dict} + new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: new_item["cache_control"] = item["cache_control"] return new_item @staticmethod def _transform_input_image_item_to_image_item( - item: dict[str, Any], + item: Mapping[str, str], ) -> ChatCompletionImageObject: """ Transform a Responses API input_image item to a Chat Completion image item @@ -1173,8 +1229,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_content_to_chat_completion_content( - content: Any, - ) -> str | list[str | dict[str, Any]]: + content: object, + ) -> str | list[str | dict[str, object]]: """ Transform a Responses API content into a Chat Completion content @@ -1188,7 +1244,7 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(content, str): return content elif isinstance(content, list): - content_list: Final[list[str | dict[str, Any]]] = [] + content_list: Final[list[str | dict[str, object]]] = [] for item in content: if isinstance(item, str): content_list.append(item) @@ -1198,8 +1254,8 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) + image_block = _STR_KEY_DICT_ADAPTER.validate_python( + dict(LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)) ) if "cache_control" in item: image_block["cache_control"] = item["cache_control"] @@ -1209,7 +1265,7 @@ class LiteLLMCompletionResponsesConfig: text_value = item.get("text") if text_value is None: continue - content_block: dict[str, Any] = { + content_block: dict[str, object] = { "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( item.get("type") or "text" ), @@ -1299,7 +1355,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(typed_tool.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - chat_completion_tool: dict[str, Any] = { + chat_completion_tool: dict[str, object] = { "type": "function", "function": { "name": typed_tool.get("name") or "", @@ -1340,7 +1396,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to Responses API request tool format. Inverse of @@ -1348,7 +1404,7 @@ class LiteLLMCompletionResponsesConfig: """ if chat_completion_tools is None or not chat_completion_tools: return [] - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: if not isinstance(tool, dict): result.append(tool) @@ -1358,7 +1414,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(fn.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - responses_tool: dict[str, Any] = { + responses_tool: dict[str, object] = { "type": "function", "name": fn.get("name") or "", "description": fn.get("description") or "", @@ -1510,7 +1566,7 @@ class LiteLLMCompletionResponsesConfig: def convert_response_function_tool_call_to_chat_completion_tool_call( tool_call_item: Any, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format. @@ -1536,7 +1592,7 @@ class LiteLLMCompletionResponsesConfig: else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {}) ) - function_dict: Final[dict[str, Any]] = { + function_dict: Final[dict[str, object]] = { "name": tool_call_item.name, "arguments": tool_call_item.arguments, } @@ -1544,7 +1600,7 @@ class LiteLLMCompletionResponsesConfig: if provider_specific_fields: function_dict["provider_specific_fields"] = provider_specific_fields - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( getattr(tool_call_item, "id", None), getattr(tool_call_item, "call_id", None), @@ -1561,9 +1617,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: "ResponseApplyPatchToolCall", index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1581,7 +1637,7 @@ class LiteLLMCompletionResponsesConfig: import json operation_dict: Final = tool_call_item.operation.model_dump() - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": tool_call_item.call_id, "function": { "name": "apply_patch", @@ -1795,9 +1851,11 @@ class LiteLLMCompletionResponsesConfig: if not images: return image_generation_items - for idx, image_item in enumerate(images): + for idx, image_item in enumerate(_DICT_ITEMS_LIST_ADAPTER.validate_python(images)): # Extract base64 from data URL - image_url = image_item.get("image_url", {}).get("url", "") + image_url = _TEXT_ADAPTER.validate_python( + _ANY_KEY_DICT_ADAPTER.validate_python(image_item.get("image_url", {})).get("url", "") + ) base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2048,8 +2106,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_text_format_to_response_format( - text_param: dict[str, Any] | Any, - ) -> dict[str, Any] | None: + text_param: object, + ) -> dict[str, object] | None: """ Transform Responses API text.format parameter to Chat Completion response_format parameter. diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..2e1e1a44594 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,7 +5,7 @@ import json import time import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, + ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -112,7 +113,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: - if isinstance(error_obj, dict): + if _is_json_object(error_obj): raw_message = error_obj.get("message") raw_type = error_obj.get("type") raw_code = error_obj.get("code") @@ -243,7 +244,9 @@ class BaseResponsesAPIStreamingIterator: # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: - response_object: Final = getattr(openai_responses_api_chunk, "response", None) + response_object: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_object is not None: response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, @@ -279,7 +282,9 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part: Final = getattr(openai_responses_api_chunk, "part", None) + _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr( + openai_responses_api_chunk, "part", None + ) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( @@ -302,7 +307,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item: Final = getattr(openai_responses_api_chunk, "item", None) + item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None) if item: encrypted_content: Final = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): @@ -324,9 +329,11 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None) + response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Final[Any | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) @@ -414,7 +421,9 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None + response_obj: Final[ResponsesAPIResponse | None] = ( + getattr(self.completed_response, "response", None) if self.completed_response else None + ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) @@ -429,7 +438,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -506,7 +515,7 @@ class BaseResponsesAPIStreamingIterator: return request_kwargs = getattr(caching_handler, "request_kwargs", None) - if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: + if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() preset_cache_key = getattr(caching_handler, "preset_cache_key", None) @@ -606,7 +615,7 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: Final[dict[str, Any]] = {} + request_payload: Final[dict[str, object]] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: @@ -695,11 +704,15 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: object, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ - hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None) + hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = ( + getattr(iterator, "_call_post_streaming_deployment_hook", None) + ) if hook_fn is None: return chunk return await hook_fn(chunk) @@ -1019,7 +1032,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _dump_response_object(obj: Any) -> dict[str, Any]: if hasattr(obj, "model_dump"): return obj.model_dump() - if isinstance(obj, dict): + if _is_json_object(obj): return obj return {} @@ -1684,7 +1697,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final[Mapping[str, object]] = json.loads(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1925,7 +1938,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, Any], + completed_event: dict[str, object], ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2065,7 +2078,7 @@ class ManagedResponsesWebSocketHandler: Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested: Final = msg_obj.get("response") - response_params: Final[dict[str, Any]] = ( + response_params: Final[dict[str, object]] = ( nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { @@ -2076,7 +2089,7 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: dict[str, Any], + call_kwargs: dict[str, object], previous_response_id: str | None, current_messages: list[dict[str, object]], prior_history: list[dict[str, object]], @@ -2129,7 +2142,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a8af4eabb3f..68a6451e273 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3121 + "limit": 3114 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 834 }, "ANN201": { - "limit": 2032 + "limit": 2031 }, "ANN202": { "limit": 865 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1630 + "limit": 1555 }, "ASYNC230": { "limit": 11 @@ -81,7 +81,7 @@ "limit": 1 }, "C901": { - "limit": 315 + "limit": 314 }, "D419": { "limit": 6 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1240 + "limit": 1238 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0a0cfe9a617..3a670bc7345 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23235 + "limit": 23149 }, "LIT002": { - "limit": 27176 + "limit": 27166 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1091 + "limit": 1086 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16769 + "limit": 16760 }, "LIT011": { "limit": 5598 From fb7861fbfd9273ff081d496745456a772b2a36a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:41:06 -0700 Subject: [PATCH 66/74] build(lint): count deleted files toward check triggers --- CLAUDE.md | 2 +- scripts/pre_commit_lint.sh | 22 ++++++++----- tests/test_litellm/test_pre_commit_lint.py | 36 ++++++++++++++++++++++ 3 files changed, 52 insertions(+), 8 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 0354e3def53..abd5993eb1a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -41,7 +41,7 @@ Python max line length is 120, not 88 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing -Run `make check` (formerly `make pre-commit`, which still works as an alias) before every commit, merge commits included. It runs the CI-gating lint scoped to your staged files, so stage everything you intend to commit first; it warns about changed files you left unstaged and names the checks that were skipped because of them. With nothing staged it instead checks the working tree's diff against the merge base with origin/litellm_internal_staging, which is how you predict the CI lint on an already-committed branch, e.g. right after a merge commit +Run `make check` (formerly `make pre-commit`, which still works as an alias) before every commit, merge commits included. It runs the CI-gating lint scoped to your staged files, so stage everything you intend to commit first; it warns about changed files you left unstaged and names the checks that were skipped because of them. With nothing staged it instead checks the working tree's diff against the merge base with origin/litellm_internal_staging, which is how you predict the CI lint on an already-committed branch, e.g. right after a merge commit. Deleted files count toward which checks run (a deletion alone can turn CI red) in both modes `make check` saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index 1bf7fe17832..afe55603466 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -44,7 +44,7 @@ fi repo_root=$(git rev-parse --show-toplevel) cd "$repo_root" -staged=$(git diff --cached --name-only --diff-filter=ACMR) +staged=$(git diff --cached --name-only --diff-filter=ACMRD) unstaged=$(git diff --name-only) untracked=$(git ls-files --others --exclude-standard) @@ -57,7 +57,7 @@ else echo " Fix: git fetch origin litellm_internal_staging" >&2 exit 1 } - scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMR "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) + scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) if [ -z "$scope" ]; then echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" exit 0 @@ -68,6 +68,12 @@ fi scope_match() { printf '%s\n' "$scope" | grep -E "$1" || true; } +existing_files() { + while IFS= read -r f; do + if [ -f "$f" ]; then printf '%s\n' "$f"; fi + done +} + litellm_py_pattern='^litellm/.*\.py$' e2e_py_pattern='^tests/e2e/.*\.py$' spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' @@ -80,15 +86,17 @@ ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' litellm_py_files=$(scope_match "$litellm_py_pattern") e2e_py_files=$(scope_match "$e2e_py_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. -fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' || true) +fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types # (Prisma schema and configs included, not just Python) plus the generator and its # lockfiles, so match that whole trigger set rather than a Python subset. spec_files=$(scope_match "$spec_pattern") # CI's frontend-lint runs prettier over a wider extension set than eslint; keep that # split so this flags exactly what the job would. -ui_prettier_files=$(scope_match "$ui_prettier_pattern") -ui_eslint_files=$(scope_match "$ui_eslint_pattern") +ui_prettier_changed=$(scope_match "$ui_prettier_pattern") +ui_eslint_changed=$(scope_match "$ui_eslint_pattern") +ui_prettier_files=$(printf '%s\n' "$ui_prettier_changed" | existing_files) +ui_eslint_files=$(printf '%s\n' "$ui_eslint_changed" | existing_files) # CI lints the committed tree, so with staged files this script predicts CI for # what you have STAGED (every trigger above reads `git diff --cached`). The tools @@ -116,7 +124,7 @@ if [ -n "$staged" ]; then } warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" - warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_files" + warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi @@ -214,7 +222,7 @@ dashboard_checks() { lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; return 1; } } -if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then +if [ -n "$ui_prettier_changed" ] || [ -n "$ui_eslint_changed" ]; then dash_log=$(mktemp) set -m dashboard_checks > "$dash_log" 2>&1 & diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 12b8d338b49..35d98903226 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -181,6 +181,42 @@ def test_nothing_staged_includes_untracked_files_in_scope(tmp_path: Path) -> Non assert "linting Python" in proc.stdout +def test_nothing_staged_deletion_only_branch_triggers_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").unlink() + _commit_all(repo, "delete module") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing to check" not in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" in proc.stdout + assert "ruff format --check" not in proc.stdout + + +def test_staged_deletion_triggers_checks_without_feeding_missing_files_to_tools(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + subprocess.run(["git", "rm", "-q", "litellm/foo.py"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged" not in proc.stdout + assert "linting Python" in proc.stdout + assert "ruff format --check" not in proc.stdout + + +def test_deleted_dashboard_file_still_triggers_dashboard_lint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "ui" / "litellm-dashboard" / "src" / "app.ts").unlink() + _commit_all(repo, "delete dashboard file") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "linting dashboard" in proc.stdout + + def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") From f6df762b2537b797b2562af44ac9107ed4fe5c77 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:45:43 -0700 Subject: [PATCH 67/74] test: roll back live router replay membership between tests (#36278) Since #35491, every Router joins the module-global _live_routers weak set at construction, and every model cost map swap replays the deployments of every member on top of the freshly adopted map. #36039 isolated the register_model ledger half of that replay but not this half: under pytest-xdist, a Router created by an earlier test in the same worker that was still referenced (or simply not yet garbage collected) re-registered its deployments during TestPriceDataReloadIntegration::test_distributed_reload_check_function, and register_model hydrated the sparse mocked gpt-3.5-turbo entry into a full ModelInfo dict, failing the exact-equality assert (reruns cannot help since the polluting router survives in the worker process) The autouse isolate_litellm_state fixture now snapshots _live_routers before each test and restores its membership on teardown, so a test's routers stop contributing to cost map rebuilds once the test ends. A canary pair in test_conftest_isolation.py asserts the rollback --- tests/test_litellm/conftest.py | 8 +++++++ tests/test_litellm/test_conftest_isolation.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index c0993051d33..0dc8f56f3ce 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -19,6 +19,7 @@ sys.path.insert( import asyncio import litellm +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( @@ -244,6 +245,8 @@ def isolate_litellm_state(): for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } + original_live_routers = set(litellm_router_module._live_routers) + # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for # JSON logging and do not restore them, which breaks later caplog-based tests. logger_state = {} @@ -313,6 +316,11 @@ def isolate_litellm_state(): litellm_utils_module._runtime_registered_model_cost.clear() litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + for _router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(_router) + for _router in original_live_routers: + litellm_router_module._live_routers.add(_router) + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/test_litellm/test_conftest_isolation.py index 88889ad7740..15183e68f66 100644 --- a/tests/test_litellm/test_conftest_isolation.py +++ b/tests/test_litellm/test_conftest_isolation.py @@ -1,9 +1,15 @@ import litellm +from litellm import Router +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module CANARY_MODEL = "conftest-isolation-canary-model" +class _CanaryRouterHolder: + router: Router | None = None + + def test_register_model_ledger_entry_is_scoped_to_this_test(): litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost @@ -11,3 +17,20 @@ def test_register_model_ledger_entry_is_scoped_to_this_test(): def test_register_model_ledger_entry_was_rolled_back(): assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost + + +def test_live_router_membership_is_scoped_to_this_test(): + _CanaryRouterHolder.router = Router( + model_list=[ + { + "model_name": "conftest-isolation-canary-router", + "litellm_params": {"model": "openai/conftest-isolation-canary-backend", "api_key": "sk-canary"}, + } + ] + ) + assert _CanaryRouterHolder.router in litellm_router_module._live_routers + + +def test_live_router_membership_was_rolled_back(): + assert _CanaryRouterHolder.router is not None + assert _CanaryRouterHolder.router not in litellm_router_module._live_routers From 0d7f7c689a58aeea6cd26b8cb9fba90b31be8955 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 8 Aug 2026 12:19:22 -0700 Subject: [PATCH 68/74] test: repair stale CircleCI contracts --- litellm/proxy/db/autorouter_session_rollup.py | 47 +++++++++++++++++ .../auto_router_endpoints.py | 50 +------------------ tests/agent_tests/test_a2a_agent.py | 2 +- .../test_proxy_budget_reset.py | 24 ++++++--- .../base_responses_api.py | 4 +- .../spend/test_autorouter_session_rollup.py | 16 +++--- .../src/components/team/TeamInfo.test.tsx | 23 +-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 8 files changed, 81 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index b1f074c26b7..9732b1d7402 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -33,6 +33,53 @@ if TYPE_CHECKING: CACHE_TTL_5M_SECONDS: Final = 300 CACHE_TTL_1H_SECONDS: Final = 3600 +AUTOROUTER_BENCHMARKS_SQL: Final = """ +WITH windowed AS ( + SELECT * FROM "LiteLLM_AutoRouterSession" + WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +), +tier_maps AS ( + SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns + FROM ( + SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns + FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv + GROUP BY router_name, router_type, kv.key + ) per_tier + GROUP BY router_name, router_type +) +SELECT + agg.*, + COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns +FROM ( +SELECT + router_name, + router_type, + COUNT(*)::int AS sessions, + COALESCE(SUM(turns), 0)::int AS turns, + COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns, + COALESCE(SUM(covered_turns), 0)::int AS covered_turns, + COALESCE(SUM(cache_hits), 0)::int AS cache_hits, + COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns, + COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits, + COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns, + COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits, + COALESCE(SUM(return_turns), 0)::int AS return_turns, + COALESCE(SUM(return_hits), 0)::int AS return_hits, + COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses, + COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses, + COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns, + COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns, + COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, + COALESCE(SUM(spend), 0)::float8 AS spend, + COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, + COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds +FROM windowed +GROUP BY router_name, router_type +) agg +LEFT JOIN tier_maps USING (router_name, router_type) +ORDER BY agg.spend DESC +""" + @dataclass(frozen=True, slots=True) class AutoRouterTurnTransaction: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index d4221845b0c..8b6aafea751 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -26,6 +26,7 @@ from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.repositories.team_repository import TeamRepository from litellm.router_strategy.complexity_router import ComplexityRouter @@ -285,53 +286,6 @@ class _SessionAggRow(BaseModel): _SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow]) -_BENCHMARKS_SQL: Final = """ -WITH windowed AS ( - SELECT * FROM "LiteLLM_AutoRouterSession" - WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp -), -tier_maps AS ( - SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns - FROM ( - SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns - FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv - GROUP BY router_name, router_type, kv.key - ) per_tier - GROUP BY router_name, router_type -) -SELECT - agg.*, - COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns -FROM ( -SELECT - router_name, - router_type, - COUNT(*)::int AS sessions, - COALESCE(SUM(turns), 0)::int AS turns, - COALESCE(SUM(unordered_turns), 0)::int AS unordered_turns, - COALESCE(SUM(covered_turns), 0)::int AS covered_turns, - COALESCE(SUM(cache_hits), 0)::int AS cache_hits, - COALESCE(SUM(same_model_turns), 0)::int AS same_model_turns, - COALESCE(SUM(same_model_hits), 0)::int AS same_model_hits, - COALESCE(SUM(first_visit_turns), 0)::int AS first_visit_turns, - COALESCE(SUM(first_visit_hits), 0)::int AS first_visit_hits, - COALESCE(SUM(return_turns), 0)::int AS return_turns, - COALESCE(SUM(return_hits), 0)::int AS return_hits, - COALESCE(SUM(return_expired_misses), 0)::int AS return_expired_misses, - COALESCE(SUM(return_within_ttl_misses), 0)::int AS return_within_ttl_misses, - COALESCE(SUM(ttl_5m_turns), 0)::int AS ttl_5m_turns, - COALESCE(SUM(ttl_1h_turns), 0)::int AS ttl_1h_turns, - COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, - COALESCE(SUM(spend), 0)::float8 AS spend, - COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, - COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds -FROM windowed -GROUP BY router_name, router_type -) agg -LEFT JOIN tier_maps USING (router_name, router_type) -ORDER BY agg.spend DESC -""" - def _parse_benchmark_day(value: str) -> datetime: try: @@ -455,7 +409,7 @@ async def get_auto_router_benchmarks( raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date") raw_rows: Final = await prisma_client.db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, start_day.isoformat(), (end_day + timedelta(days=1)).isoformat(), ) diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py index f5ad9601369..1f72ced64f1 100644 --- a/tests/agent_tests/test_a2a_agent.py +++ b/tests/agent_tests/test_a2a_agent.py @@ -40,7 +40,7 @@ class MockA2AClient: name="mock-agent", url="http://mock-agent.local" ) - async def send_message(self, request): + async def send_message(self, request, *, context=None): from a2a.compat.v0_3.conversions import pb2_v10 for text in ("hel", "hello"): diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 44da3ea06a0..00d5380b2f4 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -622,8 +622,12 @@ async def test_service_logger_keys_success(): logger success hook is called with the correct event metadata and no exception is logged. """ keys = [ - {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"}, - {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"}, + _attrify( + {"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"} + ), + _attrify( + {"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"} + ), ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=keys) @@ -740,8 +744,12 @@ async def test_service_logger_users_success(): the correct metadata and no exception is logged. """ users = [ - {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"}, - {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"}, + _attrify( + {"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"} + ), + _attrify( + {"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"} + ), ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=users) @@ -853,8 +861,12 @@ async def test_service_logger_teams_success(): the proper metadata and nothing is logged as an exception. """ teams = [ - {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"}, - {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"}, + _attrify( + {"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"} + ), + _attrify( + {"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"} + ), ] prisma_client = MagicMock() prisma_client.get_data = AsyncMock(return_value=teams) diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index 407091a65b3..f5751aa79e8 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -338,7 +338,7 @@ class BaseResponsesAPITest(ABC): ) assert result is not None assert result.id == response.id - assert result.output == response.output + assert result.output_text == response.output_text else: raise ValueError("response is not a ResponsesAPIResponse") else: @@ -352,7 +352,7 @@ class BaseResponsesAPITest(ABC): ) assert result is not None assert result.id == response.id - assert result.output == response.output + assert result.output_text == response.output_text else: raise ValueError("response is not a ResponsesAPIResponse") diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 65b70f13a3b..7bc61c40ea0 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -12,8 +12,10 @@ from typing import Final import pytest -from litellm.proxy.db.autorouter_session_rollup import UPSERT_AUTOROUTER_SESSION_SQL -from litellm.proxy.management_endpoints.auto_router_endpoints import _BENCHMARKS_SQL +from litellm.proxy.db.autorouter_session_rollup import ( + AUTOROUTER_BENCHMARKS_SQL, + UPSERT_AUTOROUTER_SESSION_SQL, +) pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -164,7 +166,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -186,7 +188,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality") rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -248,7 +250,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -275,7 +277,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d ) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) @@ -289,7 +291,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None) rows = await db.query_raw( - _BENCHMARKS_SQL, + AUTOROUTER_BENCHMARKS_SQL, (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), ) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 513719a2ad9..50e10285148 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -919,7 +919,7 @@ describe("TeamInfoView", () => { }); }; - it("prefills pairs from team metadata, hides UI-managed keys, and round-trips typed values on save", async () => { + it("should preserve metadata types and hide managed keys", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(networking.teamInfoCall).mockResolvedValue( createMockTeamData({ @@ -964,27 +964,6 @@ describe("TeamInfoView", () => { expect(updateArg.model_tpm_limit).toEqual({ "gpt-4": 100 }); }); - it("includes a newly added pair in the team update", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); - vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); - - renderWithProviders(); - await openSettingsEditor(user); - - await user.click(screen.getByRole("button", { name: /add key-value pair/i })); - await user.type(screen.getByPlaceholderText("Key"), "cost_center"); - await user.type(screen.getByPlaceholderText("Value"), "eng-1"); - - await user.click(screen.getByRole("button", { name: /save changes/i })); - - await waitFor(() => { - expect(networking.teamUpdateCall).toHaveBeenCalled(); - }); - - expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ cost_center: "eng-1" }); - }); - it("should keep declared keys as ordinary prefilled rows and submit the edited value", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(useTeamMetadataSchema).mockReturnValue({ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a75c23da1cf..2fc70883c5e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24313,7 +24313,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From cfd64d45a85fba07e83a5a734b6ab5b1bcc8af44 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 19:28:57 +0000 Subject: [PATCH 69/74] fix(ui): show team BYOK models in team fallback settings (#36241) * fix(ui): show team BYOK models in team fallback settings Team router settings loaded fallback options from /model_group/info, which resolves models without a team, so a team's own BYOK deployments were never selectable in its own fallback config. Load the team-scoped listing when a team id is present. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): ignore stale team model responses in router settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): use react-query for fallback model listing in router settings accordion --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: ryan-crabbe-berri --- .../RouterSettingsAccordion.test.tsx | 71 +++++++++++++++++-- .../RouterSettingsAccordion.tsx | 27 +++---- .../llm_calls/fetch_models.test.tsx | 33 +++++++++ .../src/components/llm_calls/fetch_models.tsx | 12 +++- .../src/components/team/TeamInfo.tsx | 1 + 5 files changed, 118 insertions(+), 26 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx index a70b7602e5b..5ac3b8b2b64 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -1,7 +1,9 @@ -import { act, fireEvent, render, screen } from "@testing-library/react"; -import type { ReactNode } from "react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import type { ReactElement, ReactNode } from "react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; +import { fetchAvailableModels, fetchAvailableModelsForTeam } from "@/components/llm_calls/fetch_models"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./RouterSettingsAccordion"; vi.mock("../networking", () => ({ @@ -9,11 +11,14 @@ vi.mock("../networking", () => ({ })); vi.mock("@/components/llm_calls/fetch_models", () => ({ - fetchAvailableModels: vi.fn().mockResolvedValue([]), + fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "global-model" }]), + fetchAvailableModelsForTeam: vi.fn().mockResolvedValue([{ model_group: "openai/*" }, { model_group: "gpt-5" }]), })); vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({ - FallbackSelectionForm: () => null, + FallbackSelectionForm: ({ availableModels }: { availableModels: string[] }) => ( +
{availableModels.join(",")}
+ ), })); vi.mock("@tremor/react", () => ({ @@ -39,9 +44,19 @@ vi.mock("../router_settings/RouterSettingsForm", () => ({ ), })); +const renderWithQueryClient = (ui: ReactElement) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render(ui, { + wrapper: ({ children }: { children: ReactNode }) => ( + {children} + ), + }); +}; + describe("RouterSettingsAccordion", () => { beforeEach(() => { - vi.useFakeTimers(); + vi.clearAllMocks(); + vi.useFakeTimers({ shouldAdvanceTime: true }); }); afterEach(() => { @@ -58,7 +73,7 @@ describe("RouterSettingsAccordion", () => { it("debounces propagation and calls onChange once with the last value", async () => { const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); - render(); + renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); @@ -81,9 +96,51 @@ describe("RouterSettingsAccordion", () => { expect(onChange.mock.calls[0][0].router_settings.routing_strategy).toBe("usage-based-routing"); }); + it("offers the team's own models, including team-scoped BYOK ones, when a teamId is given", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("gpt-5,openai/*"); + }); + expect(fetchAvailableModelsForTeam).toHaveBeenCalledWith("test-token", "team-123"); + expect(fetchAvailableModels).not.toHaveBeenCalled(); + }); + + it("falls back to the proxy-wide model listing when no teamId is given", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("global-model"); + }); + expect(fetchAvailableModelsForTeam).not.toHaveBeenCalled(); + }); + + it("ignores a stale team's model response that resolves after a newer team was selected", async () => { + const resolvers: ((models: { model_group: string }[]) => void)[] = []; + vi.mocked(fetchAvailableModelsForTeam).mockImplementation( + () => new Promise((resolve) => resolvers.push(resolve)) as Promise<{ model_group: string }[]>, + ); + + const { rerender } = renderWithQueryClient(); + await waitFor(() => expect(resolvers).toHaveLength(1)); + + rerender(); + await waitFor(() => expect(resolvers).toHaveLength(2)); + + await act(async () => { + resolvers[1]([{ model_group: "fast-team-model" }]); + resolvers[0]([{ model_group: "slow-team-model" }]); + }); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("fast-team-model"); + }); + expect(screen.getByTestId("available-models")).not.toHaveTextContent("slow-team-model"); + }); + it("does not call onChange when unmounted mid-wait", async () => { const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); - const { unmount } = render(); + const { unmount } = renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 08b917e302f..56227abe9ea 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -1,12 +1,13 @@ import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; +import { useQuery } from "@tanstack/react-query"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { getRouterSettingsCall } from "../networking"; import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks"; import { FallbackSelectionForm } from "../Settings/RouterSettings/Fallbacks/FallbackSelectionForm"; import { FallbackGroup } from "../Settings/RouterSettings/Fallbacks/FallbackGroupConfig"; -import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +import { fetchAvailableModels, fetchAvailableModelsForTeam, ModelGroup } from "@/components/llm_calls/fetch_models"; export interface RouterSettingsAccordionValue { router_settings: { @@ -30,6 +31,7 @@ interface RouterSettingsAccordionProps { value?: RouterSettingsAccordionValue; onChange?: (value: RouterSettingsAccordionValue) => void; modelData?: any; + teamId?: string | null; } export interface RouterSettingsAccordionRef { @@ -39,7 +41,7 @@ export interface RouterSettingsAccordionRef { const PROPAGATE_WAIT_MS = 100; const RouterSettingsAccordion = forwardRef( - ({ accessToken, value, onChange, modelData }, ref) => { + ({ accessToken, value, onChange, modelData, teamId }, ref) => { const [formValue, setFormValue] = useState({ routerSettings: {}, selectedStrategy: null, @@ -47,7 +49,6 @@ const RouterSettingsAccordion = forwardRef([]); const [fallbackGroups, setFallbackGroups] = useState([]); - const [modelInfo, setModelInfo] = useState([]); const [availableRoutingStrategies, setAvailableRoutingStrategies] = useState([]); const [routerFieldsMetadata, setRouterFieldsMetadata] = useState<{ [key: string]: any }>({}); const [routingStrategyDescriptions, setRoutingStrategyDescriptions] = useState<{ [key: string]: string }>({}); @@ -175,21 +176,11 @@ const RouterSettingsAccordion = forwardRef { - if (!accessToken) { - return; - } - const loadModels = async () => { - try { - const uniqueModels = await fetchAvailableModels(accessToken); - setModelInfo(uniqueModels); - } catch (error) { - console.error("Error fetching model info for fallbacks:", error); - } - }; - loadModels(); - }, [accessToken]); + const { data: modelInfo = [] } = useQuery({ + queryKey: ["fallbackAvailableModels", accessToken, teamId ?? null], + queryFn: () => (teamId ? fetchAvailableModelsForTeam(accessToken, teamId) : fetchAvailableModels(accessToken)), + enabled: Boolean(accessToken), + }); // Helper function to build router_settings from current state const buildRouterSettings = (): RouterSettingsAccordionValue["router_settings"] => { diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx new file mode 100644 index 00000000000..bd691c7f629 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx @@ -0,0 +1,33 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { modelAvailableCall } from "@/components/networking"; +import { fetchAvailableModelsForTeam } from "./fetch_models"; + +vi.mock("@/components/networking", () => ({ + modelAvailableCall: vi.fn(), + modelHubCall: vi.fn(), +})); + +const modelAvailableCallMock = vi.mocked(modelAvailableCall); + +describe("fetchAvailableModelsForTeam", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("requests the models scoped to the team so team-only BYOK models are included", async () => { + modelAvailableCallMock.mockResolvedValue({ + data: [{ id: "all-proxy-models" }, { id: "openai/*" }, { id: "gpt-5-mini" }, { id: "openai/*" }], + }); + + const models = await fetchAvailableModelsForTeam("token", "team-123"); + + expect(modelAvailableCallMock).toHaveBeenCalledWith("token", "", "", false, "team-123"); + expect(models).toEqual([{ model_group: "gpt-5-mini" }, { model_group: "openai/*" }]); + }); + + it("returns an empty list when the team has no models", async () => { + modelAvailableCallMock.mockResolvedValue({ data: [] }); + + expect(await fetchAvailableModelsForTeam("token", "team-123")).toEqual([]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 0de98330c2e..a1690b1307e 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -1,12 +1,22 @@ // fetch_models.ts -import { modelHubCall } from "@/components/networking"; +import { excludeProxyWideSentinel } from "@/components/key_team_helpers/fetch_available_models_team_key"; +import { modelAvailableCall, modelHubCall } from "@/components/networking"; export interface ModelGroup { model_group: string; mode?: string; } +export const fetchAvailableModelsForTeam = async (accessToken: string, teamId: string): Promise => { + const response = await modelAvailableCall(accessToken, "", "", false, teamId); + const modelNames: string[] = (response?.data ?? []).map((model: { id: string }) => model.id); + + return excludeProxyWideSentinel(Array.from(new Set(modelNames))) + .sort((a, b) => a.localeCompare(b)) + .map((model) => ({ model_group: model })); +}; + /** * Fetches available models using modelHubCall and formats them for the selection dropdown. */ diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index e57bf59b2da..6c6fdcaedd6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1215,6 +1215,7 @@ const TeamInfoView: React.FC = ({ From 4150248095bd5a43e44d70c3b3d94eea74069611 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 8 Aug 2026 12:38:23 -0700 Subject: [PATCH 70/74] chore: remove pre-commit rule some users do not use make pre-commit as it is a multi-minute process. I personally use it but I want users themselves to decide whether to pre-commit before each commit or not, based on what works best for them --- CLAUDE.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index abd5993eb1a..b0da2970fc8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -41,9 +41,7 @@ Python max line length is 120, not 88 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing -Run `make check` (formerly `make pre-commit`, which still works as an alias) before every commit, merge commits included. It runs the CI-gating lint scoped to your staged files, so stage everything you intend to commit first; it warns about changed files you left unstaged and names the checks that were skipped because of them. With nothing staged it instead checks the working tree's diff against the merge base with origin/litellm_internal_staging, which is how you predict the CI lint on an already-committed branch, e.g. right after a merge commit. Deleted files count toward which checks run (a deletion alone can turn CI red) in both modes - -`make check` saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice +`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in From 12aeb53aec57b2967a50cb826bf0ef2167bc632c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 8 Aug 2026 12:40:00 -0700 Subject: [PATCH 71/74] fix(otel): mark v2 server spans as failed for pre-call errors (#34546) * fix(otel): mark v2 server spans as failed for pre-call errors (LIT-4780) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): authenticate malformed-body requests before rejecting them (LIT-4780) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(auth): cover malformed-body rejection when auth error is recovered Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(auth): skip authorization for a request whose body never parsed Deferring the parse failure ran the full auth phase, including budget reservation, whose reserved amount is only released by the endpoint's post call path; the endpoint never runs, so malformed requests leaked reservations and locked a budgeted key out. Authorization now runs only when the body parsed, and a parse failure with a rejected key keeps returning the 400 it returned before. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: shivam Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/logger.py | 20 ++- litellm/proxy/auth/user_api_key_auth.py | 170 ++++++++++++------ .../integrations/otel/test_otel_v2_logger.py | 35 +++- .../proxy/auth/test_user_api_key_auth.py | 111 ++++++++++++ 4 files changed, 269 insertions(+), 67 deletions(-) diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e273289168a..2c83406afed 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -33,6 +33,7 @@ from litellm.integrations.otel.model.payloads import ( is_mcp_list_tools, is_mcp_tool_call, ) +from litellm.integrations.otel.model.semconv import Error from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service from litellm.integrations.otel.model.utils import to_ns from litellm.integrations.otel.plumbing.context import ( @@ -634,18 +635,23 @@ class OpenTelemetryV2(CustomLogger): """Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a failure that dies before any LLM-call span exists (malformed body, auth / validation rejection). Called from the proxy's global exception handler via - ``_close_dangling_otel_server_span``. The instrumentor still owns the span's - status and lifecycle, so this only decorates it — never sets status, never - ends it — and emits no exception event, matching v1's SERVER-span behavior - and avoiding a duplicate of the event ``async_post_call_failure_hook`` or - the ``auth`` phase span already records.""" + ``_close_dangling_otel_server_span``, which swallows the exception into a + ``JSONResponse`` so the instrumentor never sees it and leaves the span + ``UNSET``; the status is set here instead (v1 did the same from the handler) + so a failed request reads as failed and not merely as a span carrying an + error message. The instrumentor still owns the span's lifecycle, so this + never ends it. The exception event is recorded only when nothing stamped + this span already — ``async_post_call_failure_hook`` and the ``auth`` phase + span record their own, and a second event would duplicate it — while the + attributes are always restamped so ``error.code`` stays pinned to the real + response status.""" if span is None or not is_recordable_span(span): return + already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ()) stamp_error( span, _span_error_from_exception(exception, status_code=status_code), - record_event=False, - set_status=False, + record_event=not already_stamped, ) async def async_post_call_failure_hook( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9dc450befc2..4baa7b99a4f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1044,6 +1044,22 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None: request.state.parent_otel_span = parent_otel_span +async def _read_request_body_deferring_parse_failure( + request: Request, +) -> tuple[dict, ProxyException | None]: + """Parse the body, returning a parse failure instead of raising it. + + A body that fails to parse is still a request from a known caller, so auth + must run (resolving identity onto the request's trace) before the 400 goes + out; the caller re-raises the returned exception once identity is seeded. + """ + try: + parsed_body: Final = await _read_request_body(request=request) + except ProxyException as parse_exception: + return {}, parse_exception # mutable-ok: request_data is a plain dict across the whole auth path + return populate_request_with_path_params(request_data=parsed_body, request=request), None + + async def _user_api_key_auth_builder( request: Request, api_key: str, @@ -2516,6 +2532,72 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> ) +async def _authorize_authenticated_request( + user_api_key_auth_obj: UserAPIKeyAuth, + request: Request, + request_data: dict, + route: str, + api_key: str, +) -> UserAPIKeyAuth | None: + """Authorize an already-authenticated request: disabled-route check, the single + ``common_checks`` gate (which also reserves budget), and end-user fallback + resolution. Returns the auth object the exception handler recovered when a check + failed but the request may proceed anyway, else ``None``. + """ + ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## + RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) + + # Single authorization point. Builder paths MUST NOT call common_checks. + # Route through the same exception handler the builder uses so + # authorization failures (ProxyException, or plain Exception from + # admin-only-route / model-access / budget checks) surface as + # ProxyException consistently with pre-refactor behavior. + try: + await _run_centralized_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request=request, + request_data=request_data, + route=route, + ) + except Exception as e: + return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + e=e, + request=request, + request_data=request_data, + route=route, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + api_key=api_key, + resolved_identity=user_api_key_auth_obj, + ) + + # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return + # paths (no master key, /user/auth route, JWT short-circuits) that bypass + # the end-user resolution block. If those paths produced an auth obj + # without an ``end_user_id`` set, fall back to extracting from the request + # body so spend logs are still attributed correctly. Validation honours + # ``litellm.validate_end_user_id_in_db``. + if user_api_key_auth_obj.end_user_id is None: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) + if raw_end_user_id is not None: + resolved_end_user_id: Final = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if resolved_end_user_id is not None: + user_api_key_auth_obj.end_user_id = resolved_end_user_id + return None + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2536,8 +2618,7 @@ async def user_api_key_auth( # close, and the trace never reaches the backend. _ensure_parent_otel_span_on_request_state(request) - request_data = await _read_request_body(request=request) - request_data = populate_request_with_path_params(request_data=request_data, request=request) + request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request) route: Final[str] = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED @@ -2545,69 +2626,41 @@ async def user_api_key_auth( # triggers (key/user/team object reads) nest under it instead of flattening # onto the server span. No-op when OTel V2 isn't active. with phase_span(f"auth {route}"): - user_api_key_auth_obj: Final = await _user_api_key_auth_builder( - request=request, - api_key=api_key, - azure_api_key_header=azure_api_key_header, - anthropic_api_key_header=anthropic_api_key_header, - google_ai_studio_api_key_header=google_ai_studio_api_key_header, - azure_apim_header=azure_apim_header, - request_data=request_data, - custom_litellm_key_header=custom_litellm_key_header, - ) + try: + user_api_key_auth_obj: Final = await _user_api_key_auth_builder( + request=request, + api_key=api_key, + azure_api_key_header=azure_api_key_header, + anthropic_api_key_header=anthropic_api_key_header, + google_ai_studio_api_key_header=google_ai_studio_api_key_header, + azure_apim_header=azure_apim_header, + request_data=request_data, + custom_litellm_key_header=custom_litellm_key_header, + ) + except Exception: + # The body was read first, so a caller who sent both a malformed body and + # a rejected key used to get the 400; the response is unchanged, and the + # auth failure is still recorded on the trace by the handler that ran. + if body_parse_exception is not None: + raise body_parse_exception + raise user_api_key_auth_obj.budget_reservation = None - ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## - RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) - - # Single authorization point. Builder paths MUST NOT call common_checks. - # Route through the same exception handler the builder uses so - # authorization failures (ProxyException, or plain Exception from - # admin-only-route / model-access / budget checks) surface as - # ProxyException consistently with pre-refactor behavior. - try: - await _run_centralized_common_checks( + # A body that never parsed is authenticated (so the trace carries identity + # and this ``auth`` span) but not authorized: there is no model to check it + # against, and budget reservation would increment live spend counters that + # only the endpoint's post-call path releases; the endpoint never runs, since + # the parse failure is raised below. + if body_parse_exception is None: + recovered_auth_obj: Final = await _authorize_authenticated_request( user_api_key_auth_obj=user_api_key_auth_obj, request=request, request_data=request_data, route=route, - ) - except Exception as e: - return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( - e=e, - request=request, - request_data=request_data, - route=route, - parent_otel_span=user_api_key_auth_obj.parent_otel_span, api_key=api_key, - resolved_identity=user_api_key_auth_obj, ) - - # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return - # paths (no master key, /user/auth route, JWT short-circuits) that bypass - # the end-user resolution block. If those paths produced an auth obj - # without an ``end_user_id`` set, fall back to extracting from the request - # body so spend logs are still attributed correctly. Validation honours - # ``litellm.validate_end_user_id_in_db``. - if user_api_key_auth_obj.end_user_id is None: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) - if raw_end_user_id is not None: - resolved_end_user_id: Final = await resolve_and_validate_end_user_id( - raw_end_user_id=raw_end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth_obj.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - if resolved_end_user_id is not None: - user_api_key_auth_obj.end_user_id = resolved_end_user_id + if recovered_auth_obj is not None: + return recovered_auth_obj # Identity is now resolved. Seed it AFTER the auth span closes so the Baggage # persists on the request task (detaching the span's context token inside the @@ -2619,6 +2672,9 @@ async def user_api_key_auth( ) user_api_key_auth_obj.request_route = normalize_request_route(route) + if body_parse_exception is not None: + raise body_parse_exception + # Resolve caller identity once, here at the seam, into a single per-request # Principal projected off the key object the builder already fetched (no # second lookup). Downstream consumers read identity off this instead of diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 2573ad5a375..82b074220fa 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -1195,8 +1195,13 @@ def test_async_post_call_failure_hook_skips_a_transport_that_already_answered(): def test_record_error_attributes_on_span_decorates_without_ending(): """PATH A: a failure that dies before any LLM-call span (malformed body, validation) is stamped onto the instrumentor-owned SERVER span. The method must - not end the span or emit a duplicate exception event, and must pin error.code - to the real response status (not the exception's own code).""" + not end the span, and must pin error.code to the real response status (not the + exception's own code). + + LIT-4780: the instrumentor never sees the exception (the proxy handler turns it + into a JSONResponse), so nothing else marks the span as failed; the status and + the exception event have to come from here or the trace shows the error message + on an otherwise successful-looking request.""" logger, exporter = _logger() server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422) @@ -1206,7 +1211,31 @@ def test_record_error_attributes_on_span_decorates_without_ending(): assert span.attributes["error.type"] == "ProxyException" assert span.attributes["error.message"] == "Invalid JSON body" assert span.attributes["litellm.provider.error.code"] == "422" - assert all(e.name != "exception" for e in span.events) + assert span.status.status_code is StatusCode.ERROR + assert [e.name for e in span.events] == ["exception"] + + +def test_record_error_attributes_on_span_does_not_duplicate_an_already_stamped_error(): + """A failure that already went through ``async_post_call_failure_hook`` reaches + the exception handler too; the second stamp must keep one exception event while + still repinning error.code to the real response status.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("Authentication Error, invalid key", 401) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth() + ) + ) + logger.record_error_attributes_on_span(server, exc, 400) + server.end() + (span,) = exporter.get_finished_spans() + assert [e.name for e in span.events] == ["exception"] + assert span.attributes["litellm.provider.error.code"] == "400" + assert span.status.status_code is StatusCode.ERROR def test_record_error_attributes_on_span_ignores_below_400_and_missing_span(): diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6fc44ef7519..60d9689dc0b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4786,6 +4786,117 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error(): + """Regression (LIT-4780): a body that fails to parse must still be authenticated + first, so the rejected request's trace carries the caller's key / team / user + identity instead of an anonymous root span. The parse error is re-raised + unchanged once identity is seeded.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="team-1") + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = b'{}{"model": "gpt-4o"}' + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ) as mock_builder, + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ) as mock_common_checks, + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.seed_request_identity", + ) as mock_seed, + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer sk-test") + + assert "Invalid JSON payload" in str(exc_info.value.message) + assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST) + mock_builder.assert_awaited_once() + assert mock_seed.call_args.args[0] is builder_token + # authorization must not run for a request that is about to be rejected: + # ``common_checks`` reserves budget against live spend counters that only the + # endpoint's post-call path releases, and the endpoint never runs here + mock_common_checks.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_malformed_body_with_rejected_key_still_returns_the_parse_error(): + """The body is read before the key is authenticated, so a caller who sends both a + malformed body and a key that fails auth gets the 400. Authenticating the request + first (LIT-4780) must not turn that into the auth status code.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = b'{}{"model": "gpt-4o"}' + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + side_effect=ProxyException( + message="Authentication Error, invalid key", + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer sk-bad") + + assert "Invalid JSON payload" in str(exc_info.value.message) + assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + def _proxy_attrs_for_db_lookup(): """Minimal proxy_server attributes for driving the real ``_user_api_key_auth_builder`` down to the DB key lookup.""" From 1a40a673942a52830155a8d414c046a4f9376223 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 8 Aug 2026 12:54:23 -0700 Subject: [PATCH 72/74] fix: stabilize generated user role ordering --- litellm/proxy/_types.py | 4 ++-- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1fc05ac4653..fa89df39c5f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4557,10 +4557,10 @@ class DefaultInternalUserParams(LiteLLMPydanticObjectBase): user_role: ( Literal[ - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, ] | None ) = Field( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2fc70883c5e..a75c23da1cf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24313,7 +24313,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams From ff5f8132d178ee700d23bff3e97d639ea5027cbd Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 19:55:37 +0000 Subject: [PATCH 73/74] docs: clarify guideline priority ordering in CLAUDE.md Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- CLAUDE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index b0da2970fc8..f1bb46c1fd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -9,7 +9,7 @@ Don't assume that the existing code is correct or the right way of doing things - easy to maintain/change - modern -In that order of importance +In descending order of importance When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate From e35ee4e5fa3e25a5f750fa0ee23525a8277a4f62 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 8 Aug 2026 13:02:29 -0700 Subject: [PATCH 74/74] feat(router): independent, default-on deployment affinity for the auto-router (#36146) --- litellm/constants.py | 1 + litellm/proxy/common_utils/callback_utils.py | 3 +- litellm/proxy/litellm_pre_call_utils.py | 2 + litellm/router.py | 96 ++-- .../complexity_router/complexity_router.py | 50 +- .../complexity_router/config.py | 30 +- .../deployment_affinity_check.py | 266 ++++++++--- litellm/types/router.py | 1 + .../proxy/test_litellm_pre_call_utils.py | 2 + .../router_strategy/test_complexity_router.py | 97 ++++ .../test_deployment_affinity_check.py | 9 +- .../test_session_id_affinity.py | 441 +++++++++++++++++- tests/test_litellm/test_router.py | 65 +++ type-discipline-budget.json | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 +- 15 files changed, 964 insertions(+), 111 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 30d3bb1f26e..3b91f23fe39 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1322,6 +1322,7 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks" LITELLM_METADATA_FIELD: Final = "litellm_metadata" OLD_LITELLM_METADATA_FIELD: Final = "metadata" RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" +SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 22200567012..fbf28e223c1 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger -from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -425,6 +425,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( "_guardrail_pipelines", "_pipeline_managed_guardrails", PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 83ae59ef050..c48fee96646 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -18,6 +18,7 @@ from litellm.constants import ( INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -226,6 +227,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "applied_policies", "policy_sources", "routing_decision", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, "standard_logging_object", "proxy_server_request", diff --git a/litellm/router.py b/litellm/router.py index 1a76eb5d59b..feaf69a44ae 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -46,6 +46,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function @@ -135,6 +136,7 @@ from litellm.router_utils.handle_error import ( from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, + warn_on_unknown_model_group_affinity_flags, ) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( build_io_token_rate_limit_headers, @@ -603,6 +605,10 @@ class Router: # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. self._zero_cost_cache: dict[str, bool] = {} + self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds + self.model_group_affinity_config = model_group_affinity_config + warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) + if model_list is not None: # set_model_list will build indices automatically self.set_model_list(model_list) @@ -744,7 +750,6 @@ class Router: litellm.failure_callback = [self.deployment_callback_on_failure] self.routing_strategy_args = routing_strategy_args self.provider_budget_config = provider_budget_config - self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.router_budget_logger: RouterBudgetLimiting | None = None if RouterBudgetLimiting.should_init_router_budget_limiter( model_list=model_list, provider_budget_config=self.provider_budget_config @@ -766,7 +771,6 @@ class Router: ) self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy - self.model_group_affinity_config: dict[str, list[str]] | None = model_group_affinity_config self.allowed_fails_policy: AllowedFailsPolicy | None = None if allowed_fails_policy is not None: @@ -789,21 +793,8 @@ class Router: # If model_group_affinity_config is set but no global affinity checks were # enabled, we still need the DeploymentAffinityCheck callback (with global # flags all False) so per-group config can activate affinity per model group. - if self.model_group_affinity_config and not any( - isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or []) - ): - if self.optional_callbacks is None: - self.optional_callbacks = [] - affinity_callback: Final = DeploymentAffinityCheck( - cache=self.cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=False, - enable_responses_api_affinity=False, - enable_session_id_affinity=False, - model_group_affinity_config=self.model_group_affinity_config, - ) - self.optional_callbacks.append(affinity_callback) - litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + if self.model_group_affinity_config: + self._ensure_deployment_affinity_callback() if self.alerting_config is not None: self._initialize_alerting() @@ -1662,6 +1653,28 @@ class Router: _move_before_deployment_affinity(self.optional_callbacks, ec_callback) _move_before_deployment_affinity(litellm.callbacks, ec_callback) + def _ensure_deployment_affinity_callback(self) -> None: + """Register the DeploymentAffinityCheck callback (global flags all False) if absent. + + Needed when nothing enabled a global affinity flag but affinity can still + activate per request: per-group `model_group_affinity_config` entries, or the + session-affinity marker a complexity router stamps at pre-routing time. + """ + if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])): + return + if self.optional_callbacks is None: + self.optional_callbacks = [] + affinity_callback: Final = DeploymentAffinityCheck( + cache=self.cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=False, + model_group_affinity_config=self.model_group_affinity_config, + ) + self.optional_callbacks.append(affinity_callback) + litellm.logging_callback_manager.add_litellm_callback(affinity_callback) + def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None): if optional_pre_call_checks is None: return @@ -7683,6 +7696,8 @@ class Router: strategy=complexity_router, strategy_label="Complexity-router", ) + if complexity_router._uses_deployment_pin: + self._ensure_deployment_affinity_callback() def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" @@ -11190,6 +11205,9 @@ class Router: router_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs) if router_strategy is None: self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) + self._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None + ) return None pre_routing_hook_response: Final = await router_strategy.async_pre_routing_hook( @@ -11203,6 +11221,11 @@ class Router: request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), ) + self._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None), + ) # `model` (the alias, e.g. "smart-router") is never the deployment actually # called - apply the alias's own litellm_params (besides `model` itself, @@ -11234,21 +11257,40 @@ class Router: to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ - if routing_decision is None: + Router._stamp_or_clear_metadata_key( + request_kwargs=request_kwargs, + key="routing_decision", + value=( + None + if routing_decision is None + else Router._redact_prompt_text_if_needed( + request_kwargs=request_kwargs, routing_decision=routing_decision + ) + ), + ) + + @staticmethod + def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None: + """Write a proxy-internal metadata key for THIS routing attempt, or clear it. + + Fallbacks and retries re-enter the pre-routing hook with the same + `request_kwargs`, so every attempt must write or clear, never just write; + a value left behind by an earlier attempt would be attributed to this one. + `get_or_create_metadata_bucket` is the single owner of "which dict holds + proxy-internal metadata": it picks `litellm_metadata` when present (so the + value never lands in the `metadata` dict that routes like /v1/messages + forward to the provider) and replaces a non-dict value rather than silently + skipping the write. Clearing pops from BOTH buckets so a request whose + bucket resolution changed between attempts cannot resurrect a stale value. + """ + if value is None: for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): if isinstance(bucket, dict): - bucket.pop("routing_decision", None) + bucket.pop(key, None) return - # `get_or_create_metadata_bucket` is the single owner of "which dict holds - # proxy-internal metadata": it picks `litellm_metadata` when present (so the - # decision never lands in the `metadata` dict that routes like /v1/messages - # forward to the provider) and replaces a non-dict value rather than silently - # skipping the write. _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed( - request_kwargs=request_kwargs, routing_decision=routing_decision - ) + metadata_bucket[key] = value @staticmethod def _redact_prompt_text_if_needed( diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f6ced0bb9d2..32d252f3f68 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1651,6 +1651,28 @@ class ComplexityRouter(CustomLogger): caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + @property + def _uses_tier_pin(self) -> bool: + return bool(self.config.session_affinity and not self.config.plugins) + + @property + def _uses_deployment_pin(self) -> bool: + """session_affinity implies the deployment pin: a session frozen onto one model + group but load-balanced across its deployments would still go cache-cold, which + is the exact failure both flags exist to prevent.""" + return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins) + + def _with_session_deployment_affinity( + self, response: PreRoutingHookResponse | None + ) -> PreRoutingHookResponse | None: + if response is None or not self._uses_deployment_pin: + return response + return response.model_copy( + update={ # mutable-ok: model_copy types update as a plain dict + "session_affinity_ttl_seconds": self.config.session_affinity_ttl_seconds + } + ) + async def async_pre_routing_hook( self, model: str, @@ -1685,7 +1707,7 @@ class ComplexityRouter(CustomLogger): resolved_messages: Final = self._resolve_messages(messages, request_kwargs) conversation_continuing: Final = _conversation_is_continuing(resolved_messages) - use_session_affinity: Final = self.config.session_affinity and not self.config.plugins + use_session_affinity: Final = self._uses_tier_pin session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None @@ -1724,17 +1746,19 @@ class ComplexityRouter(CustomLogger): "ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model ) has_original_messages: Final = messages is not None and len(messages) > 0 - return PreRoutingHookResponse( - model=routed_model, - messages=messages if has_original_messages else None, - routing_decision=self._build_routing_decision( - routed_model=routed_model, - cause=cause, - tier=self._tier_for_model(routed_model), - escalation_keyword=pin_escalation_keyword, - escalated=escalated, - conversation_continuing=conversation_continuing, - ), + return self._with_session_deployment_affinity( + PreRoutingHookResponse( + model=routed_model, + messages=messages if has_original_messages else None, + routing_decision=self._build_routing_decision( + routed_model=routed_model, + cause=cause, + tier=self._tier_for_model(routed_model), + escalation_keyword=pin_escalation_keyword, + escalated=escalated, + conversation_continuing=conversation_continuing, + ), + ) ) response: Final = await self._classify_and_route( @@ -1752,7 +1776,7 @@ class ComplexityRouter(CustomLogger): value=response.model, ttl=self.config.session_affinity_ttl_seconds, ) - return response + return self._with_session_deployment_affinity(response) async def _classify_and_route( self, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 69609a973b0..0999af66fd8 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -508,13 +508,39 @@ class ComplexityRouterConfig(BaseModel): "session's first turn and reuse it for every later turn, skipping re-classification. " "Off by default so every turn is classified on its own merits and routed to the cheapest " "adequate tier. Set True to keep a multi-turn session on one model, which preserves " - "provider prompt caches and avoids cross-model conversation-history errors." + "provider prompt caches and avoids cross-model conversation-history errors. Always " + "implies the deployment pin regardless of deployment_affinity: the session sticks to " + "one deployment of the pinned model, since freezing the model while re-shuffling its " + "deployments would still go cache-cold." + ), + ) + deployment_affinity: bool = Field( + default=True, + description=( + "When True and a session_id is resolvable on the request, pin the deployment chosen " + "inside each routed model group and reuse it whenever the session returns to that " + "group, without pinning which group the session routes to. Independent of " + "session_affinity, which pins the model group instead (and always carries this " + "deployment pin with it): with session_affinity off, " + "every turn is still classified on its own merits while a session that escalates to a " + "stronger tier and comes back still lands on the deployment it used before, which is " + "what keeps a provider prompt cache warm. Pins are held per model group, so switching " + "tiers does not disturb the pin left behind in the previous group. On by default " + "because re-shuffling a conversation across deployments of the same model discards " + "that cache for no benefit; set False to keep every turn load-balanced across the " + "group, which is what a deployment set with tight per-deployment rate limits wants. " + "Inert when no session_id is resolvable, since there is nothing to key a pin on, and " + "suppressed when plugins are configured, for the same reason session_affinity is." ), ) session_affinity_ttl_seconds: int = Field( default=3600, gt=0, - description="TTL for the session affinity pin; refreshed on every cache hit", + description=( + "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the " + "session_affinity model pin and the deployment_affinity deployment pin, so it measures " + "idle time for the session's routing decisions rather than total session length" + ), ) plugins: list[RoutingPlugin] | None = Field( diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index b4408d17ffb..7fb90ab89de 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,12 +13,15 @@ where routing to a consistent deployment is still beneficial. """ import hashlib +import json +from collections.abc import Mapping, Sequence from typing import Any, Final, cast from typing_extensions import TypedDict from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -29,6 +32,47 @@ class DeploymentAffinityCacheValue(TypedDict): model_id: str +VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( + { + "deployment_affinity", + "responses_api_deployment_check", + "session_affinity", + "encrypted_content_affinity", + } +) + + +def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapping[str, Sequence[str]] | None) -> None: + """`model_group_affinity_config` is one Router-level config consumed by two callbacks: + DeploymentAffinityCheck acts on three of the flags and EncryptedContentAffinityCheck + on the fourth, so typo detection lives here at the schema, not inside either consumer. + """ + if model_group_affinity_config is None: + return + for group, flags in model_group_affinity_config.items(): + unknown = set(flags) - VALID_MODEL_GROUP_AFFINITY_FLAGS + if unknown: + verbose_router_logger.warning( + "model_group_affinity_config: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", + unknown, + group, + VALID_MODEL_GROUP_AFFINITY_FLAGS, + ) + + +_CLAIM_PIN_SCRIPT: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + + class DeploymentAffinityCheck(CustomLogger): """ Router deployment affinity callback. @@ -38,14 +82,6 @@ class DeploymentAffinityCheck(CustomLogger): """ CACHE_KEY_PREFIX = "deployment_affinity:v1" - VALID_FLAGS = frozenset( - { - "deployment_affinity", - "responses_api_deployment_check", - "session_affinity", - "encrypted_content_affinity", - } - ) def __init__( self, @@ -63,15 +99,6 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_responses_api_affinity = enable_responses_api_affinity self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} - for group, flags in self.model_group_affinity_config.items(): - unknown = set(flags) - self.VALID_FLAGS - if unknown: - verbose_router_logger.warning( - "DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s", - unknown, - group, - self.VALID_FLAGS, - ) def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -218,8 +245,13 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}" @classmethod - def get_session_affinity_cache_key(cls, model_group: str, session_id: str) -> str: - return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{session_id}" + def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str: + """Session pins are scoped by the caller's hashed API key so two callers reusing + the same client-supplied session_id cannot read or steer each other's pin. + `"unscoped"` covers direct Router usage with no authenticated caller, matching + the complexity router's own session pin key.""" + hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" + return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" @staticmethod def _get_user_key_from_metadata_dict(metadata: dict) -> str | None: @@ -278,6 +310,97 @@ class DeploymentAffinityCheck(CustomLogger): return session_id return None + @staticmethod + def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None: + """TTL from the session-affinity marker the Router stamps at pre-routing time + when an auto-router routed this request with session_affinity enabled. + Marker presence enables session pinning for this request only; anything that + is not a positive int is treated as absent.""" + for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): + ttl = metadata.get(SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY) + if isinstance(ttl, int) and not isinstance(ttl, bool) and ttl > 0: + return ttl + return None + + @staticmethod + def _pinned_model_id(stored: object) -> str | None: + """Deployment id held by a stored pin, for both the dict shape this writes and the + bare string older writers left behind. None when the value is neither.""" + if isinstance(stored, dict): + model_id: Final = stored.get("model_id") + return str(model_id) if model_id is not None else None + if isinstance(stored, str): + return stored + return None + + def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: + """The one owner of authoritative local pin writes: a plain set keeps a live + key's original expiry (`allow_ttl_override`), so the entry is replaced to make + the TTL real. Every local pin write goes through here so the redis-winner sync + and the pod-local claim can never disagree about expiry again.""" + self.cache.in_memory_cache.delete_cache(cache_key) + self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) + + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: + """First-writer-wins pin write: store `pin_value` only when the key is absent and + return the deployment id the key holds afterwards, so a caller learns whether it won + by comparing against its own id, and None when the stored value is one no reader can + interpret. Concurrent claimers converge on the + first write instead of the last. Re-claiming with the stored value refreshes its + TTL, the same keepalive the complexity router's model pin documents: an active + session must not lose its pin mid-conversation just because it outlives the + original write, so `session_affinity_ttl_seconds` bounds idle time, not total + session length. On Redis one Lua script does the get-or-set-or-refresh + atomically (same registration seam the rate limiters use) and the in-memory + tier is synchronized to the winner; without Redis, and whenever Redis is + unreachable, the pod-local check-and-set below stands in and is atomic because it + runs synchronously on the event loop. Degrading to a pod-local claim rather than + propagating the fault is what keeps same-pod stickiness through a Redis blip: the + caller only logs this result, so an escaping error would leave the session with no + pin at all and reshuffle every turn for the outage, which is worse than losing + cross-pod agreement. The redis tier is + resolved per call because the proxy attaches it after Router construction + (`Router._update_redis_cache`); the compiled script is cached per event loop + underneath the registration seam. + """ + redis_cache: Final = self.cache.redis_cache + if redis_cache is not None: + try: + claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) + raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds))) + decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw + if not isinstance(decoded, str): + return pin_value["model_id"] + try: + winner: object = json.loads(decoded) + except json.JSONDecodeError: + winner = decoded + self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds) + return self._pinned_model_id(winner) + except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the pod-local claim, never unpins + verbose_router_logger.debug( + "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e + ) + + return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds) + + def _claim_pin_in_memory( + self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int + ) -> str | None: + """Pod-local half of the claim, used when no Redis tier is attached and as the + fallback when the Redis claim fails. Mirrors the Lua script exactly, including + the keepalive: re-claiming with the stored value slides the idle window through + `_set_local_pin`. Both branches stay synchronous, hence atomic on the event + loop.""" + existing: Final = self.cache.in_memory_cache.get_cache(cache_key) + if existing is not None: + existing_model_id: Final = self._pinned_model_id(existing) + if existing_model_id == pin_value["model_id"]: + self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) + return existing_model_id + self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) + return pin_value["model_id"] + @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: for deployment in healthy_deployments: @@ -334,12 +457,21 @@ class DeploymentAffinityCheck(CustomLogger): if stable_model_map_key is None: return typed_healthy_deployments + session_affinity_active: Final = ( + enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None + ) + user_key: Final = ( + self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) + if (session_affinity_active or enable_user_key) + else None + ) + # 2) Session-id -> deployment affinity - if enable_session_id: + if session_affinity_active: session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs) if session_id is not None: session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=stable_model_map_key, session_id=session_id + model_group=stable_model_map_key, session_id=session_id, user_key=user_key ) session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key) @@ -371,7 +503,6 @@ class DeploymentAffinityCheck(CustomLogger): if not enable_user_key: return typed_healthy_deployments - user_key: Final = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if user_key is None: return typed_healthy_deployments @@ -438,18 +569,22 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id, ) = self._get_effective_flags(deployment_model_name) - if not enable_user_key and not enable_session_id: + marker_session_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs=kwargs) + session_affinity_active: Final = enable_session_id or marker_session_ttl is not None + + if not enable_user_key and not session_affinity_active: return None - user_key = None - if enable_user_key: - user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + user_key: Final = ( + self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + if (enable_user_key or session_affinity_active) + else None + ) + session_id: Final = ( + self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if session_affinity_active else None + ) - session_id = None - if enable_session_id: - session_id = self._get_session_id_from_request_kwargs(request_kwargs=kwargs) - - if user_key is None and session_id is None: + if not ((enable_user_key and user_key is not None) or session_id is not None): return None model_info = kwargs.get("model_info") @@ -473,22 +608,31 @@ class DeploymentAffinityCheck(CustomLogger): verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.") return None - if user_key is not None: + pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id)) + + if enable_user_key and user_key is not None: try: cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key) - await self.cache.async_set_cache( - cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), - ttl=self.ttl_seconds, - ) - - verbose_router_logger.debug( - "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", - deployment_model_name, - model_id, - self.ttl_seconds, - self._shorten_for_logs(user_key), + claimed_user_pin: Final = await self._claim_pin( + cache_key=cache_key, + pin_value=pin_value, + ttl_seconds=self.ttl_seconds, ) + if claimed_user_pin == pin_value["model_id"]: + verbose_router_logger.debug( + "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", + deployment_model_name, + model_id, + self.ttl_seconds, + self._shorten_for_logs(user_key), + ) + else: + verbose_router_logger.debug( + "DeploymentAffinityCheck: affinity pin already claimed model_map_key=%s existing=%s ours=%s", + deployment_model_name, + claimed_user_pin, + model_id, + ) except Exception as e: # Non-blocking: affinity is a best-effort optimization. verbose_router_logger.debug( @@ -500,21 +644,31 @@ class DeploymentAffinityCheck(CustomLogger): # Also persist Session-ID affinity if enabled and session-id is provided if session_id is not None: try: + session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds session_cache_key: Final = self.get_session_affinity_cache_key( - model_group=deployment_model_name, session_id=session_id + model_group=deployment_model_name, session_id=session_id, user_key=user_key ) - await self.cache.async_set_cache( - session_cache_key, - DeploymentAffinityCacheValue(model_id=str(model_id)), - ttl=self.ttl_seconds, - ) - verbose_router_logger.debug( - "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", - deployment_model_name, - model_id, - self.ttl_seconds, - session_id, + claimed_session_pin: Final = await self._claim_pin( + cache_key=session_cache_key, + pin_value=pin_value, + ttl_seconds=session_affinity_ttl, ) + if claimed_session_pin == pin_value["model_id"]: + verbose_router_logger.debug( + "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", + deployment_model_name, + model_id, + session_affinity_ttl, + session_id, + ) + else: + verbose_router_logger.debug( + "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s", + deployment_model_name, + claimed_session_pin, + model_id, + session_id, + ) except Exception as e: verbose_router_logger.debug( "DeploymentAffinityCheck: failed to set session affinity cache. model_map_key=%s error=%s", diff --git a/litellm/types/router.py b/litellm/types/router.py index 4280da08cbb..e166d844735 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -816,6 +816,7 @@ class PreRoutingHookResponse(BaseModel): model: str messages: list[dict[str, Any]] | None routing_decision: StandardLoggingRoutingDecision | None = None + session_affinity_ttl_seconds: int | None = None _PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6d6fd2e5507..22f9e6bb67a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -676,6 +676,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "applied_policies": ["spoofed-policy"], "policy_sources": {"spoofed-policy": "request"}, "routing_decision": {"cause": "forged", "routed_model": "spoofed"}, + "_session_deployment_affinity_ttl": 999999, "internal_call_origin": "autorouter_classifier", "_guardrail_pipelines": [{"name": "spoofed"}], "_pipeline_managed_guardrails": ["evaded"], @@ -719,6 +720,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "applied_policies", "policy_sources", "routing_decision", + "_session_deployment_affinity_ttl", "internal_call_origin", "_guardrail_pipelines", "_pipeline_managed_guardrails", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 356556f3563..94b6b68855b 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3414,6 +3414,103 @@ class TestSessionAffinity: def _request_kwargs(session_id: str) -> Dict: return {"metadata": {"session_id": session_id}} + @pytest.mark.asyncio + async def test_hook_response_carries_session_affinity_ttl_on_classify_and_pin_paths( + self, mock_router_instance, session_affinity_config + ): + """The hook response's session_affinity_ttl_seconds is what the Router stamps as + the deployment-affinity marker, so both the classify path (turn 1) and the + session-pin path (turn 2) must carry the configured TTL.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**session_affinity_config, "session_affinity_ttl_seconds": 321}, + ) + request_kwargs = self._request_kwargs("marker-session") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.session_affinity_ttl_seconds == 321 + assert second.session_affinity_ttl_seconds == 321 + + @pytest.mark.parametrize( + "session_affinity,deployment_affinity,plugins,tier_pinned,deployment_pinned", + [ + (False, False, False, False, False), + (False, True, False, False, True), + (True, False, False, True, True), + (True, True, False, True, True), + (False, True, True, False, False), + (True, True, True, False, False), + ], + ) + @pytest.mark.asyncio + async def test_tier_pin_and_deployment_pin_are_independently_gated( + self, + mock_router_instance, + basic_config, + session_affinity, + deployment_affinity, + plugins, + tier_pinned, + deployment_pinned, + ): + """deployment_affinity pins the deployment inside each routed group without pinning which + group the session routes to, so with session_affinity off the tier must still reclassify + on every turn while the marker the Router stamps is still emitted. Turn 1 classifies + REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one + does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline.""" + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "session_affinity": session_affinity, + "deployment_affinity": deployment_affinity, + **({"plugins": [_DummyPlugin()]} if plugins else {}), + }, + ) + request_kwargs = self._request_kwargs("matrix-session") + first = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + second = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert first.model == "o1-preview" + assert second.model == ("o1-preview" if tier_pinned else "gpt-4o-mini") + assert (first.session_affinity_ttl_seconds is not None) is deployment_pinned + assert (second.session_affinity_ttl_seconds is not None) is deployment_pinned + + @pytest.mark.asyncio + async def test_hook_response_has_no_session_affinity_ttl_when_disabled_or_plugins( + self, mock_router_instance, basic_config, session_affinity_config + ): + mock_router_instance.cache = DualCache() + disabled_router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "deployment_affinity": False}, + ) + plugin_router = ComplexityRouter( + model_name="test-router-plugins", + litellm_router_instance=mock_router_instance, + complexity_router_config={**session_affinity_config, "plugins": [_DummyPlugin()]}, + ) + disabled = await disabled_router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-off"), messages=self.SIMPLE_MESSAGE + ) + with_plugins = await plugin_router.async_pre_routing_hook( + model="test-model", request_kwargs=self._request_kwargs("s-plugins"), messages=self.SIMPLE_MESSAGE + ) + assert disabled.session_affinity_ttl_seconds is None + assert with_plugins.session_affinity_ttl_seconds is None + @pytest.mark.asyncio async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): """Regression: session_affinity defaults to False, so a shared session_id must NOT diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index e29adda3328..428eb0ceafd 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -465,8 +465,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id. """ - cache = AsyncMock() - cache.async_set_cache = AsyncMock() + cache = DualCache() callback = DeploymentAffinityCheck( cache=cache, @@ -489,11 +488,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): model_group="claude-sonnet-4-5@20250929", user_key="user-key-abc", ) - cache.async_set_cache.assert_called_once_with( - expected_cache_key, - {"model_id": "model-id-123"}, - ttl=123, - ) + assert await cache.async_get_cache(key=expected_cache_key) == {"model_id": "model-id-123"} @pytest.mark.asyncio diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index 0bcb0247aad..4053e6d118b 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,6 +10,7 @@ import json import litellm from litellm.caching.dual_cache import DualCache +from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -163,7 +164,7 @@ async def test_async_session_id_affinity_priority_over_user_key(): await callback.cache.async_set_cache( DeploymentAffinityCheck.get_session_affinity_cache_key( - "model_group", "session1" + "model_group", "session1", user_key="user1" ), {"model_id": "deployment-2"}, ) @@ -180,3 +181,439 @@ async def test_async_session_id_affinity_priority_over_user_key(): assert len(filtered) == 1 assert filtered[0]["model_info"]["id"] == "deployment-2" + + +MOCK_RESPONSES_API_RESPONSE = { + "id": "resp_mock-resp-456", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "azure/computer-use-preview", + "output": [], + "usage": { + "input_tokens": 5, + "output_tokens": 10, + "total_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} + + +def _smart_router(session_affinity=True, ttl_seconds=777, deployment_affinity=True): + return litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "target-group", + "complexity_router_config": { + "session_affinity": session_affinity, + "deployment_affinity": deployment_affinity, + "session_affinity_ttl_seconds": ttl_seconds, + "tiers": { + "SIMPLE": "target-group", + "MEDIUM": "target-group", + "COMPLEX": "target-group", + "REASONING": "target-group", + }, + }, + }, + }, + { + "model_name": "target-group", + "litellm_params": { + "model": "azure/computer-use-preview-1", + "api_key": "mock-api-key-1", + "api_version": "mock-api-version", + "api_base": "https://mock-endpoint-1.openai.azure.com", + }, + "model_info": {"id": "deployment-1", "base_model": "computer-use-preview"}, + }, + { + "model_name": "target-group", + "litellm_params": { + "model": "azure/computer-use-preview-2", + "api_key": "mock-api-key-2", + "api_version": "mock-api-version-2", + "api_base": "https://mock-endpoint-2.openai.azure.com", + }, + "model_info": {"id": "deployment-2", "base_model": "computer-use-preview"}, + }, + ], + ) + + +def _session_pin_key(session_id, user_key): + return DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="target-group", session_id=session_id, user_key=user_key + ) + + +def _cleanup_router_callbacks(router): + for callback in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + + +async def _one_turn(router, model, session_id, key_hash): + """One request with the shuffle forced to deployment-1, so any other landing + deployment can only come from a pin read.""" + with ( + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post, + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=lambda seq: seq[0], + ), + ): + mock_post.return_value = MockResponse(MOCK_RESPONSES_API_RESPONSE, 200) + response = await router.aresponses( + model=model, + input=f"turn for {session_id} {key_hash}", + litellm_metadata={"session_id": session_id, "user_api_key_hash": key_hash}, + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it(): + """Turn 1 persists a key-scoped deployment pin; a pin seeded to the deployment + the shuffle would never pick is then followed, proving the read path.""" + router = _smart_router() + try: + served = await _one_turn(router, "smart-router", "write-session", "key-1") + assert await router.cache.async_get_cache(key=_session_pin_key("write-session", "key-1")) == { + "model_id": served + } + assert await router.cache.async_get_cache(key=_session_pin_key("write-session", None)) is None + + await router.cache.async_set_cache( + key=_session_pin_key("read-session", "key-1"), value={"model_id": "deployment-2"} + ) + assert await _one_turn(router, "smart-router", "read-session", "key-1") == "deployment-2" + finally: + _cleanup_router_callbacks(router) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,key_hash", + [ + ("target-group", "key-1"), + ("smart-router", "key-2"), + ], + ids=["direct-group-call", "different-api-key"], +) +async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash): + """The pin binds (auto-routed request, api key, session): a direct call to the + group and a different key reusing the session id must both ignore it.""" + router = _smart_router() + try: + await router.cache.async_set_cache( + key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"} + ) + assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1" + finally: + _cleanup_router_callbacks(router) + + +@pytest.mark.asyncio +async def test_marker_write_uses_marker_ttl_and_writes_only_the_session_pin(): + """The write hook honors the marker's TTL over the callback default and writes + no user-key entry when only session affinity is active.""" + import time as time_module + + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "deployment-1"}, + "metadata": { + "deployment_model_name": "target-group", + "session_id": "ttl-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777, + }, + }, + call_type=None, + ) + + session_key = _session_pin_key("ttl-session", "key-1") + assert cache.in_memory_cache.cache_dict == {session_key: {"model_id": "deployment-1"}} + assert cache.in_memory_cache.ttl_dict[session_key] == pytest.approx(time_module.time() + 777, abs=5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_marker", ["777", True, -5, 0, None]) +async def test_malformed_marker_values_do_not_enable_session_affinity(bad_marker): + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "deployment-1"}, + "metadata": { + "deployment_model_name": "target-group", + "session_id": "bad-marker-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: bad_marker, + }, + }, + call_type=None, + ) + + assert cache.in_memory_cache.cache_dict == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enable_user_key", [False, True], ids=["session-pin", "user-key-pin"]) +async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_key): + """Two overlapping first requests select different deployments before either + write lands. Pins are first-writer-wins claims, so the second write must leave + the stored pin unchanged instead of flipping it.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=enable_user_key, + enable_responses_api_affinity=False, + ) + + def racing_kwargs(deployment_id): + metadata = {"deployment_model_name": "target-group", "user_api_key_hash": "key-1"} + if not enable_user_key: + metadata["session_id"] = "racing-session" + metadata[SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] = 777 + return {"model_info": {"id": deployment_id}, "metadata": metadata} + + await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-1"), call_type=None) + await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-2"), call_type=None) + + pinned_key = ( + DeploymentAffinityCheck.get_affinity_cache_key(model_group="target-group", user_key="key-1") + if enable_user_key + else _session_pin_key("racing-session", "key-1") + ) + assert await cache.async_get_cache(key=pinned_key) == {"model_id": "deployment-1"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stored_pin", + [{"model_id": "deployment-1"}, "deployment-1"], + ids=["dict-pin", "legacy-string-pin"], +) +async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployment(stored_pin): + """The pod-local claim mirrors the Lua keepalive: the winning deployment's + re-claim extends the pin's expiry, a losing deployment's claim touches neither + the value nor the expiry, so no-Redis setups keep stickiness across an active + session and ttl bounds idle time there too. Sameness is judged on the pinned + model id, so a legacy string pin written by the Redis branch slides the same.""" + import time as time_module + + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + pin_key = _session_pin_key("slide-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10) + first_expiry = cache.in_memory_cache.ttl_dict[pin_key] + + reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777) + assert reclaimed == "deployment-1" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry + + lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10) + assert lost == "deployment-1" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + + +@pytest.mark.asyncio +async def test_claim_pin_uses_redis_attached_after_construction(): + """The proxy attaches Redis via Router._update_redis_cache after the Router (and + this callback) are built. The claim must resolve the redis tier per call, or pins + silently stay pod-local and cross-pod first-writer-wins is lost.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + captured = {} + + async def fake_runner(keys, args, client=None): + captured["keys"] = keys + captured["args"] = args + return b'{"model_id": "other-pod-winner"}' + + late_redis = MagicMock() + late_redis.async_register_script = MagicMock(return_value=fake_runner) + cache.redis_cache = late_redis + + import time as time_module + + pin_key = _session_pin_key("late-redis-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) + + claimed = await callback._claim_pin( + cache_key=pin_key, + pin_value={"model_id": "our-deployment"}, + ttl_seconds=777, + ) + + assert claimed == "other-pod-winner" + assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) + assert captured["keys"] == (pin_key,) + assert captured["args"] == ('{"model_id": "our-deployment"}', 777) + assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == { + "model_id": "other-pod-winner" + } + + +@pytest.mark.asyncio +async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): + """A Redis outage must cost cross-pod agreement, never same-pod stickiness. The write + hook only logs this result, so an escaping error would leave the session unpinned and + reshuffle every turn for the whole outage. DualCache's write path, which this claim + replaced, wrote the in-memory tier before ever touching Redis.""" + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + + async def exploding_runner(keys, args, client=None): + raise ConnectionError("redis is down") + + down_redis = MagicMock() + down_redis.async_register_script = MagicMock(return_value=exploding_runner) + cache.redis_cache = down_redis + + key = _session_pin_key("outage-session", "key-1") + claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777) + + assert claimed == "our-deployment" + assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"} + + second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777) + assert second == "our-deployment" + + +@pytest.mark.asyncio +async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups(): + """Wildcard deployments keep the literal pattern as model_name on both the read + path and the write path, so the marker-gated pin round-trips through one key.""" + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + ) + request_kwargs = { + "model_info": {"id": "wild-deployment-2"}, + "metadata": { + "deployment_model_name": "openai/*", + "session_id": "wild-session", + "user_api_key_hash": "key-1", + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777, + }, + } + + await callback.async_pre_call_deployment_hook(kwargs=request_kwargs, call_type=None) + filtered = await callback.async_filter_deployments( + model="openai/gpt-4o", + healthy_deployments=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": f"wild-deployment-{i}"}, + } + for i in (1, 2) + ], + messages=[], + request_kwargs=request_kwargs, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["wild-deployment-2"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,session_affinity,deployment_affinity,expect_marker", + [ + ("smart-router", False, True, True), + ("smart-router", True, False, True), + ("smart-router", False, False, False), + ("target-group", False, True, False), + ], + ids=[ + "deployment-affinity-stamps", + "session-affinity-implies-deployment-pin", + "both-off-no-stamp", + "non-auto-routed-clears", + ], +) +async def test_pre_routing_hook_stamps_or_clears_the_marker_per_attempt( + model, session_affinity, deployment_affinity, expect_marker +): + """Every routing attempt writes or clears the marker, so a fallback from an + auto-routed group to a plain group cannot carry a stale marker. session_affinity + implies the deployment pin: a session frozen onto one group must not re-shuffle + across that group's deployments.""" + router = _smart_router(session_affinity=session_affinity, deployment_affinity=deployment_affinity) + try: + request_kwargs = { + "metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111}, + "litellm_metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111}, + } + await router.async_pre_routing_hook( + model=model, + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "Hello"}], + ) + if expect_marker: + assert request_kwargs["litellm_metadata"][SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] == 777 + else: + assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["metadata"] + assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["litellm_metadata"] + finally: + _cleanup_router_callbacks(router) + + +def test_complexity_router_with_deployment_affinity_registers_affinity_callback(): + enabled = _smart_router() + session_only = _smart_router(session_affinity=True, deployment_affinity=False) + disabled = _smart_router(session_affinity=False, deployment_affinity=False) + try: + assert [ + (cb.enable_user_key_affinity, cb.enable_responses_api_affinity, cb.enable_session_id_affinity) + for cb in enabled.optional_callbacks or [] + if isinstance(cb, DeploymentAffinityCheck) + ] == [(False, False, False)] + assert any(isinstance(cb, DeploymentAffinityCheck) for cb in session_only.optional_callbacks or []) + assert not any(isinstance(cb, DeploymentAffinityCheck) for cb in disabled.optional_callbacks or []) + finally: + _cleanup_router_callbacks(enabled) + _cleanup_router_callbacks(session_only) + _cleanup_router_callbacks(disabled) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index da2d78edb73..67fa827a8e4 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7550,3 +7550,68 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): assert capture.messages, "the fallback failure path did not log at ERROR" assert huge_message not in "".join(capture.messages) assert max(len(message) for message in capture.messages) < 5_000 +def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets(): + request_kwargs = {"metadata": {}} + litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7) + assert request_kwargs["metadata"]["probe"] == 7 + + stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}} + litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None) + assert "probe" not in stale_kwargs["metadata"] + assert "probe" not in stale_kwargs["litellm_metadata"] + + +@pytest.mark.parametrize( + "complexity_router_config,expect_callback", + [ + ({"tiers": {"SIMPLE": "gpt-4o"}}, True), + ({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False}, False), + ({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False, "session_affinity": True}, True), + ], +) +def test_complexity_router_registers_affinity_callback_for_deployment_pin(complexity_router_config, expect_callback): + """The marker the complexity router stamps is inert unless a DeploymentAffinityCheck is + registered to read it, so deployment_affinity has to pull the callback in, and its default-on + means a bare config registers one. Opting out must skip the callback entirely rather than + register a filter that can never fire, including when session_affinity is on, since the two + pins are independent.""" + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + { + "model_name": "my-complexity-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": complexity_router_config, + }, + }, + ] + ) + try: + registered = any(isinstance(cb, DeploymentAffinityCheck) for cb in router.optional_callbacks or []) + assert registered is expect_callback + finally: + for cb in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(cb) + + +def test_ensure_deployment_affinity_callback_is_idempotent(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = litellm.Router(model_list=[]) + try: + router._ensure_deployment_affinity_callback() + router._ensure_deployment_affinity_callback() + affinity_callbacks = [ + cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) + ] + assert len(affinity_callbacks) == 1 + finally: + for cb in router.optional_callbacks or []: + litellm.logging_callback_manager.remove_callback_from_all_lists(cb) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3a670bc7345..d621e85f09b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16760 + "limit": 16758 }, "LIT011": { "limit": 5598 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a75c23da1cf..b361a16a0e6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31695,6 +31695,12 @@ export interface components { * @description Default model to use if tier cannot be determined */ default_model?: string | null; + /** + * Deployment Affinity + * @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is. + * @default true + */ + deployment_affinity: boolean; /** * Dimension Weights * @description Weights for each scoring dimension @@ -31752,13 +31758,13 @@ export interface components { semantic_keyword_matching: boolean; /** * Session Affinity - * @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. + * @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. Always implies the deployment pin regardless of deployment_affinity: the session sticks to one deployment of the pinned model, since freezing the model while re-shuffling its deployments would still go cache-cold. * @default false */ session_affinity: boolean; /** * Session Affinity Ttl Seconds - * @description TTL for the session affinity pin; refreshed on every cache hit + * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length * @default 3600 */ session_affinity_ttl_seconds: number;