From 6ed1c6b420e197b3327cf37e9d54f3a6cd9e17fd Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Fri, 26 Jun 2026 13:25:28 +0000 Subject: [PATCH 001/121] fix(deps): bump langgraph-checkpoint to 4.1.1 to resolve OSV vulnerability --- uv.lock | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index cac1696bf34..8c9e20eeb21 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-20T23:16:25.061268Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" [manifest] @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.0" +version = "4.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] [[package]] @@ -3281,6 +3281,7 @@ dependencies = [ { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, + { name = "langgraph-checkpoint" }, { name = "openai" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -3505,6 +3506,7 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, + { name = "langgraph-checkpoint", specifier = "==4.1.1" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, From 3f5186f9afcced38bdcd8a6095c11e6a8e206627 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Fri, 26 Jun 2026 13:39:29 +0000 Subject: [PATCH 002/121] fix(ocr): use defensive getattr in load_rust_ocr --- litellm/ocr/rust_bridge.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 1e3312c1473..253b35cb689 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -97,7 +97,7 @@ def load_rust_ocr() -> RustOcr | None: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, litellm_python_bridge.ocr) + return cast(RustOcr, getattr(litellm_python_bridge, "ocr", None)) def load_rust_aocr() -> RustAocr | None: From fc36825dfd68ab3e5b142a401810de84a452fe62 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Tue, 21 Jul 2026 06:16:02 +0000 Subject: [PATCH 003/121] fix(batches): support AWS Bedrock batch cancellation via StopModelInvocationJob (#33986) --- litellm/batches/main.py | 9 +- litellm/llms/bedrock/batches/handler.py | 162 +++++++++--------------- tests/test_bedrock_cancel_batch.py | 55 ++++++++ 3 files changed, 124 insertions(+), 102 deletions(-) create mode 100644 tests/test_bedrock_cancel_batch.py diff --git a/litellm/batches/main.py b/litellm/batches/main.py index f124882b5a4..0c2cf15d385 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -1087,9 +1087,16 @@ def cancel_batch( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "bedrock": + from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler + + response = BedrockBatchesHandler.cancel_batch( + batch_id=batch_id, + **kwargs, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format( + message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'bedrock' are supported.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c071f331337..55236eae525 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -6,9 +6,7 @@ from openai.types.batch import Metadata as OpenAIBatchMetadata from litellm.types.utils import LiteLLMBatch -# AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses. -# Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response` -# so create / retrieve return consistent statuses. +# AWS Bedrock model-invocation-job statuses -> OpenAI Batch statuses. _BEDROCK_MIJ_STATUS_TO_OPENAI = { "Submitted": "validating", "Validating": "validating", @@ -44,17 +42,6 @@ def _extract_job_id_from_arn(arn: str) -> Optional[str]: def _predict_output_file_uri( output_prefix: str, input_uri: str, job_id: Optional[str] ) -> Optional[str]: - """ - Compute the deterministic per-job result file URI Bedrock writes to. - - Bedrock lays results out as:: - - //.out - - We compute it client-side so OpenAI-style ``client.files.content(output_file_id)`` - works without an extra S3 ``ListObjectsV2`` round-trip. Returns ``None`` if we - don't have enough info; callers should fall back to the bare prefix. - """ if not output_prefix or not input_uri or not job_id: return None if not output_prefix.endswith("/"): @@ -76,40 +63,76 @@ def _to_epoch(value: Any) -> Optional[int]: class BedrockBatchesHandler: - """ - Handler for Bedrock Batches. + """Handler for Bedrock Batches.""" - Specific providers/models needed some special handling. + @staticmethod + def cancel_batch( + batch_id: str, + aws_region_name: Optional[str] = None, + logging_obj=None, + **kwargs, + ) -> "LiteLLMBatch": + """ + Cancel an AWS Bedrock batch model invocation job using StopModelInvocationJob. + """ + try: + import boto3 + from botocore.exceptions import ClientError + except ImportError as exc: + raise ImportError( + "Missing boto3/botocore to call bedrock. Run 'pip install boto3'." + ) from exc - E.g. Twelve Labs Embedding Async Invoke - """ + region = ( + aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" + ) + + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + creds = BedrockBatchesConfig().get_credentials( + aws_access_key_id=kwargs.get("aws_access_key_id"), + aws_secret_access_key=kwargs.get("aws_secret_access_key"), + aws_session_token=kwargs.get("aws_session_token"), + aws_region_name=region, + aws_session_name=kwargs.get("aws_session_name"), + aws_profile_name=kwargs.get("aws_profile_name"), + aws_role_name=kwargs.get("aws_role_name"), + aws_web_identity_token=kwargs.get("aws_web_identity_token"), + aws_sts_endpoint=kwargs.get("aws_sts_endpoint"), + aws_external_id=kwargs.get("aws_external_id"), + ) + + client = boto3.client( + "bedrock", + region_name=region, + aws_access_key_id=creds.access_key, + aws_secret_access_key=creds.secret_key, + aws_session_token=creds.token, + ) + + try: + client.stop_model_invocation_job(jobIdentifier=batch_id) + except ClientError as e: + # Idempotency: if job is already Stopping/Stopped/Completed, swallow ValidationException + if e.response.get("Error", {}).get("Code") != "ValidationException": + raise e + + return BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=batch_id, + aws_region_name=region, + logging_obj=logging_obj, + **kwargs, + ) @staticmethod def _handle_async_invoke_status( batch_id: str, aws_region_name: str, logging_obj=None, **kwargs ) -> "LiteLLMBatch": - """ - Handle async invoke status check for AWS Bedrock. - - This is for Twelve Labs Embedding Async Invoke. - - Args: - batch_id: The async invoke ARN - aws_region_name: AWS region name - **kwargs: Additional parameters - - Returns: - dict: Status information including status, output_file_id (S3 URL), etc. - """ import asyncio - from litellm.llms.bedrock.embed.embedding import BedrockEmbedding async def _async_get_status(): - # Create embedding handler instance embedding_handler = BedrockEmbedding() - - # Get the status of the async invoke job status_response = await embedding_handler._get_async_invoke_status( invocation_arn=batch_id, aws_region_name=aws_region_name, @@ -117,18 +140,13 @@ class BedrockBatchesHandler: **kwargs, ) - # Transform response to a LiteLLMBatch object - from litellm.types.utils import LiteLLMBatch - openai_batch_metadata: OpenAIBatchMetadata = { - "output_file_id": status_response["outputDataConfig"][ - "s3OutputDataConfig" - ]["s3Uri"], + "output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"], "failure_message": status_response.get("failureMessage") or "", "model_arn": status_response["modelArn"], } - result = LiteLLMBatch( + return LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=status_response["status"], @@ -151,10 +169,6 @@ class BedrockBatchesHandler: input_file_id="", ) - return result - - # Since this function is called from within an async context via run_in_executor, - # we need to create a new event loop in a thread to avoid conflicts import concurrent.futures def run_in_thread(): @@ -176,37 +190,6 @@ class BedrockBatchesHandler: logging_obj=None, **kwargs, ) -> "LiteLLMBatch": - """ - Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch - inference jobs (the ARN type returned by ``CreateModelInvocationJob``). - - ``CreateModelInvocationJob`` lives on the Bedrock **control plane** - (``bedrock..amazonaws.com``), distinct from the data-plane - ``bedrock-runtime`` endpoint that serves Twelve Labs async-invoke ARNs. - The two ARN families therefore can't share a handler — see - ``litellm/batches/main.py`` for the dispatch. - - Args: - batch_id: A ``arn:aws:bedrock:::model-invocation-job/`` - ARN (or just the trailing job id; both are accepted by - ``GetModelInvocationJob``). - aws_region_name: Region for the boto3 ``bedrock`` client. If omitted, - we fall back to parsing the region out of ``batch_id`` itself. - logging_obj: Optional litellm logging object. - **kwargs: Optional AWS credential overrides - (``aws_access_key_id``, ``aws_secret_access_key``, - ``aws_session_token``, ``aws_profile_name``, - ``aws_role_name``, ``aws_session_name``, - ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``). Unknown keys are ignored. - - Returns: - ``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that - ``request_counts`` is always ``(0, 0, 0)`` because - ``GetModelInvocationJob`` does not surface per-record counts; - callers that need accurate counts should parse - ``manifest.json.out`` from the output S3 prefix. - """ try: import boto3 except ImportError as exc: @@ -214,15 +197,10 @@ class BedrockBatchesHandler: "Missing boto3 to call bedrock. Run 'pip install boto3'." ) from exc - # Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default). region = ( aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1" ) - # Resolve credentials through the same path the rest of the bedrock - # provider uses, so model_list / env / role-assumption configs are - # honored. We instantiate BedrockBatchesConfig (which extends - # BaseAWSLLM) lazily to avoid a circular import at module load. from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig creds = BedrockBatchesConfig().get_credentials( @@ -247,10 +225,6 @@ class BedrockBatchesHandler: ) if logging_obj is not None: - # Use the bare job id in the logged URL so we don't double up the - # `model-invocation-job/` segment when `batch_id` is a full ARN. - # `GetModelInvocationJob` accepts either form, but only the bare id - # produces a sensible-looking URL in logs. url_path_id = _extract_job_id_from_arn(batch_id) or batch_id logging_obj.pre_call( input=batch_id, @@ -291,26 +265,12 @@ class BedrockBatchesHandler: .get("s3Uri", "") ) - # Bedrock returns the output *prefix* the user supplied at job creation. - # Actual results land at //.out — we - # surface that single-file URI as `output_file_id` so the OpenAI-style - # download flow works without an extra S3 listing call. We deliberately - # do NOT fall back to the bare prefix when prediction fails: a prefix - # is not a downloadable object, so handing it back as `output_file_id` - # would reproduce the very NoSuchKey bug this handler exists to fix. - # The bare prefix is preserved in metadata for callers that want the - # `manifest.json.out` or want to do their own listing. job_arn = response.get("jobArn", batch_id) job_id = _extract_job_id_from_arn(job_arn) output_file_uri = _predict_output_file_uri(output_prefix, input_uri, job_id) completed_at = _to_epoch(response.get("endTime")) - # Note: metadata uses "" (not None) for unknown URIs to satisfy the - # OpenAI Batch metadata schema, which is `dict[str, str]`. The - # `output_file_id` field on the LiteLLMBatch itself does carry None - # correctly (see below), so callers should branch on that, not on - # `metadata["output_file_uri"]`. openai_batch_metadata: OpenAIBatchMetadata = { "model_arn": response.get("modelId", ""), "job_arn": job_arn, diff --git a/tests/test_bedrock_cancel_batch.py b/tests/test_bedrock_cancel_batch.py new file mode 100644 index 00000000000..e52cc2c447e --- /dev/null +++ b/tests/test_bedrock_cancel_batch.py @@ -0,0 +1,55 @@ +from unittest.mock import MagicMock, patch +import pytest + +import litellm +from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler + + +@patch("boto3.client") +def test_bedrock_cancel_batch_handler(mock_boto_client): + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + mock_client_instance.stop_model_invocation_job.return_value = {} + mock_client_instance.get_model_invocation_job.return_value = { + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + "status": "Stopping", + "submitTime": 1700000000, + "lastModifiedTime": 1700000100, + } + + res = BedrockBatchesHandler.cancel_batch( + batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + aws_region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + mock_client_instance.stop_model_invocation_job.assert_called_once_with( + jobIdentifier="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id" + ) + assert res.status == "cancelling" + + +@patch("boto3.client") +def test_litellm_cancel_batch_bedrock_dispatcher(mock_boto_client): + mock_client_instance = MagicMock() + mock_boto_client.return_value = mock_client_instance + + mock_client_instance.stop_model_invocation_job.return_value = {} + mock_client_instance.get_model_invocation_job.return_value = { + "jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + "status": "Stopped", + "submitTime": 1700000000, + "lastModifiedTime": 1700000100, + } + + res = litellm.cancel_batch( + batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id", + custom_llm_provider="bedrock", + aws_region_name="us-east-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + + assert res.status == "cancelled" From 163ab6e34b5e1b59675d850f86e0f96cdbce4d64 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Tue, 21 Jul 2026 06:24:17 +0000 Subject: [PATCH 004/121] fix(batches): refine bedrock cancel_batch type hints and validation error handling --- litellm/batches/main.py | 2 +- litellm/llms/bedrock/batches/handler.py | 9 +++++++-- litellm/ocr/rust_bridge.py | 2 +- uv.lock | 10 ++++------ 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 0c2cf15d385..9d4cb29e926 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -944,7 +944,7 @@ async def acancel_batch( def cancel_batch( batch_id: str, model: Optional[str] = None, - custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai", + custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai", "bedrock"], str] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 55236eae525..c710bcfdafa 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -113,8 +113,13 @@ class BedrockBatchesHandler: try: client.stop_model_invocation_job(jobIdentifier=batch_id) except ClientError as e: - # Idempotency: if job is already Stopping/Stopped/Completed, swallow ValidationException - if e.response.get("Error", {}).get("Code") != "ValidationException": + error_code = e.response.get("Error", {}).get("Code") + error_msg = e.response.get("Error", {}).get("Message", "").lower() + if error_code == "ValidationException" and any( + term in error_msg for term in ["stop", "terminal", "completed", "already"] + ): + pass + else: raise e return BedrockBatchesHandler._handle_model_invocation_job_status( diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 253b35cb689..1e3312c1473 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -97,7 +97,7 @@ def load_rust_ocr() -> RustOcr | None: import litellm_python_bridge except ImportError: return None - return cast(RustOcr, getattr(litellm_python_bridge, "ocr", None)) + return cast(RustOcr, litellm_python_bridge.ocr) def load_rust_aocr() -> RustAocr | None: diff --git a/uv.lock b/uv.lock index 8c9e20eeb21..cac1696bf34 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-06-20T23:16:25.061268Z" exclude-newer-span = "P3D" [manifest] @@ -3160,15 +3160,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.1" +version = "4.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, + { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, ] [[package]] @@ -3281,7 +3281,6 @@ dependencies = [ { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, - { name = "langgraph-checkpoint" }, { name = "openai" }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -3506,7 +3505,6 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4.0" }, { name = "jsonschema", specifier = ">=4.0.0,<5.0" }, { name = "langfuse", marker = "extra == 'proxy-runtime'", specifier = ">=2.59.7,<3.0" }, - { name = "langgraph-checkpoint", specifier = "==4.1.1" }, { name = "litellm-enterprise", marker = "extra == 'proxy'", editable = "enterprise" }, { name = "litellm-proxy-extras", marker = "extra == 'proxy'", editable = "litellm-proxy-extras" }, { name = "llm-sandbox", marker = "extra == 'proxy-runtime'", specifier = ">=0.3.39,<1.0" }, From ebf6167d8acf48499e294ecf3a7642b4112913eb Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Mon, 3 Aug 2026 16:50:58 -0400 Subject: [PATCH 005/121] fix(anthropic): stop emitting empty thinking blocks on the Responses adapter OpenAI emits a reasoning output item on every reasoning turn, but only emits reasoning_summary_text deltas when a summary was requested and actually produced. The Anthropic /v1/messages Responses stream adapter opened the thinking content block eagerly on response.output_item.added, so a summary-less reasoning item surfaced as {"type": "thinking", "thinking": ""}. Clients persist that in their session transcript and replay it on the next turn; an Anthropic model then rejects the request with "each thinking block must contain thinking", which is what users hit when a resumed session falls back to the default Anthropic model. Open the thinking block on the first non-empty summary delta instead, and only emit content_block_stop for items that actually have an open block. --- .../responses_adapters/streaming_iterator.py | 83 +++++++------------ ...t_responses_adapters_streaming_iterator.py | 79 +++++++++++++++++- 2 files changed, 106 insertions(+), 56 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index f12dd979338..c588e791cd9 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -68,6 +68,19 @@ class AnthropicResponsesStreamWrapper: self._current_block_index += 1 return self._current_block_index + def _open_block(self, item_id: str | None, content_block: dict[str, Any]) -> int: + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": content_block, + } + ) + return block_idx + def _process_event(self, event: Any) -> None: """Convert one Responses API event into zero or more Anthropic chunks queued for emission.""" event_type = getattr(event, "type", None) @@ -93,47 +106,22 @@ class AnthropicResponsesStreamWrapper: item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item_type == "message": - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - } - ) + self._open_block(item_id, {"type": "text", "text": ""}) elif item_type == "function_call": call_id: Final = ( getattr(item, "call_id", None) or (item.get("call_id") if isinstance(item, dict) else None) or "" ) name = getattr(item, "name", None) or (item.get("name") if isinstance(item, dict) else None) or "" - block_idx = self._next_block_index() if item_id: - self._item_id_to_block_index[item_id] = block_idx self._pending_tool_ids[item_id] = call_id - self._chunk_queue.append( + self._open_block( + item_id, { - "type": "content_block_start", - "index": block_idx, - "content_block": { - "type": "tool_use", - "id": call_id, - "name": name, - "input": {}, - }, - } - ) - elif item_type == "reasoning": - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "thinking", "thinking": ""}, - } + "type": "tool_use", + "id": call_id, + "name": name, + "input": {}, + }, ) return @@ -146,16 +134,7 @@ class AnthropicResponsesStreamWrapper: # Some providers (e.g. LMStudio) skip response.output_item.added, # so no text block is open yet; synthesize content_block_start # instead of emitting a delta with index -1 - block_idx = self._next_block_index() - if item_id: - self._item_id_to_block_index[item_id] = block_idx - self._chunk_queue.append( - { - "type": "content_block_start", - "index": block_idx, - "content_block": {"type": "text", "text": ""}, - } - ) + block_idx = self._open_block(item_id, {"type": "text", "text": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -169,11 +148,11 @@ class AnthropicResponsesStreamWrapper: if event_type == "response.reasoning_summary_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) - if item_id - else self._current_block_index - ) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + if not delta: + return + block_idx = self._open_block(item_id, {"type": "thinking", "thinking": ""}) self._chunk_queue.append( { "type": "content_block_delta", @@ -207,11 +186,9 @@ class AnthropicResponsesStreamWrapper: item_id = ( getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) if item else None ) - block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) - if item_id - else self._current_block_index - ) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + return self._chunk_queue.append( { "type": "content_block_stop", diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 73b58e71009..b1ae865fde1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -76,6 +76,78 @@ class TestProcessEventResponseCreatedGuard: assert len(message_starts) == 1 +class TestReasoningItemWithoutSummaryText: + """Regression: a reasoning item whose summary never produces text must not + surface as a thinking content block. + + OpenAI emits ``response.output_item.added`` with ``type: "reasoning"`` on + every reasoning turn, but only emits + ``response.reasoning_summary_text.delta`` when a summary was requested and + the model actually produced one. Eagerly opening the block on + ``output_item.added`` left ``{"type": "thinking", "thinking": ""}`` in the + assistant turn, which clients persist in their session transcript. Replaying + that transcript against an Anthropic model (what ``claude --resume`` does + once the resumed session falls back to the default Anthropic model) fails + with:: + + 400 invalid_request_error - messages.2.content.0.thinking: + each thinking block must contain thinking + + So the thinking block is opened on the first non-empty summary delta. + """ + + @staticmethod + def _gpt_turn(reasoning_summary_deltas: list) -> list: + return [ + {"type": "response.created"}, + {"type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}}, + *( + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": delta} + for delta in reasoning_summary_deltas + ), + {"type": "response.output_item.done", "item": {"type": "reasoning", "id": "rs_1"}}, + {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hello"}, + {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, + ] + + def test_reasoning_without_summary_emits_no_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=[])) + + assert not [ + c for c in chunks if c["type"] == "content_block_start" and c["content_block"]["type"] == "thinking" + ] + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ("content_block_stop", 0), + ] + assert chunks[1]["content_block"] == {"type": "text", "text": ""} + + def test_reasoning_with_only_empty_summary_deltas_emits_no_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["", ""])) + + assert not [c for c in chunks if c["type"] == "content_block_delta" and c["delta"]["type"] == "thinking_delta"] + assert not [ + c for c in chunks if c["type"] == "content_block_start" and c["content_block"]["type"] == "thinking" + ] + + def test_reasoning_with_summary_text_still_emits_a_thinking_block(self): + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["Weigh", "ing options"])) + + assert [(c["type"], c.get("index")) for c in chunks[1:]] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ("content_block_delta", 0), + ("content_block_stop", 0), + ("content_block_start", 1), + ("content_block_delta", 1), + ("content_block_stop", 1), + ] + assert chunks[1]["content_block"] == {"type": "thinking", "thinking": ""} + assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" + + class TestProcessEventTextDeltaWithoutOutputItemAdded: """Streams that skip response.output_item.added (e.g. LMStudio) must still open a text block before any delta and never emit index -1.""" @@ -110,12 +182,13 @@ class TestProcessEventTextDeltaWithoutOutputItemAdded: "type": "response.output_item.added", "item": {"type": "reasoning", "id": "rs_1"}, }, + {"type": "response.reasoning_summary_text.delta", "item_id": "rs_1", "delta": "hm"}, {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, ] ) - assert chunks[1]["type"] == "content_block_start" - assert chunks[1]["content_block"] == {"type": "text", "text": ""} - assert [c["index"] for c in chunks[1:]] == [1, 1] + assert chunks[2]["type"] == "content_block_start" + assert chunks[2]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks[2:]] == [1, 1] def test_process_event_registered_item_id_does_not_synthesize_start(self): chunks = _process_all( From 1d8a642e0683e13be122c532436e0919d6c540f5 Mon Sep 17 00:00:00 2001 From: heathriel Date: Wed, 22 Jul 2026 08:41:01 -0700 Subject: [PATCH 006/121] fix(fireworks_ai): support router slugs via routers/ prefix Bare fireworks_ai/ only resolved to accounts/fireworks/models/, so Fireworks routers (served at accounts/fireworks/routers/, e.g. glm-latest and firerouter) could not be reached without passing the full resource id. Add a shared resolve_fireworks_resource_name helper that maps an explicit routers/ or models/ segment to the right resource path, keeps the existing -fast router heuristic, and defaults bare slugs to models/ for backward compatibility. Wire it into both the chat and text-completion transforms, which had drifted (completion lacked router handling entirely) --- .../llms/fireworks_ai/chat/transformation.py | 18 ++++---- litellm/llms/fireworks_ai/common_utils.py | 11 +++++ .../fireworks_ai/completion/transformation.py | 7 +-- .../test_fireworks_ai_chat_transformation.py | 43 ++++++++++++++++++ ..._fireworks_ai_completion_transformation.py | 34 ++++++++++++++ .../test_fireworks_ai_common_utils.py | 45 +++++++++++++++++++ type-discipline-budget.json | 2 +- 7 files changed, 146 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py create mode 100644 tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index a796aa47b70..26f0caefacd 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -39,7 +39,11 @@ from ...openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, ) -from ..common_utils import FireworksAIException, FireworksAIMixin +from ..common_utils import ( + FireworksAIException, + FireworksAIMixin, + resolve_fireworks_resource_name, +) def _extract_fireworks_hidden_params(payload: dict) -> dict: @@ -459,12 +463,10 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): litellm_params: dict, headers: dict, ) -> dict: - if not model.startswith("accounts/") and "#" not in model: - if model.endswith("-fast"): - model = f"accounts/fireworks/routers/{model}" - else: - model = f"accounts/fireworks/models/{model}" - messages = self._transform_messages_helper(messages=messages, model=model, litellm_params=litellm_params) + resolved_model: Final = resolve_fireworks_resource_name(model) + messages = self._transform_messages_helper( + messages=messages, model=resolved_model, litellm_params=litellm_params + ) if "tools" in optional_params and optional_params["tools"] is not None: tools: Final = self._transform_tools(tools=optional_params["tools"]) optional_params["tools"] = tools @@ -478,7 +480,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): "include_usage": True, } return super().transform_request( - model=model, + model=resolved_model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 143dd151027..e07e7a26f9e 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -29,6 +29,17 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: return None +def resolve_fireworks_resource_name(model: str) -> str: + stripped: Final = model.removeprefix("fireworks_ai/") + if stripped.startswith("accounts/") or "#" in stripped: + return stripped + if stripped.startswith(("routers/", "models/")): + return f"accounts/fireworks/{stripped}" + if stripped.endswith("-fast"): + return f"accounts/fireworks/routers/{stripped}" + return f"accounts/fireworks/models/{stripped}" + + class FireworksAIMixin: """ Common Base Config functions across Fireworks AI Endpoints diff --git a/litellm/llms/fireworks_ai/completion/transformation.py b/litellm/llms/fireworks_ai/completion/transformation.py index c141e097d3a..c460510f39c 100644 --- a/litellm/llms/fireworks_ai/completion/transformation.py +++ b/litellm/llms/fireworks_ai/completion/transformation.py @@ -4,7 +4,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from ...base_llm.completion.transformation import BaseTextCompletionConfig from ...openai.completion.utils import _transform_prompt -from ..common_utils import FireworksAIMixin +from ..common_utils import FireworksAIMixin, resolve_fireworks_resource_name class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig): @@ -50,11 +50,8 @@ class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig ) -> dict: prompt: Final = _transform_prompt(messages=messages) - if not model.startswith("accounts/") and "#" not in model: - model = f"accounts/fireworks/models/{model}" - data: Final = { - "model": model, + "model": resolve_fireworks_resource_name(model), "prompt": prompt, **optional_params, } diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 94945ed4bfb..87908ef60c3 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1282,3 +1282,46 @@ def test_streaming_surfaces_fireworks_response_fields(): assert surfaced["fireworks_raw_outputs"] == [raw_output] assert surfaced["fireworks_perf_metrics"] == {"prompt-tokens": 5} assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3] + + +def test_transform_request_routes_router_slug(): + config = FireworksAIConfig() + + data = config.transform_request( + model="routers/glm-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/routers/glm-latest" + + +def test_transform_request_bare_slug_stays_model(): + config = FireworksAIConfig() + + data = config.transform_request( + model="glm-4p6", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/models/glm-4p6" + + +def test_transform_request_direct_route_passthrough(): + config = FireworksAIConfig() + model = "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c" + + data = config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == model diff --git a/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py b/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py new file mode 100644 index 00000000000..996f1fd975b --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py @@ -0,0 +1,34 @@ +import os +import sys + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.fireworks_ai.completion.transformation import ( + FireworksAITextCompletionConfig, +) + + +def test_transform_text_completion_request_routes_router_slug(): + config = FireworksAITextCompletionConfig() + + data = config.transform_text_completion_request( + model="routers/glm-latest", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/routers/glm-latest" + + +def test_transform_text_completion_request_bare_slug_stays_model(): + config = FireworksAITextCompletionConfig() + + data = config.transform_text_completion_request( + model="glm-4p6", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + headers={}, + ) + + assert data["model"] == "accounts/fireworks/models/glm-4p6" diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py new file mode 100644 index 00000000000..4af395baf41 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py @@ -0,0 +1,45 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name + + +@pytest.mark.parametrize( + "model, expected", + [ + ("routers/glm-latest", "accounts/fireworks/routers/glm-latest"), + ("routers/firerouter", "accounts/fireworks/routers/firerouter"), + ("fireworks_ai/routers/glm-latest", "accounts/fireworks/routers/glm-latest"), + ("models/glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("fireworks_ai/models/glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("fireworks_ai/glm-4p6", "accounts/fireworks/models/glm-4p6"), + ("kimi-k2p6-fast", "accounts/fireworks/routers/kimi-k2p6-fast"), + ( + "accounts/fireworks/routers/glm-latest", + "accounts/fireworks/routers/glm-latest", + ), + ( + "accounts/fireworks/models/glm-4p6", + "accounts/fireworks/models/glm-4p6", + ), + ( + "fireworks_ai/accounts/fireworks/routers/glm-latest", + "accounts/fireworks/routers/glm-latest", + ), + ( + "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c", + "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c", + ), + ( + "glm-4p6#accounts/gitlab/deployments/2fb7764c", + "glm-4p6#accounts/gitlab/deployments/2fb7764c", + ), + ], +) +def test_resolve_fireworks_resource_name(model, expected): + assert resolve_fireworks_resource_name(model) == expected diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..d9038e20df9 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -30,6 +30,6 @@ "limit": 16783 }, "LIT011": { - "limit": 5602 + "limit": 5599 } } From 8d6b8d2ce94b712b941fd117e1279454f880db50 Mon Sep 17 00:00:00 2001 From: LHMQ878 <72402929@cityu-dg.edu.cn> Date: Fri, 7 Aug 2026 10:47:42 +0800 Subject: [PATCH 007/121] fix(proxy): register WebSocket passthrough for OpenAI prefixes create_websocket_passthrough_route existed but /openai and /openai_passthrough only registered HTTP methods, so WS upgrades were rejected at routing. Add catch-all websocket routes mirroring the HTTP passthrough target construction. Fixes #36088 --- .../llm_passthrough_endpoints.py | 47 ++++++++++++++++++- .../test_openai_ws_passthrough_routes.py | 15 ++++++ 2 files changed, 61 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 40c49df26cf..f9409ab366f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -27,7 +27,7 @@ from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth, user_api_key_auth_websocket from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -1934,6 +1934,51 @@ async def openai_proxy_route( ) +@router.websocket("/openai_passthrough/{endpoint:path}") +@router.websocket("/openai/{endpoint:path}") +async def openai_websocket_proxy_route( + websocket: WebSocket, + endpoint: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), +): + """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" + base_target_url = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" + openai_api_key = passthrough_endpoint_router.get_credentials( + custom_llm_provider=litellm.LlmProviders.OPENAI.value, + region_name=None, + ) + if openai_api_key is None: + await websocket.close(code=1011) + raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") + + encoded_endpoint = httpx.URL(endpoint).path + if not encoded_endpoint.startswith("/"): + encoded_endpoint = "/" + encoded_endpoint + base_url = httpx.URL(base_target_url) + updated_url = BaseOpenAIPassThroughHandler._join_url_paths( + base_url=base_url, + path=encoded_endpoint, + custom_llm_provider=litellm.LlmProviders.OPENAI, + ) + # HTTP(S) base -> WS(S) target for the upgrade. + if updated_url.startswith("https://"): + wss_target = "wss://" + updated_url[len("https://") :] + elif updated_url.startswith("http://"): + wss_target = "ws://" + updated_url[len("http://") :] + else: + wss_target = updated_url + + return await websocket_passthrough_request( + websocket=websocket, + target=wss_target, + custom_headers={"Authorization": f"Bearer {openai_api_key}"}, + user_api_key_dict=user_api_key_dict, + forward_headers=True, + endpoint=f"/openai/{endpoint}", + accept_websocket=True, + ) + + class BaseOpenAIPassThroughHandler: @staticmethod async def _base_openai_pass_through_handler( diff --git a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py new file mode 100644 index 00000000000..e0184c6c428 --- /dev/null +++ b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py @@ -0,0 +1,15 @@ +"""OpenAI passthrough must register WebSocket catch-all routes (#36088).""" + +from starlette.routing import WebSocketRoute + +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router + + +def test_openai_websocket_passthrough_routes_registered(): + ws_paths = { + route.path + for route in router.routes + if isinstance(route, WebSocketRoute) + } + assert "/openai/{endpoint:path}" in ws_paths + assert "/openai_passthrough/{endpoint:path}" in ws_paths From 55b52970eac10e36d6563cbff775c0840aacb117 Mon Sep 17 00:00:00 2001 From: LHMQ878 <72402929@cityu-dg.edu.cn> Date: Fri, 7 Aug 2026 11:09:27 +0800 Subject: [PATCH 008/121] fix(proxy): preserve OpenAI WS query params and provider auth Forward realtime model query string, keep OPENAI_API_KEY (forward_headers=False), satisfy ruff strict gates, sync dashboard OpenAPI types, and cover the behavior in tests. --- .../llm_passthrough_endpoints.py | 18 +++-- .../test_openai_ws_passthrough_routes.py | 43 ++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 76 +++++++++++++++++++ 3 files changed, 128 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f9409ab366f..31a8836abf8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -9,7 +9,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. import json import os import re -from typing import Any, Final, cast +from typing import Annotated, Any, Final, cast import httpx from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket @@ -1939,8 +1939,8 @@ async def openai_proxy_route( async def openai_websocket_proxy_route( websocket: WebSocket, endpoint: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), -): + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], +) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" base_target_url = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" openai_api_key = passthrough_endpoint_router.get_credentials( @@ -1949,7 +1949,7 @@ async def openai_websocket_proxy_route( ) if openai_api_key is None: await websocket.close(code=1011) - raise Exception("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") + raise ValueError("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") encoded_endpoint = httpx.URL(endpoint).path if not encoded_endpoint.startswith("/"): @@ -1960,7 +1960,6 @@ async def openai_websocket_proxy_route( path=encoded_endpoint, custom_llm_provider=litellm.LlmProviders.OPENAI, ) - # HTTP(S) base -> WS(S) target for the upgrade. if updated_url.startswith("https://"): wss_target = "wss://" + updated_url[len("https://") :] elif updated_url.startswith("http://"): @@ -1968,12 +1967,17 @@ async def openai_websocket_proxy_route( else: wss_target = updated_url - return await websocket_passthrough_request( + query_string = websocket.url.query + if query_string: + separator = "&" if "?" in wss_target else "?" + wss_target = f"{wss_target}{separator}{query_string}" + + await websocket_passthrough_request( websocket=websocket, target=wss_target, custom_headers={"Authorization": f"Bearer {openai_api_key}"}, user_api_key_dict=user_api_key_dict, - forward_headers=True, + forward_headers=False, endpoint=f"/openai/{endpoint}", accept_websocket=True, ) diff --git a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py index e0184c6c428..9101cd4b780 100644 --- a/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/test_litellm/proxy/test_openai_ws_passthrough_routes.py @@ -1,8 +1,14 @@ -"""OpenAI passthrough must register WebSocket catch-all routes (#36088).""" +"""OpenAI passthrough must register WebSocket catch-all routes (#36088).""" from starlette.routing import WebSocketRoute +from unittest.mock import AsyncMock, MagicMock, patch -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router +import pytest + +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + openai_websocket_proxy_route, + router, +) def test_openai_websocket_passthrough_routes_registered(): @@ -13,3 +19,36 @@ def test_openai_websocket_passthrough_routes_registered(): } assert "/openai/{endpoint:path}" in ws_paths assert "/openai_passthrough/{endpoint:path}" in ws_paths + + +@pytest.mark.asyncio +async def test_openai_websocket_forwards_query_and_keeps_provider_auth(): + websocket = MagicMock() + websocket.url.query = "model=gpt-4o-realtime-preview" + websocket.close = AsyncMock() + user = MagicMock() + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="sk-provider", + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.BaseOpenAIPassThroughHandler._join_url_paths", + return_value="https://api.openai.com/v1/realtime", + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.websocket_passthrough_request", + new_callable=AsyncMock, + ) as mock_ws, + ): + await openai_websocket_proxy_route( + websocket=websocket, + endpoint="v1/realtime", + user_api_key_dict=user, + ) + + kwargs = mock_ws.await_args.kwargs + assert kwargs["target"] == "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview" + assert kwargs["custom_headers"] == {"Authorization": "Bearer sk-provider"} + assert kwargs["forward_headers"] is False diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 752572f9863..d175f94c634 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8483,6 +8483,26 @@ export interface paths { patch?: never; trace?: never; }; + "/openai/": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: openai_websocket_proxy_route + * @description WebSocket connection endpoint + */ + get: operations["websocket_openai_websocket_proxy_route_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/openai/deployments/{model}/chat/completions": { parameters: { query?: never; @@ -8983,6 +9003,26 @@ export interface paths { patch: operations["openai_proxy_route_openai__endpoint__patch"]; trace?: never; }; + "/openai_passthrough/": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: openai_websocket_proxy_route + * @description WebSocket connection endpoint + */ + get: operations["websocket_openai_websocket_proxy_route_get_2"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/openai_passthrough/{endpoint}": { parameters: { query?: never; @@ -46427,6 +46467,24 @@ export interface operations { }; }; }; + websocket_openai_websocket_proxy_route_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; chat_completion_openai_deployments__model__chat_completions_post: { parameters: { query?: never; @@ -47286,6 +47344,24 @@ export interface operations { }; }; }; + websocket_openai_websocket_proxy_route_get_2: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; openai_proxy_route_openai_passthrough__endpoint__get: { parameters: { query?: never; From ae4ee365902478c5141f166f94e739a8860674ff Mon Sep 17 00:00:00 2001 From: LHMQ878 <72402929@cityu-dg.edu.cn> Date: Fri, 7 Aug 2026 11:21:25 +0800 Subject: [PATCH 009/121] fix(proxy): satisfy type-discipline Final/mutable rules on OpenAI WS route --- .../llm_passthrough_endpoints.py | 40 ++++++++++--------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 31a8836abf8..9ff71c77130 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1942,8 +1942,8 @@ async def openai_websocket_proxy_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" - base_target_url = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" - openai_api_key = passthrough_endpoint_router.get_credentials( + base_target_url: Final = os.getenv("OPENAI_API_BASE") or "https://api.openai.com/" + openai_api_key: Final = passthrough_endpoint_router.get_credentials( custom_llm_provider=litellm.LlmProviders.OPENAI.value, region_name=None, ) @@ -1951,31 +1951,33 @@ async def openai_websocket_proxy_route( await websocket.close(code=1011) raise ValueError("Required 'OPENAI_API_KEY' in environment to make pass-through calls to OpenAI.") - encoded_endpoint = httpx.URL(endpoint).path - if not encoded_endpoint.startswith("/"): - encoded_endpoint = "/" + encoded_endpoint - base_url = httpx.URL(base_target_url) - updated_url = BaseOpenAIPassThroughHandler._join_url_paths( + raw_path: Final = httpx.URL(endpoint).path + encoded_endpoint: Final = raw_path if raw_path.startswith("/") else f"/{raw_path}" + base_url: Final = httpx.URL(base_target_url) + updated_url: Final = BaseOpenAIPassThroughHandler._join_url_paths( base_url=base_url, path=encoded_endpoint, custom_llm_provider=litellm.LlmProviders.OPENAI, ) - if updated_url.startswith("https://"): - wss_target = "wss://" + updated_url[len("https://") :] - elif updated_url.startswith("http://"): - wss_target = "ws://" + updated_url[len("http://") :] - else: - wss_target = updated_url - - query_string = websocket.url.query - if query_string: - separator = "&" if "?" in wss_target else "?" - wss_target = f"{wss_target}{separator}{query_string}" + wss_base: Final = ( + "wss://" + updated_url[len("https://") :] + if updated_url.startswith("https://") + else "ws://" + updated_url[len("http://") :] + if updated_url.startswith("http://") + else updated_url + ) + query_string: Final = websocket.url.query + wss_target: Final = ( + f"{wss_base}{'&' if '?' in wss_base else '?'}{query_string}" if query_string else wss_base + ) + custom_headers: Final = { + "Authorization": f"Bearer {openai_api_key}" + } # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers await websocket_passthrough_request( websocket=websocket, target=wss_target, - custom_headers={"Authorization": f"Bearer {openai_api_key}"}, + custom_headers=custom_headers, user_api_key_dict=user_api_key_dict, forward_headers=False, endpoint=f"/openai/{endpoint}", From 7a9c38ed0f1302f9064962ed3f951a41ca6a9bd3 Mon Sep 17 00:00:00 2001 From: LHMQ878 <72402929@cityu-dg.edu.cn> Date: Fri, 7 Aug 2026 11:21:46 +0800 Subject: [PATCH 010/121] fix(proxy): place mutable-ok on OpenAI WS headers dict literal --- .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9ff71c77130..9ef5f22ec4b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1970,9 +1970,9 @@ async def openai_websocket_proxy_route( wss_target: Final = ( f"{wss_base}{'&' if '?' in wss_base else '?'}{query_string}" if query_string else wss_base ) - custom_headers: Final = { + custom_headers: Final = { # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers "Authorization": f"Bearer {openai_api_key}" - } # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers + } await websocket_passthrough_request( websocket=websocket, From 02d4f8d6a86eae6d02393d29f99e72f0712eb234 Mon Sep 17 00:00:00 2001 From: LHMQ878 <72402929@cityu-dg.edu.cn> Date: Fri, 7 Aug 2026 11:45:21 +0800 Subject: [PATCH 011/121] style(proxy): ruff-format OpenAI websocket passthrough route --- .../proxy/pass_through_endpoints/llm_passthrough_endpoints.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9ef5f22ec4b..216d966b800 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1967,9 +1967,7 @@ async def openai_websocket_proxy_route( else updated_url ) query_string: Final = websocket.url.query - wss_target: Final = ( - f"{wss_base}{'&' if '?' in wss_base else '?'}{query_string}" if query_string else wss_base - ) + wss_target: Final = f"{wss_base}{'&' if '?' in wss_base else '?'}{query_string}" if query_string else wss_base custom_headers: Final = { # mutable-ok: websocket_passthrough_request requires a plain dict of upstream headers "Authorization": f"Bearer {openai_api_key}" } From 1b2430b8b6c468a017bea31d8d7d7b3164976d1f Mon Sep 17 00:00:00 2001 From: devin-ai-integration Date: Mon, 10 Aug 2026 10:25:37 +0000 Subject: [PATCH 012/121] fix(bedrock): report uploaded size in the FileObject returned by managed batch uploads --- litellm/llms/bedrock/files/transformation.py | 25 ++++++--- .../test_bedrock_files_transformation.py | 54 +++++++++++++++++++ 2 files changed, 72 insertions(+), 7 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index bd3570d50a3..d5e7957a1a9 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -62,6 +62,10 @@ from ..common_utils import BedrockError, resolve_s3_encryption_key_id # Same pattern as the `upload_url` handoff in `transform_create_file_request`. S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers" +# litellm_params key carrying the size of the body uploaded to S3, handed from +# `transform_create_file_request` to `transform_create_file_response`. +UPLOAD_CONTENT_LENGTH_PARAM: Final = "_s3_upload_content_length" + def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]: return MappingProxyType(dict(items)) @@ -154,6 +158,18 @@ def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str: return bucket_name +def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int: + """ + S3 answers PutObject with an empty body, so the stored object size comes from the + signed request recorded by `transform_create_file_request`, not the response headers. + """ + uploaded_size: Final = litellm_params.get(UPLOAD_CONTENT_LENGTH_PARAM) + if isinstance(uploaded_size, int): + return uploaded_size + response_content_length: Final = raw_response.headers.get("Content-Length", "0") + return int(response_content_length) if response_content_length.isdigit() else 0 + + class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Config for Bedrock Files - handles S3 uploads for Bedrock batch processing @@ -861,6 +877,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) litellm_params["upload_url"] = api_base + litellm_params[UPLOAD_CONTENT_LENGTH_PARAM] = len(file_content.encode("utf-8")) # Return a dict that tells the HTTP handler exactly what to do return { @@ -1018,12 +1035,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Transform S3 File upload response into OpenAI-style FileObject """ - # For S3 uploads, we typically get an ETag and other metadata - response_headers: Final = raw_response.headers - # Extract S3 object information from the response - # S3 PUT object returns ETag and other metadata in headers - content_length: Final = response_headers.get("Content-Length", "0") - # Use the actual upload URL that was used for the S3 upload upload_url: Final = litellm_params.get("upload_url") file_id: str = "" @@ -1038,7 +1049,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): filename=filename, created_at=int(time.time()), # Current timestamp status="uploaded", - bytes=int(content_length) if content_length.isdigit() else 0, + bytes=_uploaded_object_size(litellm_params=litellm_params, raw_response=raw_response), object="file", ) 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 270add48e0e..e0b235d68fd 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 @@ -586,6 +586,60 @@ class TestBedrockFilesTransformation: assert "x-amz-server-side-encryption" not in headers assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + def test_create_file_response_reports_uploaded_object_size(self): + """ + S3 answers PutObject with an empty body, so the returned FileObject must report the + size of the body that was uploaded instead of the response's Content-Length (always 0). + """ + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + litellm_params: dict = {"s3_bucket_name": "litellm-batch-bucket"} + 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=litellm_params, + ) + assert isinstance(request, dict) + uploaded_size = len(request["data"].encode("utf-8")) + assert uploaded_size > 0 + + file_object = config.transform_create_file_response( + model=None, + raw_response=httpx.Response( + status_code=200, + headers={"Content-Length": "0", "ETag": '"abc123"'}, + content=b"", + ), + logging_obj=MagicMock(), + litellm_params=litellm_params, + ) + + assert file_object.bytes == uploaded_size + def test_openai_passthrough_still_works(self): """ Regression test: ensure OpenAI-compatible models (e.g. gpt-oss) From c019ce53e3daa6cde544aef6b1d468d719311a25 Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 11:15:45 -0400 Subject: [PATCH 013/121] feat(ui): add user ID request log filter Co-Authored-By: Codex --- .../view_logs/RequestLogsFilters.test.tsx | 82 ++++++++++++++++++- .../view_logs/RequestLogsFilters.tsx | 51 +++++++++++- .../components/view_logs/RequestLogsPanel.tsx | 3 +- .../components/view_logs/RequestLogsTable.tsx | 12 ++- .../components/view_logs/log_filter_logic.tsx | 1 + 5 files changed, 143 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index 82b179b2654..0d50effabb4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -14,6 +14,10 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useInfiniteModelInfo: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ + useInfiniteUsers: vi.fn(), +})); + vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers", () => ({ useInfiniteSpendLogEndUsers: vi.fn(), })); @@ -21,6 +25,7 @@ vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers", () => ({ import { useInfiniteSpendLogEndUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers"; import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; const emptyInfiniteQuery = { data: { pages: [], pageParams: [] }, @@ -32,10 +37,16 @@ const emptyInfiniteQuery = { const LOGS_WINDOW = { start_date: "2026-07-23 00:00:00", end_date: "2026-07-24 00:00:00" }; -function renderFilters(filters: Record = {}) { +function renderFilters(filters: Record = {}, showUserIdFilter = true) { const set = vi.fn(); renderWithProviders( - filters[id]} set={set} teams={[]} logsWindow={LOGS_WINDOW} />, + filters[id]} + set={set} + teams={[]} + logsWindow={LOGS_WINDOW} + showUserIdFilter={showUserIdFilter} + />, ); return { set }; } @@ -50,6 +61,9 @@ describe("RequestLogsFilters", () => { vi.mocked(useInfiniteModelInfo).mockReturnValue( emptyInfiniteQuery as unknown as ReturnType, ); + vi.mocked(useInfiniteUsers).mockReturnValue( + emptyInfiniteQuery as unknown as ReturnType, + ); vi.mocked(useInfiniteSpendLogEndUsers).mockReturnValue( emptyInfiniteQuery as unknown as ReturnType, ); @@ -62,6 +76,7 @@ describe("RequestLogsFilters", () => { "Team ID", "Status", "Key Alias", + "User ID", "End User", "Error Code", "Error Message", @@ -74,6 +89,59 @@ describe("RequestLogsFilters", () => { } }); + it("places User ID between Key Alias and End User", async () => { + renderFilters(); + + const labels = ["Key Alias", "User ID", "End User"].map((label) => screen.getByText(label)); + expect(labels[0].compareDocumentPosition(labels[1]) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + expect(labels[1].compareDocumentPosition(labels[2]) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); + }); + + it("selects a user by display name while storing the user ID filter", async () => { + vi.mocked(useInfiniteUsers).mockReturnValue({ + ...emptyInfiniteQuery, + data: { + pages: [ + { + users: [{ user_id: "user-1", user_alias: "Alice", user_email: "alice@example.com" }], + page: 1, + page_size: 50, + total: 1, + total_pages: 1, + }, + ], + pageParams: [1], + }, + } as unknown as ReturnType); + const user = userEvent.setup(); + const { set } = renderFilters(); + + await user.click(await screen.findByPlaceholderText("Search an internal user")); + expect(await screen.findByText("Alice")).toBeInTheDocument(); + expect(screen.getByText("alice@example.com | User ID: user-1")).toBeInTheDocument(); + await user.click(screen.getByText("Alice")); + + expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.USER_ID, "user-1"); + }); + + it("pushes the User ID picker query to the paginated user lookup", async () => { + const user = userEvent.setup(); + renderFilters(); + + const input = await screen.findByPlaceholderText("Search an internal user"); + await user.click(input); + await user.type(input, "alice@example.com"); + + await waitFor(() => expect(useInfiniteUsers).toHaveBeenCalledWith(50, "alice@example.com")); + }); + + it("does not show or query the User ID filter for non-admin request logs", () => { + renderFilters({}, false); + + expect(screen.queryByText("User ID")).not.toBeInTheDocument(); + expect(useInfiniteUsers).not.toHaveBeenCalled(); + }); + it("scopes the Key Alias lookup to the selected team", async () => { renderFilters({ [LOG_FILTER_IDS.TEAM_ID]: "team-42" }); @@ -164,7 +232,15 @@ describe("RequestLogsFilters", () => { it("scopes the End User lookup to the window the logs table is showing", async () => { const otherWindow = { start_date: "2026-01-01 00:00:00", end_date: "2026-01-02 00:00:00" }; - renderWithProviders( undefined} set={vi.fn()} teams={[]} logsWindow={otherWindow} />); + renderWithProviders( + undefined} + set={vi.fn()} + teams={[]} + logsWindow={otherWindow} + showUserIdFilter + />, + ); await waitFor(() => expect(useInfiniteSpendLogEndUsers).toHaveBeenCalledWith(otherWindow, 50, undefined)); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx index 2005b868cd6..47e0bad6f62 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx @@ -5,6 +5,7 @@ import { useMemo, useState } from "react"; import { useInfiniteSpendLogEndUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers"; import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import { DataTableFilterField } from "@/components/shared/DataTable"; import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; @@ -144,6 +145,46 @@ function ModelFilterField({ value, onChange }: { value: string; onChange: (value ); } +function UserIdFilterField({ value, onChange }: { value: string; onChange: (value: string | undefined) => void }) { + const [search, setSearch] = useState(""); + const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteUsers( + PAGE_SIZE, + emptyToUndefined(search), + ); + + const options = useMemo(() => { + const seen = new Set(); + return (data?.pages ?? []).flatMap((page) => + page.users.flatMap((user) => { + if (!user.user_id || seen.has(user.user_id)) return []; + seen.add(user.user_id); + const label = user.user_alias || user.user_email || user.user_id; + const email = user.user_email && user.user_email !== label ? user.user_email : ""; + const sublabel = + user.user_id === label ? email : [email, `User ID: ${user.user_id}`].filter(Boolean).join(" | "); + return [{ label, value: user.user_id, sublabel }]; + }), + ); + }, [data]); + + return ( + + onChange(emptyToUndefined(next))} + onSearchChange={setSearch} + onLoadMore={() => void fetchNextPage()} + hasNextPage={hasNextPage} + isLoading={isLoading} + isFetchingNextPage={isFetchingNextPage} + placeholder="Search an internal user" + emptyText="No users found" + /> + + ); +} + function EndUserFilterField({ value, onChange, @@ -243,9 +284,10 @@ interface RequestLogsFiltersProps { set: (columnId: string, value: unknown) => void; teams: Team[]; logsWindow: LogsWindow; + showUserIdFilter: boolean; } -export function RequestLogsFilters({ get, set, teams, logsWindow }: RequestLogsFiltersProps) { +export function RequestLogsFilters({ get, set, teams, logsWindow, showUserIdFilter }: RequestLogsFiltersProps) { const valueOf = (id: string): string => asString(get(id)); const setter = (id: string) => (next: string | undefined) => set(id, next); @@ -279,6 +321,13 @@ export function RequestLogsFilters({ get, set, teams, logsWindow }: RequestLogsF teamId={valueOf(LOG_FILTER_IDS.TEAM_ID)} /> + {showUserIdFilter && ( + + )} + void; teams: Team[]; logsWindow: LogsWindow; + showUserIdFilter: boolean; toolbarChildren?: ReactNode; } @@ -69,6 +70,7 @@ export function RequestLogsTable({ onSessionClick, teams, logsWindow, + showUserIdFilter, toolbarChildren, }: RequestLogsTableProps) { const [filtersOpen, setFiltersOpen] = useState(false); @@ -122,7 +124,15 @@ export function RequestLogsTable({ title="Filters" description="Narrow down request logs" > - {({ get, set }) => } + {({ get, set }) => ( + + )} )} diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index e1089c6a16c..78ecb52c184 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -36,6 +36,7 @@ export const LOG_FILTER_LABELS: Record = { [LOG_FILTER_IDS.TEAM_ID]: "Team ID", [LOG_FILTER_IDS.STATUS]: "Status", [LOG_FILTER_IDS.KEY_ALIAS]: "Key Alias", + [LOG_FILTER_IDS.USER_ID]: "User ID", [LOG_FILTER_IDS.END_USER]: "End User", [LOG_FILTER_IDS.ERROR_CODE]: "Error Code", [LOG_FILTER_IDS.ERROR_MESSAGE]: "Error Message", From 151bdbb2a9d6cfb3ccf758a98b04cfa02779a959 Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 11:21:48 -0400 Subject: [PATCH 014/121] test(ui): cover user filter pagination Co-Authored-By: Codex --- .../view_logs/RequestLogsFilters.test.tsx | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index 0d50effabb4..fd53949c87a 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -135,6 +135,38 @@ describe("RequestLogsFilters", () => { await waitFor(() => expect(useInfiniteUsers).toHaveBeenCalledWith(50, "alice@example.com")); }); + it("loads the next page when the User ID list is scrolled near the end", async () => { + const fetchNextPage = vi.fn(); + vi.mocked(useInfiniteUsers).mockReturnValue({ + ...emptyInfiniteQuery, + fetchNextPage, + hasNextPage: true, + data: { + pages: [ + { + users: [{ user_id: "user-1", user_alias: "Alice", user_email: "alice@example.com" }], + page: 1, + page_size: 50, + total: 51, + total_pages: 2, + }, + ], + pageParams: [1], + }, + } as unknown as ReturnType); + const user = userEvent.setup(); + renderFilters(); + + await user.click(await screen.findByPlaceholderText("Search an internal user")); + const list = await screen.findByTestId("paginated-search-select-list"); + Object.defineProperty(list, "scrollTop", { value: 90, configurable: true }); + Object.defineProperty(list, "clientHeight", { value: 10, configurable: true }); + Object.defineProperty(list, "scrollHeight", { value: 100, configurable: true }); + list.dispatchEvent(new Event("scroll", { bubbles: true })); + + await waitFor(() => expect(fetchNextPage).toHaveBeenCalled()); + }); + it("does not show or query the User ID filter for non-admin request logs", () => { renderFilters({}, false); From 297fe272ecce15832565a6cc27c13763160944bb Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 11:50:43 -0400 Subject: [PATCH 015/121] feat: scope request log user filter Add a bounded spend-log user facet for the Request Logs picker and intersect explicit user filters with the caller's own and permitted-team scope. Co-Authored-By: Codex --- litellm/proxy/_types.py | 6 +- .../management_v1/spend_logs.py | 222 +++++++++++------- .../spend_management_endpoints.py | 17 +- .../management_v1/test_spend_logs.py | 83 +++++-- .../test_spend_management_endpoints.py | 78 +++++- .../hooks/spendLogs/useSpendLogUsers.test.ts | 40 ++++ .../hooks/spendLogs/useSpendLogUsers.ts | 21 ++ .../view_logs/RequestLogsFilters.test.tsx | 71 ++---- .../view_logs/RequestLogsFilters.tsx | 41 ++-- .../components/view_logs/RequestLogsPanel.tsx | 5 - .../components/view_logs/RequestLogsTable.tsx | 12 +- .../view_logs/log_filter_logic.test.tsx | 12 +- .../components/view_logs/log_filter_logic.tsx | 5 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 60 +++++ 14 files changed, 470 insertions(+), 203 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb330d00756..d628d956e73 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -679,6 +679,7 @@ class LiteLLMRoutes(enum.Enum): # permitted teams exactly like /spend/logs/ui — it belongs to the same # access tier, not to customer management. "/management/v1/spend_logs/end_users", + "/management/v1/spend_logs/users", "/cost/estimate", ] @@ -871,12 +872,13 @@ class LiteLLMRoutes(enum.Enum): # PROXY_ADMIN_VIEW_ONLY — the route gate must match). "/customer/list", "/customer/info", - # UI Logs page detail drawer (single + session) and the end-user filter - # facet. The list endpoint `/spend/logs/ui` is covered via + # UI Logs page detail drawer (single + session) and the filter facets. + # The list endpoint `/spend/logs/ui` is covered via # spend_tracking_routes below. "/spend/logs/ui/{logId}", "/spend/logs/session/ui", "/management/v1/spend_logs/end_users", + "/management/v1/spend_logs/users", # Settings / observability read endpoints exposed in admin-only # sidebar groups (Logging & Alerts, Admin Settings, Budgets, # Invitations). diff --git a/litellm/proxy/management_endpoints/management_v1/spend_logs.py b/litellm/proxy/management_endpoints/management_v1/spend_logs.py index 96e60fcfdfc..5fee8eaede3 100644 --- a/litellm/proxy/management_endpoints/management_v1/spend_logs.py +++ b/litellm/proxy/management_endpoints/management_v1/spend_logs.py @@ -1,7 +1,7 @@ """`/management/v1/spend_logs` facets.""" from datetime import datetime, timezone -from typing import Annotated, Any, Final +from typing import Annotated, Any, Final, Literal from fastapi import APIRouter, Depends, Query, Request @@ -35,7 +35,7 @@ def _as_utc(value: datetime) -> datetime: return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) -async def _end_user_scope_clause( +async def _spend_log_scope_clause( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, next_param_index: int, @@ -43,8 +43,8 @@ async def _end_user_scope_clause( """SQL predicate restricting the facet to spend logs this caller may read. Returns ``(None, ())`` for a proxy admin. Mirrors the scoping ``/spend/logs/ui`` - applies, so the dropdown can never offer an end user whose rows the caller - could not open. + applies, so a dropdown can never offer a value from a row the caller could + not open. """ from litellm.proxy.spend_tracking.spend_management_endpoints import ( _get_permitted_team_ids_for_spend_logs, @@ -77,6 +77,98 @@ async def _end_user_scope_clause( return f"({' OR '.join(clauses)})", params +async def _list_spend_log_facet( + request: Request, + user_api_key_dict: UserAPIKeyAuth, + start_time: datetime, + end_time: datetime, + q: str | None, + page: int, + page_size: int, + column: Literal["end_user", "user"], +) -> FacetListResponse: + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}database-not-connected", + title="Database not connected", + status=503, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + ) + + column_sql: Final = "end_user" if column == "end_user" else '"user"' + window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time)) + search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else () + search_clause: Final = (f"{column_sql} ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () + + scope_clause, scope_params = await _spend_log_scope_clause( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + next_param_index=len(window_params) + len(search_params) + 1, + ) + + where_parts: Final = ( + ( + "\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')", + "\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')", + f"{column_sql} IS NOT NULL", + f"{column_sql} != ''", + ) + + search_clause + + ((scope_clause,) if scope_clause is not None else ()) + ) + + # The inner LIMIT walks the startTime index newest first and bounds the + # rows DISTINCT can inspect. request_id makes the cut-off deterministic, + # and page_size + 1 reveals has_more without a COUNT(*). + params: Final = ( + window_params + + search_params + + scope_params + + (SPEND_LOGS_FACET_SCAN_CAP, page_size + 1, (page - 1) * page_size) + ) + scan_idx: Final = len(params) - 2 + facet_sql: Final = ( + f"SELECT DISTINCT {column_sql} FROM (" + f" SELECT {column_sql}" + f' FROM "LiteLLM_SpendLogs"' + f" WHERE {' AND '.join(where_parts)}" + f' ORDER BY "startTime" DESC, request_id DESC' + f" LIMIT ${scan_idx}" + f") recent" + f" ORDER BY {column_sql} ASC" + f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" + ) + rows: Final = await prisma_client.db.query_raw(facet_sql, *params) + values: Final[list[str]] = [row[column] for row in rows if row.get(column)] + has_more: Final = len(values) > page_size + + return FacetListResponse( + data=values[:page_size], + meta=PageMeta(page=page, page_size=page_size, has_more=has_more), + links=build_page_links(request=request, page=page, has_more=has_more), + ) + except ManagementProblem: + raise + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.management_endpoints.management_v1.spend_logs._list_spend_log_facet(): Exception occured - %s", + e, + ) + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}internal-server-error", + title="Internal server error", + status=500, + detail=f"Failed to list spend log {column.replace('_', ' ')}s.", + ) + ) + + @router.get( "/spend_logs/end_users", tags=["Budget & Spend Tracking"], @@ -116,85 +208,47 @@ async def list_spend_log_end_users( --header 'Authorization: Bearer sk-1234' ``` """ - try: - from litellm.proxy.proxy_server import prisma_client + return await _list_spend_log_facet( + request=request, + user_api_key_dict=user_api_key_dict, + start_time=start_time, + end_time=end_time, + q=q, + page=page, + page_size=page_size, + column="end_user", + ) - if prisma_client is None: - raise ManagementProblem( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}database-not-connected", - title="Database not connected", - status=503, - detail=CommonProxyErrors.db_not_connected_error.value, - ) - ) - window_params: Final[tuple[Any, ...]] = (_as_utc(start_time), _as_utc(end_time)) - search_params: Final[tuple[Any, ...]] = (f"%{escape_like(q)}%",) if q else () - search_clause: Final = (f"end_user ILIKE ${len(window_params) + 1} ESCAPE '\\'",) if q else () - - scope_clause, scope_params = await _end_user_scope_clause( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - next_param_index=len(window_params) + len(search_params) + 1, - ) - - where_parts: Final = ( - ( - "\"startTime\" >= ($1::timestamptz AT TIME ZONE 'UTC')", - "\"startTime\" <= ($2::timestamptz AT TIME ZONE 'UTC')", - "end_user IS NOT NULL", - "end_user != ''", - ) - + search_clause - + ((scope_clause,) if scope_clause is not None else ()) - ) - - # The inner LIMIT is the safety bound: it walks the startTime index newest - # first and stops, so DISTINCT never runs over an unbounded row set. - # request_id breaks startTime ties so the cut-off row is deterministic and - # successive OFFSET pages agree on the set they are paging through. - # page_size + 1: one row beyond the page reveals has_more without a COUNT(*). - params: Final = ( - window_params - + search_params - + scope_params - + (SPEND_LOGS_FACET_SCAN_CAP, page_size + 1, (page - 1) * page_size) - ) - scan_idx: Final = len(params) - 2 - facet_sql: Final = ( - f"SELECT DISTINCT end_user FROM (" - f" SELECT end_user" - f' FROM "LiteLLM_SpendLogs"' - f" WHERE {' AND '.join(where_parts)}" - f' ORDER BY "startTime" DESC, request_id DESC' - f" LIMIT ${scan_idx}" - f") recent" - f" ORDER BY end_user ASC" - f" LIMIT ${scan_idx + 1} OFFSET ${scan_idx + 2}" - ) - rows: Final = await prisma_client.db.query_raw(facet_sql, *params) - end_users: Final[list[str]] = [row["end_user"] for row in rows if row.get("end_user")] - has_more: Final = len(end_users) > page_size - - return FacetListResponse( - data=end_users[:page_size], - meta=PageMeta(page=page, page_size=page_size, has_more=has_more), - links=build_page_links(request=request, page=page, has_more=has_more), - ) - - except ManagementProblem: - raise - except Exception as e: - verbose_proxy_logger.exception( - "litellm.proxy.management_endpoints.management_v1.spend_logs.list_spend_log_end_users(): Exception occured - %s", - e, - ) - raise ManagementProblem( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}internal-server-error", - title="Internal server error", - status=500, - detail="Failed to list spend log end users.", - ) - ) +@router.get( + "/spend_logs/users", + tags=["Budget & Spend Tracking"], + dependencies=[Depends(user_api_key_auth), Depends(reject_unknown_query_params)], + response_model=FacetListResponse, +) +async def list_spend_log_users( + request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_time: Annotated[ + datetime, + Query(alias="filter[startTime][gte]", description="Window start (UTC when no offset is given)"), + ], + end_time: Annotated[ + datetime, + Query(alias="filter[startTime][lte]", description="Window end (UTC when no offset is given)"), + ], + q: Annotated[str | None, Query(description="Case-insensitive partial match on the internal user id")] = None, + page: Annotated[int, Query(ge=1, description="Page number")] = 1, + page_size: Annotated[int, Query(ge=1, le=100, description="Page size")] = 50, +) -> FacetListResponse: + """The distinct internal users appearing in spend logs the caller can read.""" + return await _list_spend_log_facet( + request=request, + user_api_key_dict=user_api_key_dict, + start_time=start_time, + end_time=end_time, + q=q, + page=page, + page_size=page_size, + column="user", + ) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8fb5570965b..99d870f5ad4 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2427,6 +2427,7 @@ async def ui_view_spend_logs( request_id=request_id, ) permitted_team_ids: list[str] | None = None + scope_to_caller_user = False if not is_request_id_lookup and not is_admin_view: if team_id is not None: can_view_team: Final = await _can_team_member_view_log( @@ -2440,7 +2441,6 @@ async def ui_view_spend_logs( detail={"error": f"Not authorized to view team spend for team_id={team_id}"}, ) where_conditions["team_id"] = team_id - where_conditions.pop("user", None) else: if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict): try: @@ -2451,13 +2451,20 @@ async def ui_view_spend_logs( except Exception: permitted_team_ids = [] if permitted_team_ids: - where_conditions.pop("user", None) + if user_id is None: + where_conditions.pop("user", None) where_conditions["OR"] = [ {"user": user_api_key_dict.user_id}, {"team_id": {"in": permitted_team_ids}}, ] else: - where_conditions["user"] = user_api_key_dict.user_id + if user_id is None: + where_conditions["user"] = user_api_key_dict.user_id + else: + where_conditions["AND"] = where_conditions.get("AND", []) + [ + {"user": user_api_key_dict.user_id} + ] + scope_to_caller_user = True where_conditions.pop("team_id", None) # Calculate skip value for pagination skip: Final = (page - 1) * page_size @@ -2508,6 +2515,10 @@ async def ui_view_spend_logs( sql_params.append(permitted_team_ids) p += 2 sql_conditions.append(or_clause) + elif scope_to_caller_user: + sql_conditions.append(f'"user" = ${p}') + sql_params.append(user_api_key_dict.user_id) + p += 1 if session_id is not None and isinstance(session_id, str): like_escaped_session_id: Final = session_id.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py index 79f13a6f703..35fcd3b6cd7 100644 --- a/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_spend_logs.py @@ -1,5 +1,4 @@ from datetime import datetime, timezone -from typing import List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -45,6 +44,7 @@ app.include_router(router) client = TestClient(app) END_USERS_PATH = f"{MANAGEMENT_V1_PREFIX}/spend_logs/end_users" +USERS_PATH = f"{MANAGEMENT_V1_PREFIX}/spend_logs/users" WINDOW = "filter[startTime][gte]=2026-07-23T00:00:00Z&filter[startTime][lte]=2026-07-24T00:00:00Z" @@ -65,7 +65,7 @@ def as_proxy_admin(): app.dependency_overrides.clear() -def _mock_rows(mock_prisma_client, end_users: List[str]) -> AsyncMock: +def _mock_rows(mock_prisma_client, end_users: list[str]) -> AsyncMock: query_raw = AsyncMock(return_value=[{"end_user": eu} for eu in end_users]) mock_prisma_client.db.query_raw = query_raw return query_raw @@ -82,6 +82,11 @@ def _get(query: str = WINDOW): return client.get(f"{END_USERS_PATH}{suffix}", headers={"Authorization": "Bearer k"}) +def _get_users(query: str = WINDOW): + suffix = f"?{query}" if query else "" + return client.get(f"{USERS_PATH}{suffix}", headers={"Authorization": "Bearer k"}) + + def test_returns_the_control_plane_envelope(mock_prisma_client, as_proxy_admin): """`{data, meta, links}` is the contract; a bare list or a legacy `aliases` key is not.""" _mock_rows(mock_prisma_client, ["a", "b"]) @@ -213,7 +218,7 @@ def test_requires_a_time_window(mock_prisma_client, as_proxy_admin, query): def test_rejects_a_malformed_window_as_a_problem_document(mock_prisma_client, as_proxy_admin): _mock_rows(mock_prisma_client, []) - response = _get(f"filter[startTime][gte]=yesterday&filter[startTime][lte]=2026-07-24T00:00:00Z") + response = _get("filter[startTime][gte]=yesterday&filter[startTime][lte]=2026-07-24T00:00:00Z") assert response.status_code == 400 assert response.headers["content-type"].startswith("application/problem+json") @@ -400,6 +405,49 @@ def test_q_placeholder_precedes_the_scan_limit_and_offset(mock_prisma_client, as assert query_raw.call_args.args[5:] == (11, 0) +def test_user_facet_reads_internal_users_from_spend_logs(mock_prisma_client, as_proxy_admin): + query_raw = AsyncMock(return_value=[{"user": "alice@example.com"}, {"user": "user-42"}]) + mock_prisma_client.db.query_raw = query_raw + + response = _get_users() + + assert response.status_code == 200 + assert response.json()["data"] == ["alice@example.com", "user-42"] + sql = query_raw.call_args.args[0] + assert 'SELECT DISTINCT "user"' in sql + assert '"user" IS NOT NULL' in sql + assert "end_user IS NOT NULL" not in sql + + +def test_user_facet_uses_the_same_team_scope_as_request_logs(mock_prisma_client): + query_raw = AsyncMock(return_value=[{"user": "member@example.com"}]) + mock_prisma_client.db.query_raw = query_raw + original = _as_role(LitellmUserRoles.INTERNAL_USER, user_id="team-admin-1") + try: + with patch( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + new=AsyncMock(return_value=["team-a"]), + ): + response = _get_users() + finally: + app.dependency_overrides = original + + assert response.status_code == 200 + assert '("user" = $3 OR team_id = ANY($4::text[]))' in query_raw.call_args.args[0] + assert query_raw.call_args.args[3] == "team-admin-1" + assert query_raw.call_args.args[4] == ["team-a"] + + +def test_user_facet_searches_the_internal_user_value(mock_prisma_client, as_proxy_admin): + query_raw = AsyncMock(return_value=[]) + mock_prisma_client.db.query_raw = query_raw + + _get_users(f"{WINDOW}&q=alice%40example.com") + + assert '"user" ILIKE $3 ESCAPE' in query_raw.call_args.args[0] + assert query_raw.call_args.args[3] == "%alice@example.com%" + + @pytest.mark.parametrize( "role", [ @@ -416,18 +464,19 @@ def test_is_reachable_by_every_role_that_can_open_the_logs_page(role): """ from litellm.proxy.auth.route_checks import RouteChecks - for allowed in ( - LiteLLMRoutes.internal_user_routes.value, - LiteLLMRoutes.internal_user_view_only_routes.value, - ): - assert ("/spend/logs/ui" in allowed) == (END_USERS_PATH in allowed) + for facet_path in (END_USERS_PATH, USERS_PATH): + for allowed in ( + LiteLLMRoutes.internal_user_routes.value, + LiteLLMRoutes.internal_user_view_only_routes.value, + ): + assert ("/spend/logs/ui" in allowed) == (facet_path in allowed) - if role in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY): - allowed_routes = ( - LiteLLMRoutes.internal_user_routes.value - if role == LitellmUserRoles.INTERNAL_USER - else LiteLLMRoutes.internal_user_view_only_routes.value - ) - assert RouteChecks.check_route_access(route=END_USERS_PATH, allowed_routes=allowed_routes) - else: - assert END_USERS_PATH in LiteLLMRoutes.admin_viewer_routes.value + if role in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY): + allowed_routes = ( + LiteLLMRoutes.internal_user_routes.value + if role == LitellmUserRoles.INTERNAL_USER + else LiteLLMRoutes.internal_user_view_only_routes.value + ) + assert RouteChecks.check_route_access(route=facet_path, allowed_routes=allowed_routes) + else: + assert facet_path in LiteLLMRoutes.admin_viewer_routes.value diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 057193a69db..81512cd8e66 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1321,7 +1321,7 @@ async def test_ui_view_spend_logs_internal_user_scoped_without_user_id( @pytest.mark.asyncio -async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeypatch): +async def test_ui_view_spend_logs_team_admin_can_filter_team_spend_by_user(client, monkeypatch): """ Team admins should be able to view team-wide spend when team_id is provided. """ @@ -1346,11 +1346,23 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp "startTime": datetime.datetime.now(timezone.utc).isoformat(), "model": "gpt-4", }, + { + "id": "log3", + "request_id": "req3", + "api_key": "sk-test-key", + "user": "member3", + "team_id": "team_admin_team", + "spend": 0.15, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + }, ] def filter_by_team(where): - if "team_id" in where and where["team_id"] == "team_admin_team": + if where.get("team_id") == "team_admin_team" and where.get("user") == "member1": return [mock_spend_logs[0]] + if where.get("team_id") == "team_admin_team": + return [mock_spend_logs[0], mock_spend_logs[2]] return mock_spend_logs class TeamTable: @@ -1383,6 +1395,7 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp "/spend/logs/ui", params={ "team_id": "team_admin_team", + "user_id": "member1", "start_date": start_date, "end_date": end_date, }, @@ -1398,6 +1411,66 @@ async def test_ui_view_spend_logs_team_admin_can_view_team_spend(client, monkeyp app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_user_filter_intersects_permitted_team_scope(client, monkeypatch): + member_log = { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "member@example.com", + "team_id": "team-9", + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + other_team_log = { + **member_log, + "id": "log2", + "request_id": "req2", + "team_id": "team-outside-scope", + } + seen_where = [] + + def filter_by_user_and_scope(where): + seen_where.append(where) + if where.get("user") == "member@example.com" and {"multi_team": True} in where.get("OR", []): + return [member_log] + return [member_log, other_team_log] + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([member_log, other_team_log], filter_by_user_and_scope), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=["team-9"]), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "user_id": "member@example.com", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [row["request_id"] for row in response.json()["data"]] == ["req1"] + assert any( + where.get("user") == "member@example.com" and {"multi_team": True} in where.get("OR", []) + for where in seen_where + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_pagination(client, monkeypatch): mock_spend_logs = [ @@ -1578,6 +1651,7 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): assert data["total_pages"] == 2 assert len(data["data"]) == 1 assert data["data"][0]["request_id"] == "req1" + assert data["data"][0]["user"] == "member1" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts new file mode 100644 index 00000000000..5a79bf74ce3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.test.ts @@ -0,0 +1,40 @@ +import { renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const useInfiniteQuery = vi.fn(); +vi.mock("@/lib/http/api", () => ({ $api: { useInfiniteQuery: (...args: unknown[]) => useInfiniteQuery(...args) } })); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +import { useInfiniteSpendLogUsers } from "./useSpendLogUsers"; + +const WINDOW = { start_date: "2026-07-23 00:00:00", end_date: "2026-07-24 00:00:00" }; + +describe("useInfiniteSpendLogUsers", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + }); + + it("calls the scoped spend-log user facet with the visible window", () => { + renderHook(() => useInfiniteSpendLogUsers(WINDOW, 25, "alice")); + + const expectedQuery = { + "filter[startTime][gte]": "2026-07-23 00:00:00", + "filter[startTime][lte]": "2026-07-24 00:00:00", + page_size: 25, + q: "alice", + }; + expect(useInfiniteQuery.mock.calls[0][1]).toBe("/management/v1/spend_logs/users"); + expect(useInfiniteQuery.mock.calls[0][2].params.query).toEqual(expectedQuery); + }); + + it("omits q when the search box is empty", () => { + renderHook(() => useInfiniteSpendLogUsers(WINDOW, 50, "")); + + expect(useInfiniteQuery.mock.calls[0][2].params.query).not.toHaveProperty("q"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts new file mode 100644 index 00000000000..3a82c9e9d91 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/spendLogs/useSpendLogUsers.ts @@ -0,0 +1,21 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { $api } from "@/lib/http/api"; + +import { nextPageFromLinks, type SpendLogsWindow } from "./useSpendLogEndUsers"; + +export const useInfiniteSpendLogUsers = (window: SpendLogsWindow, pageSize: number = 50, q?: string) => { + const { accessToken } = useAuthorized(); + const query = { + "filter[startTime][gte]": window.start_date, + "filter[startTime][lte]": window.end_date, + page_size: pageSize, + ...(q !== undefined && q !== "" ? { q } : {}), + }; + const options = { + pageParamName: "page", + initialPageParam: 1, + getNextPageParam: nextPageFromLinks, + enabled: Boolean(accessToken), + }; + return $api.useInfiniteQuery("get", "/management/v1/spend_logs/users", { params: { query } }, options); +}; diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index fd53949c87a..1c94e6418a0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -14,8 +14,8 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useInfiniteModelInfo: vi.fn(), })); -vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({ - useInfiniteUsers: vi.fn(), +vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogUsers", () => ({ + useInfiniteSpendLogUsers: vi.fn(), })); vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers", () => ({ @@ -23,9 +23,9 @@ vi.mock("@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers", () => ({ })); import { useInfiniteSpendLogEndUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers"; +import { useInfiniteSpendLogUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogUsers"; import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; -import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; const emptyInfiniteQuery = { data: { pages: [], pageParams: [] }, @@ -37,16 +37,10 @@ const emptyInfiniteQuery = { const LOGS_WINDOW = { start_date: "2026-07-23 00:00:00", end_date: "2026-07-24 00:00:00" }; -function renderFilters(filters: Record = {}, showUserIdFilter = true) { +function renderFilters(filters: Record = {}) { const set = vi.fn(); renderWithProviders( - filters[id]} - set={set} - teams={[]} - logsWindow={LOGS_WINDOW} - showUserIdFilter={showUserIdFilter} - />, + filters[id]} set={set} teams={[]} logsWindow={LOGS_WINDOW} />, ); return { set }; } @@ -61,8 +55,8 @@ describe("RequestLogsFilters", () => { vi.mocked(useInfiniteModelInfo).mockReturnValue( emptyInfiniteQuery as unknown as ReturnType, ); - vi.mocked(useInfiniteUsers).mockReturnValue( - emptyInfiniteQuery as unknown as ReturnType, + vi.mocked(useInfiniteSpendLogUsers).mockReturnValue( + emptyInfiniteQuery as unknown as ReturnType, ); vi.mocked(useInfiniteSpendLogEndUsers).mockReturnValue( emptyInfiniteQuery as unknown as ReturnType, @@ -97,31 +91,27 @@ describe("RequestLogsFilters", () => { expect(labels[1].compareDocumentPosition(labels[2]) & Node.DOCUMENT_POSITION_FOLLOWING).toBeTruthy(); }); - it("selects a user by display name while storing the user ID filter", async () => { - vi.mocked(useInfiniteUsers).mockReturnValue({ + it("selects an internal user value from the caller's visible spend logs", async () => { + vi.mocked(useInfiniteSpendLogUsers).mockReturnValue({ ...emptyInfiniteQuery, data: { pages: [ { - users: [{ user_id: "user-1", user_alias: "Alice", user_email: "alice@example.com" }], - page: 1, - page_size: 50, - total: 1, - total_pages: 1, + data: ["alice@example.com"], + meta: { page: 1, page_size: 50, has_more: false }, + links: { self: "", next: null }, }, ], pageParams: [1], }, - } as unknown as ReturnType); + } as unknown as ReturnType); const user = userEvent.setup(); const { set } = renderFilters(); await user.click(await screen.findByPlaceholderText("Search an internal user")); - expect(await screen.findByText("Alice")).toBeInTheDocument(); - expect(screen.getByText("alice@example.com | User ID: user-1")).toBeInTheDocument(); - await user.click(screen.getByText("Alice")); + await user.click(await screen.findByText("alice@example.com")); - expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.USER_ID, "user-1"); + expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.USER_ID, "alice@example.com"); }); it("pushes the User ID picker query to the paginated user lookup", async () => { @@ -132,28 +122,26 @@ describe("RequestLogsFilters", () => { await user.click(input); await user.type(input, "alice@example.com"); - await waitFor(() => expect(useInfiniteUsers).toHaveBeenCalledWith(50, "alice@example.com")); + await waitFor(() => expect(useInfiniteSpendLogUsers).toHaveBeenCalledWith(LOGS_WINDOW, 50, "alice@example.com")); }); it("loads the next page when the User ID list is scrolled near the end", async () => { const fetchNextPage = vi.fn(); - vi.mocked(useInfiniteUsers).mockReturnValue({ + vi.mocked(useInfiniteSpendLogUsers).mockReturnValue({ ...emptyInfiniteQuery, fetchNextPage, hasNextPage: true, data: { pages: [ { - users: [{ user_id: "user-1", user_alias: "Alice", user_email: "alice@example.com" }], - page: 1, - page_size: 50, - total: 51, - total_pages: 2, + data: ["alice@example.com"], + meta: { page: 1, page_size: 50, has_more: true }, + links: { self: "", next: "?page=2" }, }, ], pageParams: [1], }, - } as unknown as ReturnType); + } as unknown as ReturnType); const user = userEvent.setup(); renderFilters(); @@ -167,13 +155,6 @@ describe("RequestLogsFilters", () => { await waitFor(() => expect(fetchNextPage).toHaveBeenCalled()); }); - it("does not show or query the User ID filter for non-admin request logs", () => { - renderFilters({}, false); - - expect(screen.queryByText("User ID")).not.toBeInTheDocument(); - expect(useInfiniteUsers).not.toHaveBeenCalled(); - }); - it("scopes the Key Alias lookup to the selected team", async () => { renderFilters({ [LOG_FILTER_IDS.TEAM_ID]: "team-42" }); @@ -264,15 +245,7 @@ describe("RequestLogsFilters", () => { it("scopes the End User lookup to the window the logs table is showing", async () => { const otherWindow = { start_date: "2026-01-01 00:00:00", end_date: "2026-01-02 00:00:00" }; - renderWithProviders( - undefined} - set={vi.fn()} - teams={[]} - logsWindow={otherWindow} - showUserIdFilter - />, - ); + renderWithProviders( undefined} set={vi.fn()} teams={[]} logsWindow={otherWindow} />); await waitFor(() => expect(useInfiniteSpendLogEndUsers).toHaveBeenCalledWith(otherWindow, 50, undefined)); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx index 47e0bad6f62..017260230dd 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.tsx @@ -3,9 +3,9 @@ import { useMemo, useState } from "react"; import { useInfiniteSpendLogEndUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogEndUsers"; +import { useInfiniteSpendLogUsers } from "@/app/(dashboard)/hooks/spendLogs/useSpendLogUsers"; import { useInfiniteKeyAliases } from "@/app/(dashboard)/hooks/keys/useKeyAliases"; import { useInfiniteModelInfo } from "@/app/(dashboard)/hooks/models/useModels"; -import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; import { DataTableFilterField } from "@/components/shared/DataTable"; import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; @@ -145,9 +145,18 @@ function ModelFilterField({ value, onChange }: { value: string; onChange: (value ); } -function UserIdFilterField({ value, onChange }: { value: string; onChange: (value: string | undefined) => void }) { +function UserIdFilterField({ + value, + onChange, + logsWindow, +}: { + value: string; + onChange: (value: string | undefined) => void; + logsWindow: LogsWindow; +}) { const [search, setSearch] = useState(""); - const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteUsers( + const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteSpendLogUsers( + logsWindow, PAGE_SIZE, emptyToUndefined(search), ); @@ -155,14 +164,10 @@ function UserIdFilterField({ value, onChange }: { value: string; onChange: (valu const options = useMemo(() => { const seen = new Set(); return (data?.pages ?? []).flatMap((page) => - page.users.flatMap((user) => { - if (!user.user_id || seen.has(user.user_id)) return []; - seen.add(user.user_id); - const label = user.user_alias || user.user_email || user.user_id; - const email = user.user_email && user.user_email !== label ? user.user_email : ""; - const sublabel = - user.user_id === label ? email : [email, `User ID: ${user.user_id}`].filter(Boolean).join(" | "); - return [{ label, value: user.user_id, sublabel }]; + page.data.flatMap((userId) => { + if (!userId || seen.has(userId)) return []; + seen.add(userId); + return [{ label: userId, value: userId }]; }), ); }, [data]); @@ -284,10 +289,9 @@ interface RequestLogsFiltersProps { set: (columnId: string, value: unknown) => void; teams: Team[]; logsWindow: LogsWindow; - showUserIdFilter: boolean; } -export function RequestLogsFilters({ get, set, teams, logsWindow, showUserIdFilter }: RequestLogsFiltersProps) { +export function RequestLogsFilters({ get, set, teams, logsWindow }: RequestLogsFiltersProps) { const valueOf = (id: string): string => asString(get(id)); const setter = (id: string) => (next: string | undefined) => set(id, next); @@ -321,12 +325,11 @@ export function RequestLogsFilters({ get, set, teams, logsWindow, showUserIdFilt teamId={valueOf(LOG_FILTER_IDS.TEAM_ID)} /> - {showUserIdFilter && ( - - )} + void; teams: Team[]; logsWindow: LogsWindow; - showUserIdFilter: boolean; toolbarChildren?: ReactNode; } @@ -70,7 +69,6 @@ export function RequestLogsTable({ onSessionClick, teams, logsWindow, - showUserIdFilter, toolbarChildren, }: RequestLogsTableProps) { const [filtersOpen, setFiltersOpen] = useState(false); @@ -124,15 +122,7 @@ export function RequestLogsTable({ title="Filters" description="Narrow down request logs" > - {({ get, set }) => ( - - )} + {({ get, set }) => } )} diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx index 17d26dc00f3..26c5bda1593 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx @@ -45,7 +45,6 @@ const defaultProps = { userRole: "Admin" as string | null, userID: "user-1" as string | null, columnFilters: [] as ColumnFiltersState, - filterByCurrentUser: false, activeTab: "request logs", isLiveTail: false, startTime: "2025-01-01T00:00:00", @@ -181,17 +180,16 @@ describe("useLogFilterLogic", () => { }); }); - describe("filterByCurrentUser", () => { - it("scopes to the current user when no explicit user filter is set", async () => { - renderFilterHook({ filterByCurrentUser: true }); + describe("user scope", () => { + it("leaves an empty user filter for the backend to authorize", async () => { + renderFilterHook(); await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); - expect(lastCallParams()?.params).toMatchObject({ user_id: "user-1" }); + expect(lastCallParams()?.params?.user_id).toBeUndefined(); }); - it("lets an explicit user filter win over the current-user scope", async () => { + it("sends an explicit user filter for the backend to intersect with authorization", async () => { renderFilterHook({ - filterByCurrentUser: true, columnFilters: [{ id: LOG_FILTER_IDS.USER_ID, value: "someone-else" }], }); diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 78ecb52c184..474f51e93b3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -99,7 +99,6 @@ export function useLogFilterLogic({ userRole, userID, columnFilters, - filterByCurrentUser, activeTab, isLiveTail, startTime, @@ -113,7 +112,6 @@ export function useLogFilterLogic({ userRole: string | null; userID: string | null; columnFilters: ColumnFiltersState; - filterByCurrentUser: boolean | null; activeTab: string; isLiveTail: boolean; startTime: string; @@ -137,7 +135,6 @@ export function useLogFilterLogic({ endTime, isCustomDate, columnFilters, - filterByCurrentUser ? userID : null, sortBy, sortOrder, ], @@ -167,7 +164,7 @@ export function useLogFilterLogic({ team_id: getFilterValue(columnFilters, LOG_FILTER_IDS.TEAM_ID), request_id: getFilterValue(columnFilters, LOG_FILTER_IDS.REQUEST_ID), session_id: getFilterValue(columnFilters, LOG_FILTER_IDS.SESSION_ID), - user_id: userIdFilter ?? (filterByCurrentUser ? userID ?? undefined : undefined), + user_id: userIdFilter, end_user: getFilterValue(columnFilters, LOG_FILTER_IDS.END_USER), status_filter: getFilterValue(columnFilters, LOG_FILTER_IDS.STATUS), model_id: getFilterValue(columnFilters, LOG_FILTER_IDS.MODEL_ID), diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1e46ae9c577..75e222f7472 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7529,6 +7529,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/spend_logs/users": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Spend Log Users + * @description The distinct internal users appearing in spend logs the caller can read. + */ + get: operations["list_spend_log_users_management_v1_spend_logs_users_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/mcp-rest/test/connection": { parameters: { query?: never; @@ -45400,6 +45420,46 @@ export interface operations { }; }; }; + list_spend_log_users_management_v1_spend_logs_users_get: { + parameters: { + query: { + /** @description Window start (UTC when no offset is given) */ + "filter[startTime][gte]": string; + /** @description Window end (UTC when no offset is given) */ + "filter[startTime][lte]": string; + /** @description Case-insensitive partial match on the internal user id */ + q?: string | null; + /** @description Page number */ + page?: number; + /** @description Page size */ + page_size?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["FacetListResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; test_connection_mcp_rest_test_connection_post: { parameters: { query?: never; From fac2b6b56b4020b85c423bac11f5b78367b8833e Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 12:12:05 -0400 Subject: [PATCH 016/121] refactor: derive request log scope immutably Resolve the authorized own-user and permitted-team predicates once and add regression coverage for explicit-user intersection, unfiltered team scope, and team lookup failure fallback. Co-Authored-By: Codex --- .../spend_management_endpoints.py | 78 +++++++---- .../test_spend_management_endpoints.py | 125 +++++++++++++++++- 2 files changed, 174 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 99d870f5ad4..ed2ecd8325a 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2426,8 +2426,23 @@ async def ui_view_spend_logs( user_api_key_dict=user_api_key_dict, request_id=request_id, ) - permitted_team_ids: list[str] | None = None - scope_to_caller_user = False + user_scope_applies: Final = ( + not is_request_id_lookup + and not is_admin_view + and team_id is None + and _can_user_view_spend_log(user_api_key_dict=user_api_key_dict) + ) + permitted_team_ids: Final = ( + await _get_permitted_team_ids_for_spend_logs_or_empty( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + ) + if user_scope_applies + else () + ) + explicit_user_requires_caller_scope: Final = ( + user_scope_applies and not permitted_team_ids and user_id is not None + ) if not is_request_id_lookup and not is_admin_view: if team_id is not None: can_view_team: Final = await _can_team_member_view_log( @@ -2441,31 +2456,22 @@ async def ui_view_spend_logs( detail={"error": f"Not authorized to view team spend for team_id={team_id}"}, ) where_conditions["team_id"] = team_id - else: - if _can_user_view_spend_log(user_api_key_dict=user_api_key_dict): - try: - permitted_team_ids = await _get_permitted_team_ids_for_spend_logs( - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - ) - except Exception: - permitted_team_ids = [] - if permitted_team_ids: - if user_id is None: - where_conditions.pop("user", None) - where_conditions["OR"] = [ - {"user": user_api_key_dict.user_id}, - {"team_id": {"in": permitted_team_ids}}, - ] + elif user_scope_applies: + if permitted_team_ids: + if user_id is None: + where_conditions.pop("user", None) + where_conditions["OR"] = [ + {"user": user_api_key_dict.user_id}, + {"team_id": {"in": permitted_team_ids}}, + ] + else: + if user_id is None: + where_conditions["user"] = user_api_key_dict.user_id else: - if user_id is None: - where_conditions["user"] = user_api_key_dict.user_id - else: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - {"user": user_api_key_dict.user_id} - ] - scope_to_caller_user = True - where_conditions.pop("team_id", None) + where_conditions["AND"] = where_conditions.get("AND", []) + [ + {"user": user_api_key_dict.user_id} + ] + where_conditions.pop("team_id", None) # Calculate skip value for pagination skip: Final = (page - 1) * page_size @@ -2509,13 +2515,13 @@ async def ui_view_spend_logs( p += 1 # Multi-team OR filter: (user = $X OR team_id = ANY($Y)) - if permitted_team_ids is not None and len(permitted_team_ids) > 0: + if permitted_team_ids: or_clause: Final = f'("user" = ${p} OR team_id = ANY(${p + 1}::text[]))' sql_params.append(user_api_key_dict.user_id) sql_params.append(permitted_team_ids) p += 2 sql_conditions.append(or_clause) - elif scope_to_caller_user: + elif explicit_user_requires_caller_scope: sql_conditions.append(f'"user" = ${p}') sql_params.append(user_api_key_dict.user_id) p += 1 @@ -4283,3 +4289,19 @@ async def _get_permitted_team_ids_for_spend_logs( ): permitted.append(team_obj.team_id) return permitted + + +async def _get_permitted_team_ids_for_spend_logs_or_empty( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[str, ...]: + """Resolve permitted teams once, falling back to the caller's own-user scope.""" + try: + return tuple( + await _get_permitted_team_ids_for_spend_logs( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + ) + ) + except Exception: + return () diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 81512cd8e66..87bbb1c2f80 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -150,7 +150,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): return where -def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None): +def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. @@ -177,6 +177,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return [{col: value, "_count": {col: n}} for value, n in tallied.items()] async def query_raw(self, sql_query, *params): + if query_observer is not None: + query_observer(sql_query, params) if "mcp_tool_call_count" in sql_query: return [] filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) @@ -1320,6 +1322,127 @@ async def test_ui_view_spend_logs_internal_user_scoped_without_user_id( app.dependency_overrides.pop(ps.user_api_key_auth, None) +@pytest.mark.asyncio +async def test_ui_view_spend_logs_explicit_user_filter_cannot_escape_own_scope(client, monkeypatch): + caller_log = { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "caller@example.com", + "team_id": None, + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + observed_queries = [] + + def observe_query(sql_query, params): + if 'FROM "LiteLLM_SpendLogs"' in sql_query: + observed_queries.append((sql_query, params)) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([caller_log], lambda _where: [], query_observer=observe_query), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=[]), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller@example.com" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "user_id": "someone-else@example.com", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json()["data"] == [] + page_sql, page_params = next((sql, params) for sql, params in observed_queries if "SELECT\n" in sql) + assert page_sql.count('"user" = $') == 2 + assert page_params[2:4] == ("someone-else@example.com", "caller@example.com") + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_without_user_filter_includes_permitted_team_scope(client, monkeypatch): + caller_log = { + "id": "log1", + "request_id": "req1", + "api_key": "sk-test-key", + "user": "team-admin@example.com", + "team_id": None, + "spend": 0.05, + "startTime": datetime.datetime.now(timezone.utc).isoformat(), + "model": "gpt-4", + } + member_log = {**caller_log, "id": "log2", "request_id": "req2", "user": "member@example.com", "team_id": "team-9"} + outside_log = { + **caller_log, + "id": "log3", + "request_id": "req3", + "user": "outside@example.com", + "team_id": "outside-team", + } + + def filter_by_scope(where): + if {"multi_team": True} in where.get("OR", []) and "user" not in where: + return [caller_log, member_log] + return [caller_log, member_log, outside_log] + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([caller_log, member_log, outside_log], filter_by_scope), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=["team-9"]), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin@example.com" + ) + + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert [row["request_id"] for row in response.json()["data"]] == ["req1", "req2"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_permitted_team_scope_falls_back_to_own_user_when_lookup_fails(monkeypatch): + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(side_effect=RuntimeError("database unavailable")), + ) + + permitted_team_ids = await spend_management_endpoints._get_permitted_team_ids_for_spend_logs_or_empty( + prisma_client=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="caller@example.com", + ), + ) + + assert permitted_team_ids == () + + @pytest.mark.asyncio async def test_ui_view_spend_logs_team_admin_can_filter_team_spend_by_user(client, monkeypatch): """ From 19eae00d71f85405eec104d15d502a6b5bd9a68b Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 12:16:36 -0400 Subject: [PATCH 017/121] fix(ui): make per-user usage filter searchable Reuse the Global Usage user search and pagination behavior in the Per User report, including empty-result handling. Co-Authored-By: Codex --- .../components/EntityUsage/EntityUsage.tsx | 10 ++++- .../components/UsagePageView.test.tsx | 26 +++++++++++- .../_components/components/UsagePageView.tsx | 41 +++++++++++-------- .../UsageExportHeader.test.tsx | 25 +++++++++++ .../EntityUsageExport/UsageExportHeader.tsx | 22 ++++++++-- .../src/components/EntityUsageExport/index.ts | 1 + 6 files changed, 100 insertions(+), 25 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 956060fc244..6f63b56c372 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -37,7 +37,7 @@ import { Alert, Button, Tooltip } from "antd"; import React, { type ReactNode, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; -import { UsageExportHeader } from "@/components/EntityUsageExport"; +import { UsageExportHeader, type UsageFilterSelectProps } from "@/components/EntityUsageExport"; import type { EntityType } from "@/components/EntityUsageExport/types"; import { agentDailyActivityCall, @@ -97,6 +97,7 @@ interface EntityUsageProps { entityList: EntityList[] | null; premiumUser: boolean; dateValue: DateRangePickerValue; + filterSelectProps?: UsageFilterSelectProps; } const ENTITY_FETCH_FNS: Record Promise> = { @@ -120,6 +121,7 @@ const EntityUsage: React.FC = ({ entityList, userRole, dateValue, + filterSelectProps, }) => { const { teams } = useTeams(); const [selectedTags, setSelectedTags] = useState([]); @@ -678,13 +680,17 @@ const EntityUsage: React.FC = ({ dateValue={dateValue} entityType={entityType} spendData={spendData} - showFilters={entityType !== "team" && entityList !== null && entityList.length > 0} + showFilters={ + entityType !== "team" && + (filterSelectProps?.showSearch === true || (entityList !== null && entityList.length > 0)) + } filterLabel={getFilterLabel(entityType)} filterPlaceholder={getFilterPlaceholder(entityType)} selectedFilters={selectedTags} onFiltersChange={setSelectedTags} filterOptions={getAllTags() || undefined} filterMode={entityType === "user" ? "single" : "multiple"} + filterSelectProps={filterSelectProps} teams={teams || []} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index 9085cf961a9..0c6874d0786 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -46,7 +46,18 @@ vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ })); vi.mock("./EntityUsage/EntityUsage", () => ({ - default: () =>
Entity Usage
, + default: ({ + entityType, + filterSelectProps, + }: { + entityType?: string; + filterSelectProps?: { showSearch?: boolean }; + }) => ( +
+ Entity Usage + {entityType === "user" && filterSelectProps?.showSearch && Searchable user filter} +
+ ), EntityList: [], })); @@ -76,6 +87,7 @@ vi.mock("./UsageViewSelect/UsageViewSelect", async () => { React.createElement("option", { value: "customer" }, "Customer Usage"), tagOption, React.createElement("option", { value: "agent" }, "Agent Usage"), + React.createElement("option", { value: "user" }, "User Usage"), React.createElement("option", { value: "user-agent-activity" }, "User Agent Activity"), ); }; @@ -924,6 +936,18 @@ describe("UsagePage", () => { expect(mockUseInfiniteUsers).toHaveBeenCalledWith(50, undefined); }); + it("should reuse the searchable user filter in the user usage view", async () => { + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + fireEvent.change(screen.getByTestId("usage-view-select"), { target: { value: "user" } }); + + expect(await screen.findByText("Searchable user filter")).toBeInTheDocument(); + }); + it("should deduplicate users across pages", async () => { mockUseInfiniteUsers.mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 494df313ac0..9cd499dd7d2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -38,7 +38,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { all_admin_roles, internalUserRoles } from "@/utils/roles"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import CloudZeroExportModal from "@/components/cloudzero_export_modal"; -import EntityUsageExportModal from "@/components/EntityUsageExport"; +import EntityUsageExportModal, { type UsageFilterSelectProps } from "@/components/EntityUsageExport"; import { Team } from "@/components/key_team_helpers/key_list"; import { gatewayDailyActivityCall, @@ -161,6 +161,26 @@ const UsagePage: React.FC = ({ teams, organizations }) => { } }; + const userFilterSelectProps: UsageFilterSelectProps = { + showSearch: true, + filterOption: false, + onSearch: handleUserSearchChange, + searchValue: userSearchInput, + onPopupScroll: handleUserPopupScroll, + loading: isLoadingUsers, + notFoundContent: isLoadingUsers ? : "No users found", + popupRender: (menu) => ( + <> + {menu} + {isFetchingNextUsersPage && ( +
+ +
+ )} + + ), + }; + // For admins: null means global view (all users), a string means filter by that user // For non-admins: always set to their own user ID const [selectedUserId, setSelectedUserId] = useState(isAdmin ? null : userID || null); @@ -565,29 +585,13 @@ const UsagePage: React.FC = ({ teams, organizations }) => {
Filter by user Date: Thu, 13 Aug 2026 12:27:55 -0400 Subject: [PATCH 018/121] test: remove unrelated session log assertion Drop a stray assertion against a field that is not present in the session pagination fixture. Co-Authored-By: Codex --- .../proxy/spend_tracking/test_spend_management_endpoints.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 87bbb1c2f80..7052e050806 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1774,7 +1774,6 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): assert data["total_pages"] == 2 assert len(data["data"]) == 1 assert data["data"][0]["request_id"] == "req1" - assert data["data"][0]["user"] == "member1" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) From fd45fc581ec6c9be424b29d4c43daf3189bd91a6 Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 16:50:41 +0000 Subject: [PATCH 019/121] fix(model_prices): refresh deprecation dates, add grok-4.6 and gemini 3.1 flash tts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 64 ++++++++++++++++++- model_prices_and_context_window.json | 64 ++++++++++++++++++- 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3b2cdcf5ff7..af16ec7d3de 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17585,6 +17585,7 @@ "supports_tool_choice": true }, "ft:gpt-3.5-turbo-0613": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 4096, @@ -17596,6 +17597,7 @@ "supports_tool_choice": true }, "ft:gpt-3.5-turbo-1106": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -19451,7 +19453,7 @@ "uses_embed_content": true }, "gemini/gemini-embedding-001": { - "deprecation_date": "2028-05-14", + "deprecation_date": "2026-07-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "gemini", "max_input_tokens": 2048, @@ -19627,6 +19629,7 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -19938,6 +19941,7 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_token_cost": 1e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -20234,6 +20238,7 @@ "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "input_cost_per_token_priority": 1.25e-06, @@ -22369,6 +22374,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-16k": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -22503,6 +22509,7 @@ "supports_vision": true }, "gpt-4-turbo-preview": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -40723,6 +40730,48 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "xai/grok-4.6-latest": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "xai/grok-beta": { "input_cost_per_token": 5e-06, "litellm_provider": "xai", @@ -45509,6 +45558,19 @@ "rpm": 10, "gemini_audio_only_live": true }, + "gemini/gemini-3.1-flash-tts-preview": { + "input_cost_per_token": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "audio_speech", + "output_cost_per_token": 2e-05, + "source": "https://ai.google.dev/gemini-api/docs/models/gemini-3.1-flash-tts-preview", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, "litellm_provider": "gemini", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3b2cdcf5ff7..af16ec7d3de 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17585,6 +17585,7 @@ "supports_tool_choice": true }, "ft:gpt-3.5-turbo-0613": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 4096, @@ -17596,6 +17597,7 @@ "supports_tool_choice": true }, "ft:gpt-3.5-turbo-1106": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -19451,7 +19453,7 @@ "uses_embed_content": true }, "gemini/gemini-embedding-001": { - "deprecation_date": "2028-05-14", + "deprecation_date": "2026-07-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "gemini", "max_input_tokens": 2048, @@ -19627,6 +19629,7 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -19938,6 +19941,7 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_token_cost": 1e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -20234,6 +20238,7 @@ "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "input_cost_per_token_priority": 1.25e-06, @@ -22369,6 +22374,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-16k": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -22503,6 +22509,7 @@ "supports_vision": true }, "gpt-4-turbo-preview": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -40723,6 +40730,48 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-4.6": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "xai/grok-4.6-latest": { + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "xai/grok-beta": { "input_cost_per_token": 5e-06, "litellm_provider": "xai", @@ -45509,6 +45558,19 @@ "rpm": 10, "gemini_audio_only_live": true }, + "gemini/gemini-3.1-flash-tts-preview": { + "input_cost_per_token": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 8192, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "audio_speech", + "output_cost_per_token": 2e-05, + "source": "https://ai.google.dev/gemini-api/docs/models/gemini-3.1-flash-tts-preview", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, "litellm_provider": "gemini", From 5edf8e71ce13ce77f1f906a840615d3cb7f069ce Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 13:00:08 -0400 Subject: [PATCH 020/121] chore: rerun CI Generated with AI Co-Authored-By: Codex From d794b613479bc28c095eb7c449d6860d3829c72f Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 13:09:47 -0400 Subject: [PATCH 021/121] chore(ui): bump nanoid to 3.3.18 Update the transitive lockfile entry to the first patched 3.x release so OSV no longer reports GHSA-2v37-7h3g-55p8. Generated with AI Co-Authored-By: Codex --- 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 515a992bc85..b36b07631e3 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.17", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.17.tgz", - "integrity": "sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==", + "version": "3.3.18", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.18.tgz", + "integrity": "sha512-DTg4MJbGMWkfi6VZFdNt2/caMbQy4Ou+Op/hJQvGEWcnVfoA1QA+xzRKAzw9jD6+GVOOeYr/mIcuDSdug6F6+w==", "funding": [ { "type": "github", From c30b043a5145384ac05eab8c14c206ad570708fc Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 17:11:30 +0000 Subject: [PATCH 022/121] add tpm/rpm to gemini-3.1-flash-tts-preview entry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 +++- model_prices_and_context_window.json | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index af16ec7d3de..b0b6b5c33b5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45569,7 +45569,9 @@ "source": "https://ai.google.dev/gemini-api/docs/models/gemini-3.1-flash-tts-preview", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "tpm": 4000000, + "rpm": 10 }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index af16ec7d3de..b0b6b5c33b5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45569,7 +45569,9 @@ "source": "https://ai.google.dev/gemini-api/docs/models/gemini-3.1-flash-tts-preview", "supported_endpoints": [ "/v1/audio/speech" - ] + ], + "tpm": 4000000, + "rpm": 10 }, "gemini-2.5-flash-preview-tts": { "input_cost_per_token": 3e-07, From 80f49024a3ce7f1bb5c98a3eab2435b614ee348b Mon Sep 17 00:00:00 2001 From: Daniel Meismer Date: Thu, 13 Aug 2026 13:14:17 -0400 Subject: [PATCH 023/121] chore(ui): bump nanoid to 3.3.18 Update the transitive lockfile entry to the first patched 3.x release so OSV no longer reports GHSA-2v37-7h3g-55p8. Generated with AI Co-Authored-By: Codex --- 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 515a992bc85..b36b07631e3 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.17", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.17.tgz", - "integrity": "sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==", + "version": "3.3.18", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.18.tgz", + "integrity": "sha512-DTg4MJbGMWkfi6VZFdNt2/caMbQy4Ou+Op/hJQvGEWcnVfoA1QA+xzRKAzw9jD6+GVOOeYr/mIcuDSdug6F6+w==", "funding": [ { "type": "github", From 7b28476bfc9f53e3517aca49138ee3f4b8a53e2d Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 17:33:32 +0000 Subject: [PATCH 024/121] retrigger ci Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From f91e698adbd00b88c3114a276b8a3d0095302ffc Mon Sep 17 00:00:00 2001 From: MUSE Date: Tue, 21 Jul 2026 11:50:53 +0900 Subject: [PATCH 025/121] fix(batch): avoid reading a nonexistent output artifact for completed batches Completed batches that contain only failed requests do not generate an output file, leaving output_file_id unset while the failures are recorded through error_file_id instead. The completion handler attempted to read the output payload regardless of whether an output file actually existed. During retrieve polling this caused the logging pipeline to fail with "Output file id is None cannot retrieve file content", preventing normal completion bookkeeping from running. Skip output retrieval when no output file is available and return an empty batch summary (zero usage, zero cost, no model entries). The lower-level file retrieval helper still reports an error if it is called directly with an invalid or missing file identifier. Closes #33987 --- litellm/batches/batch_utils.py | 11 ++++++++++ .../test_litellm/batches/test_batch_utils.py | 21 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index e73b887ae0a..cfd37864eac 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -58,6 +58,17 @@ async def _handle_completed_batch( model_name: Optional model name litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) """ + # A completed batch whose request lines all failed has no output file - the + # results are written to a separate error_file_id and output_file_id is None. + # There is nothing to price or measure, so report an empty result set instead + # of calling _fetch_batch_output_file_content, which raises on a missing + # output file. Without this guard the logging worker crashes on every + # aretrieve_batch poll and the completed batch's zero-cost accounting is lost. + # The generic retrieval helper keeps raising for callers that explicitly ask + # for a missing output file. + if batch.output_file_id is None: + return 0.0, Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), [] + file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params) if ( diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 523b512e4cf..dbf1102b31d 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -976,6 +976,27 @@ async def test_handle_completed_batch_orchestration(monkeypatch): assert models == ["gpt-4o"] +@pytest.mark.asyncio +async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): + """ + Regression: an all-error batch completes with output_file_id=None (results go + to a separate error_file_id). _handle_completed_batch must report an empty + result set - zero cost, zero usage, no models - instead of letting the file + fetch raise "Output file id is None" on every aretrieve_batch logging poll. + """ + # The output-file fetch must not even be attempted when there is no output file. + async def _must_not_fetch(*args, **kwargs): + pytest.fail("_fetch_batch_output_file_content should not be called") + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch) + + cost, usage, models = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") + + assert cost == 0.0 + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) + assert models == [] + + @pytest.mark.asyncio async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch): raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}] From f5ccc4ebdb764a826dcf398335c03bde83610b1b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:01:35 -0700 Subject: [PATCH 026/121] feat(lint): exempt TypedDict-annotated dict literals from LIT002 --- scripts/check_type_discipline.py | 108 +++++++++++++++--- .../test_check_type_discipline.py | 49 ++++++++ type-discipline-budget.json | 2 +- 3 files changed, 142 insertions(+), 17 deletions(-) diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index ce9eb391d55..e21693b2c9b 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -18,13 +18,22 @@ LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehens Catches the unannotated seed-then-mutate pattern LIT001 cannot see (`acc = []`). Build the value in one shot and freeze it: a `tuple`/`frozenset` wrapping a generator (`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / - NamedTuple / ReadOnly TypedDict, or (if it really must be dynamic) a - MappingProxyType wrapping a dict literal or comprehension. Generator expressions - and freezing-wrapper calls (`tuple(...)`, `frozenset(...)`, + NamedTuple, a TypedDict-annotated dict literal, or (if it really must be + dynamic) a MappingProxyType wrapping a dict literal or comprehension. Generator + expressions and freezing-wrapper calls (`tuple(...)`, `frozenset(...)`, `MappingProxyType(...)`) are not construction and pass, as does the value passed directly to a wrapper: it is frozen before it can escape, though anything mutable nested inside it still counts. Annotation-internal lists - (`Callable[[int], str]`) are exempt. Suppress with `# mutable-ok: `. + (`Callable[[int], str]`) are exempt. A dict literal whose assignment is + annotated with a TypedDict (`x: Final[MyTD] = {...}`; bare `x: Final = {...}` + does not qualify) is a fixed-shape build basedpyright checks key-by-key against + fields LIT012 keeps ReadOnly, not a growable accumulator, so it is exempt along + with the dict literals nested in it (nested TypedDict fields); any other + construction inside still counts. Detection is name-based: Final/ClassVar/ + Optional (and Annotated's first argument) unwrap, and any remaining named head + outside the mutable collections and Mapping/Any/object is taken to be a + TypedDict, since a dict literal assigned to any other named type would not + survive basedpyright. Suppress with `# mutable-ok: `. LIT003 noqa suppression without rule codes or without a reason. Required shape: `# noqa: TID251 # ` LIT004 pyright/mypy ignore without bracketed codes or without a reason. @@ -138,6 +147,14 @@ MUTABLE_CONSTRUCTORS = frozenset(( # qualified `collections.deque(...)` still counts. QUALIFIED_CONSTRUCTORS = MUTABLE_CONSTRUCTORS - frozenset(("dict", "list", "set")) FREEZING_WRAPPERS = frozenset(("tuple", "frozenset", "MappingProxyType")) +# Wrappers unwrapped when deciding whether an assignment's annotation names a +# TypedDict (the LIT002 dict-literal exemption); bare, they name no type. Annotated +# is handled separately: only its first argument is type syntax. +TYPEDDICT_ANNOTATION_WRAPPERS = frozenset(("Final", "ClassVar", "Optional")) +# Heads that can type a dict literal without being a TypedDict. Every other named +# head counts as one: a dict literal assigned to any other named type would not +# survive basedpyright, which is the second gate behind this name-based check. +NON_TYPEDDICT_HEADS = MUTABLE_COLLECTIONS | frozenset(("Mapping", "Any", "object")) UNSAFE_GUARDS = frozenset(("TypeGuard", "TypeIs")) READONLY_QUALIFIER = "ReadOnly" # Qualifiers ReadOnly may nest under, in any order (PEP 705); for Annotated only the @@ -270,6 +287,14 @@ def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, . # --------------------------------------------------------------------------- # +def _head_name(node: ast.expr) -> str | None: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + def _is_literal_subscript(node: ast.AST) -> bool: if not isinstance(node, ast.Subscript): return False @@ -485,6 +510,58 @@ def _frozen_argument_ids(tree: ast.AST) -> frozenset[int]: ) +def _is_typeddict_annotation(annotation: ast.expr) -> bool: + """True iff the annotation names a TypedDict, by the name-based heuristic. + + Final/ClassVar/Optional unwrap (as does Annotated's first argument, the only + one that is type syntax), string forward references are parsed, and whatever + named head remains counts as a TypedDict unless it is a mutable collection or + Mapping/Any/object -- the heads that can type a dict literal without being + one. Bare wrappers (`x: Final = ...`) name no type and never qualify. + """ + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + try: + inner = ast.parse(annotation.value, mode="eval").body + except SyntaxError: + return False + return _is_typeddict_annotation(inner) + if isinstance(annotation, ast.Subscript): + head = _head_name(annotation.value) + if head in TYPEDDICT_ANNOTATION_WRAPPERS: + return _is_typeddict_annotation(annotation.slice) + if head == "Annotated": + first = annotation.slice.elts[0] if isinstance(annotation.slice, ast.Tuple) and annotation.slice.elts else None + return first is not None and _is_typeddict_annotation(first) + return head is not None and head not in NON_TYPEDDICT_HEADS + name = _head_name(annotation) + return ( + name is not None + and name not in NON_TYPEDDICT_HEADS + and name not in TYPEDDICT_ANNOTATION_WRAPPERS + and name != "Annotated" + ) + + +def _typeddict_build_ids(tree: ast.AST) -> frozenset[int]: + """ids() of every dict literal built under a TypedDict-annotated assignment. + + `x: Final[MyTD] = {...}` is a fixed-shape build: basedpyright checks each key + against the declared fields, which LIT012 keeps ReadOnly, so nothing here is + the seed-then-mutate accumulator LIT002 hunts. Dict literals nested in the + value (nested TypedDict fields) share the exemption; any other construction + inside it still counts, and a bare `x: Final = {...}` stays flagged. + """ + return frozenset( + id(sub) + for node in ast.walk(tree) + if isinstance(node, ast.AnnAssign) + and isinstance(node.value, ast.Dict) + and _is_typeddict_annotation(node.annotation) + for sub in ast.walk(node.value) + if isinstance(sub, ast.Dict) + ) + + def _construction_kind(node: ast.expr) -> str | None: """Human label if `node` builds a mutable collection, else None.""" if isinstance(node, ast.List): @@ -511,8 +588,14 @@ def _construction_kind(node: ast.expr) -> str | None: def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) -> Iterator[Violation]: in_annotation = _annotation_node_ids(tree) frozen_arguments = _frozen_argument_ids(tree) + typeddict_builds = _typeddict_build_ids(tree) for node in ast.walk(tree): - if not isinstance(node, ast.expr) or id(node) in in_annotation or id(node) in frozen_arguments: + if ( + not isinstance(node, ast.expr) + or id(node) in in_annotation + or id(node) in frozen_arguments + or id(node) in typeddict_builds + ): continue kind = _construction_kind(node) if kind is None or node.lineno in comments.mutable_ok_lines: @@ -521,9 +604,10 @@ def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) path, node.lineno, "LIT002", f"mutable {kind}: this builds a collection that can be grown or rewritten. " f"Build it in one shot and freeze it -- a tuple/frozenset wrapping a generator " - f"(`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / NamedTuple " - f"/ ReadOnly TypedDict, or (if it really must be dynamic) a MappingProxyType " - f"wrapping a dict literal or comprehension (suppress: `# mutable-ok: `)", + f"(`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / NamedTuple, " + f"a TypedDict-annotated dict literal (`x: Final[MyTD] = {{...}}`), or (if it " + f"really must be dynamic) a MappingProxyType wrapping a dict literal or " + f"comprehension (suppress: `# mutable-ok: `)", ) @@ -851,14 +935,6 @@ def iter_param_violations(path: Path, tree: ast.AST, comments: Comments) -> Iter # --------------------------------------------------------------------------- # -def _head_name(node: ast.expr) -> str | None: - if isinstance(node, ast.Name): - return node.id - if isinstance(node, ast.Attribute): - return node.attr - return None - - def _base_names(cls: ast.ClassDef) -> frozenset[str]: """The names of a class's bases; a subscripted base (`Foo[int]`) counts as `Foo`.""" return frozenset( diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 2870a803db8..78268a6daa3 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -199,6 +199,55 @@ def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): assert "LIT002" not in codes +def test_typeddict_annotated_dict_literal_is_exempt(tmp_path): + assert "LIT002" not in _codes( + tmp_path, "from typing import Final\nfrom foo import MyTD\nx: Final[MyTD] = {'a': 1}\n" + ) + assert "LIT002" not in _codes(tmp_path, "from foo import MyTD\nx: MyTD = {'a': 1}\n") + assert "LIT002" not in _codes(tmp_path, "from typing import Final\nx: Final['MyTD'] = {'a': 1}\n") + assert "LIT002" not in _codes(tmp_path, "import foo\nfrom typing import Final\nx: Final[foo.MyTD] = {'a': 1}\n") + + +def test_wrapped_typeddict_annotations_share_the_exemption(tmp_path): + assert "LIT002" not in _codes( + tmp_path, "from typing import Final, Optional\nx: Final[Optional[MyTD]] = {'a': 1}\n" + ) + assert "LIT002" not in _codes( + tmp_path, "from typing import Annotated, Final\nx: Final[Annotated[MyTD, 'meta']] = {'a': 1}\n" + ) + assert "LIT002" not in _codes( + tmp_path, "from typing import ClassVar\nclass C:\n x: ClassVar[MyTD] = {'a': 1}\n" + ) + + +def test_bare_final_dict_literal_still_counts(tmp_path): + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import ClassVar\nclass C:\n x: ClassVar = {'a': 1}\n") + + +def test_non_typeddict_annotations_do_not_exempt(tmp_path): + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[dict[str, int]] = {'a': 1}\n") + assert "LIT002" in _codes( + tmp_path, "from collections.abc import Mapping\nfrom typing import Final\nx: Final[Mapping[str, int]] = {'a': 1}\n" + ) + assert "LIT002" in _codes(tmp_path, "from typing import Any, Final\nx: Final[Any] = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[object] = {'a': 1}\n") + + +def test_typeddict_exemption_covers_only_dict_literals(tmp_path): + # A TypedDict cannot be built from a comprehension (its keys are fixed literals), + # and `dict(...)` is the constructor call the rule targets, so neither is exempt. + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[MyTD] = dict(a=1)\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[MyTD] = {k: 1 for k in ('a',)}\n") + + +def test_nested_dict_literals_share_the_typeddict_exemption(tmp_path): + assert "LIT002" not in _codes( + tmp_path, "from typing import Final\nx: Final[Outer] = {'inner': {'a': 1}, 'steps': ({'b': 2},)}\n" + ) + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[Outer] = {'tags': ['a']}\n") + + # --------------------------------------------------------------------------- # # Casts (LIT006) # --------------------------------------------------------------------------- # diff --git a/type-discipline-budget.json b/type-discipline-budget.json index a7286d9a89a..03191c460c0 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23001 }, "LIT002": { - "limit": 27146 + "limit": 26916 }, "LIT003": { "limit": 269 From 316732b3ae4682f652ee0faffb35b5a7878c6da8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:01:35 -0700 Subject: [PATCH 027/121] fix(scripts): unwrap PEP 604 unions in LIT002 TypedDict detection --- scripts/check_type_discipline.py | 14 +++++++++----- tests/test_litellm/test_check_type_discipline.py | 4 ++-- type-discipline-budget.json | 2 +- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index e21693b2c9b..0706c8a7bd8 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -30,7 +30,8 @@ LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehens fields LIT012 keeps ReadOnly, not a growable accumulator, so it is exempt along with the dict literals nested in it (nested TypedDict fields); any other construction inside still counts. Detection is name-based: Final/ClassVar/ - Optional (and Annotated's first argument) unwrap, and any remaining named head + Optional (and Annotated's first argument) unwrap, a PEP 604 union + (`MyTD | None`) qualifies through either arm, and any remaining named head outside the mutable collections and Mapping/Any/object is taken to be a TypedDict, since a dict literal assigned to any other named type would not survive basedpyright. Suppress with `# mutable-ok: `. @@ -514,10 +515,11 @@ def _is_typeddict_annotation(annotation: ast.expr) -> bool: """True iff the annotation names a TypedDict, by the name-based heuristic. Final/ClassVar/Optional unwrap (as does Annotated's first argument, the only - one that is type syntax), string forward references are parsed, and whatever - named head remains counts as a TypedDict unless it is a mutable collection or - Mapping/Any/object -- the heads that can type a dict literal without being - one. Bare wrappers (`x: Final = ...`) name no type and never qualify. + one that is type syntax), a PEP 604 union qualifies through either arm, string + forward references are parsed, and whatever named head remains counts as a + TypedDict unless it is a mutable collection or Mapping/Any/object -- the heads + that can type a dict literal without being one. Bare wrappers + (`x: Final = ...`) name no type and never qualify. """ if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): try: @@ -525,6 +527,8 @@ def _is_typeddict_annotation(annotation: ast.expr) -> bool: except SyntaxError: return False return _is_typeddict_annotation(inner) + if isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr): + return _is_typeddict_annotation(annotation.left) or _is_typeddict_annotation(annotation.right) if isinstance(annotation, ast.Subscript): head = _head_name(annotation.value) if head in TYPEDDICT_ANNOTATION_WRAPPERS: diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 78268a6daa3..84dd547ad80 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -218,6 +218,8 @@ def test_wrapped_typeddict_annotations_share_the_exemption(tmp_path): assert "LIT002" not in _codes( tmp_path, "from typing import ClassVar\nclass C:\n x: ClassVar[MyTD] = {'a': 1}\n" ) + assert "LIT002" not in _codes(tmp_path, "from typing import Final\nx: Final[MyTD | None] = {'a': 1}\n") + assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[dict[str, int] | None] = {'a': 1}\n") def test_bare_final_dict_literal_still_counts(tmp_path): @@ -235,8 +237,6 @@ def test_non_typeddict_annotations_do_not_exempt(tmp_path): def test_typeddict_exemption_covers_only_dict_literals(tmp_path): - # A TypedDict cannot be built from a comprehension (its keys are fixed literals), - # and `dict(...)` is the constructor call the rule targets, so neither is exempt. assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[MyTD] = dict(a=1)\n") assert "LIT002" in _codes(tmp_path, "from typing import Final\nx: Final[MyTD] = {k: 1 for k in ('a',)}\n") diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 03191c460c0..909afb0a9db 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23001 }, "LIT002": { - "limit": 26916 + "limit": 26912 }, "LIT003": { "limit": 269 From 9a9e7a58d3edb8421660aeaebdf1bfdbb4d74a21 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 12 Aug 2026 23:21:51 -0400 Subject: [PATCH 028/121] fix(spend): give a batch's cost row a primary key of its own request_id is the primary key of LiteLLM_SpendLogs and the flush inserts with skip_duplicates, so a spend log whose id already exists is dropped with no error raised and a "processed 1 spend log" line still logged. Batch cost accounting produced exactly such an id twice over, and on a proxy with message redaction enabled no batch cost row could be written at all. get_spend_logs_id derived the id by md5-hashing the response for two call types, aretrieve_batch and acreate_file. Redaction makes that hash a constant: perform_redaction returns the fixed {"text": "redacted-by-litellm"} placeholder for any shape it cannot redact, which is what a batch object and a file body both become, so every such row hashed to md5('{"text": "redacted-by-litellm"}') = 00fcbef15a3b0097e14b0ca016ed30a0 regardless of provider, user, or amount. The first row to claim that id owned it and every later row was discarded. Verified against a live proxy: four payloads spanning two providers and three distinct spend values all computed that id, and the table held one acreate_file row dating to 2025-05-25, the row that had claimed it. Keying off the batch's own identity instead is necessary but not sufficient, because creating a batch already writes an acreate_batch row under exactly that id, so the cost row becomes a duplicate of the batch's own creation row. Also verified live: after the hash was removed the poller computed and flushed a batch's cost, and the only row carrying that id was the acreate_batch row from when the batch was submitted. The id now comes from the response's own id, then the standard logging payload's id, then litellm_call_id, and a batch cost row is namespaced with a _batch_cost suffix so it cannot collide with the creation row. The middle term is what keeps this correct under redaction: that payload is built from the unredacted response, so it still carries the batch id after redaction has flattened the body. Keying the cost row to the batch rather than to the call also keeps accounting the same batch twice collapsing to one row instead of billing it twice. Every other call type still derives its key exactly as before. Cost and usage themselves are unaffected by redaction: the token columns fall back to the standard logging payload and spend comes from its response_cost, neither of which redaction touches. generate_hash_from_response had no other caller and is removed with it. --- .../spend_tracking/spend_tracking_utils.py | 43 ++--- .../test_spend_tracking_utils.py | 149 ++++++++++++++++++ 2 files changed, 163 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8d2569b2229..3146d8bccfb 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1,5 +1,3 @@ -import hashlib -import json import os import re import secrets @@ -28,6 +26,7 @@ from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.utils import ( + CallTypes, CostBreakdown, StandardLoggingGuardrailInformation, StandardLoggingMCPToolCall, @@ -144,36 +143,22 @@ def _get_spend_logs_metadata( return clean_metadata -def generate_hash_from_response(response_obj: Any) -> str: - """ - Generate a stable hash from a response object. - - Args: - response_obj: The response object to hash (can be dict, list, etc.) - - Returns: - A hex string representation of the MD5 hash - """ - try: - # Create a stable JSON string of the entire response object - # Sort keys to ensure consistent ordering - json_str: Final = json.dumps(response_obj, sort_keys=True) - - # Generate a hash of the response object - unique_hash: Final = hashlib.md5(json_str.encode()).hexdigest() - return unique_hash - except Exception: - # Return a fallback hash if serialization fails - return hashlib.md5(str(response_obj).encode()).hexdigest() +BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" def get_spend_logs_id(call_type: str, response_obj: dict, kwargs: dict) -> str | None: - if call_type == "aretrieve_batch" or call_type == "acreate_file": - # Generate a hash from the response object - id: str | None = generate_hash_from_response(response_obj) - else: - id = cast(str | None, response_obj.get("id")) or cast(str | None, kwargs.get("litellm_call_id")) - return id + standard_logging_payload = kwargs.get("standard_logging_object") + candidate_ids: Final = ( + response_obj.get("id"), + standard_logging_payload.get("id") if isinstance(standard_logging_payload, dict) else None, + kwargs.get("litellm_call_id"), + ) + resolved_id: Final = next( + (candidate for candidate in candidate_ids if isinstance(candidate, str) and candidate), None + ) + if resolved_id is not None and call_type == CallTypes.aretrieve_batch.value: + return f"{resolved_id}{BATCH_COST_REQUEST_ID_SUFFIX}" + return resolved_id def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> dict: diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9eb45c399db..6e43998a671 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -37,6 +37,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_request_body_for_spend_logs_payload, _should_store_prompts_and_responses_in_spend_logs, get_logging_payload, + get_spend_logs_id, ) from litellm.types.utils import ( StandardLoggingHiddenParams, @@ -2959,3 +2960,151 @@ def test_user_traffic_carries_no_internal_call_origin(): ) metadata = json.loads(payload["metadata"]) assert metadata["internal_call_origin"] is None + + +REDACTED_RESPONSE_PLACEHOLDER = {"text": "redacted-by-litellm"} +CONSTANT_ID_FROM_HASHED_PLACEHOLDER = "00fcbef15a3b0097e14b0ca016ed30a0" + + +@pytest.mark.parametrize("call_type", ["aretrieve_batch", "acreate_file"]) +def test_get_spend_logs_id_stays_unique_when_the_response_is_a_redaction_placeholder(call_type): + """request_id is the LiteLLM_SpendLogs primary key and the flush inserts with + skip_duplicates, so two calls must never derive the same id from identical response + content. Message redaction replaces every body it cannot redact with one fixed + placeholder, which is what a batch and a file body both become, so hashing the + response collapsed all of them onto a single id and silently dropped every row + after the first.""" + suffix = "_batch_cost" if call_type == "aretrieve_batch" else "" + first = get_spend_logs_id(call_type, dict(REDACTED_RESPONSE_PLACEHOLDER), {"litellm_call_id": "call-id-1"}) + second = get_spend_logs_id(call_type, dict(REDACTED_RESPONSE_PLACEHOLDER), {"litellm_call_id": "call-id-2"}) + + assert first == f"call-id-1{suffix}" + assert second == f"call-id-2{suffix}" + assert first != second + assert first != CONSTANT_ID_FROM_HASHED_PLACEHOLDER + assert second != CONSTANT_ID_FROM_HASHED_PLACEHOLDER + + +@pytest.mark.parametrize("call_type", ["aretrieve_batch", "acreate_file"]) +def test_get_spend_logs_id_prefers_the_response_id_for_batch_and_file_calls(call_type): + """A batch or file response that survives redaction carries its own id, so the row + keys off that rather than the per-call id.""" + expected = "batch_abc123_batch_cost" if call_type == "aretrieve_batch" else "batch_abc123" + assert get_spend_logs_id(call_type, {"id": "batch_abc123"}, {"litellm_call_id": "call-id-1"}) == expected + + +def test_get_logging_payload_gives_redacted_batch_and_file_rows_distinct_request_ids(): + """End to end at the payload level: a batch retrieve and a file create whose bodies + were both flattened to the same redaction placeholder must still produce two + insertable rows, each carrying its own spend.""" + payloads = [ + get_logging_payload( + kwargs={ + "call_type": call_type, + "model": model, + "litellm_call_id": call_id, + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=dict(REDACTED_RESPONSE_PLACEHOLDER), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + for call_type, model, call_id in ( + ("aretrieve_batch", "global.anthropic.claude-haiku-4-5-20251001-v1:0", "call-id-batch"), + ("acreate_file", "vertex_ai/gemini-2.5-flash", "call-id-file"), + ) + ] + request_ids = [payload["request_id"] for payload in payloads] + + assert request_ids == ["call-id-batch_batch_cost", "call-id-file"] + assert len(set(request_ids)) == len(request_ids) + assert CONSTANT_ID_FROM_HASHED_PLACEHOLDER not in request_ids + + +@pytest.mark.parametrize("call_type", ["aretrieve_batch", "acreate_file"]) +def test_get_spend_logs_id_keys_off_batch_identity_when_the_body_was_redacted(call_type): + """Retrieving one batch twice must produce one row, not two. Redaction strips the id + off the response body, so the identity has to come from the standard logging payload, + which is built from the unredacted response and keeps it. Falling through to the + per-call id here would write a second row carrying the same batch's full cost and + overstate spend by a multiple of how often the caller polled.""" + standard_logging_object = {"id": "batch_abc123"} + first = get_spend_logs_id( + call_type, + dict(REDACTED_RESPONSE_PLACEHOLDER), + {"litellm_call_id": "call-id-1", "standard_logging_object": standard_logging_object}, + ) + second = get_spend_logs_id( + call_type, + dict(REDACTED_RESPONSE_PLACEHOLDER), + {"litellm_call_id": "call-id-2", "standard_logging_object": standard_logging_object}, + ) + + expected = "batch_abc123_batch_cost" if call_type == "aretrieve_batch" else "batch_abc123" + assert first == second == expected + assert first != CONSTANT_ID_FROM_HASHED_PLACEHOLDER + + +def test_get_spend_logs_id_separates_distinct_batches_whose_bodies_were_both_redacted(): + """The flip side of idempotency: two different batches must not share a row just + because redaction flattened both bodies to the same placeholder.""" + ids = [ + get_spend_logs_id( + "aretrieve_batch", + dict(REDACTED_RESPONSE_PLACEHOLDER), + {"litellm_call_id": f"call-id-{index}", "standard_logging_object": {"id": batch_id}}, + ) + for index, batch_id in enumerate(("batch_first", "batch_second")) + ] + + assert ids == ["batch_first_batch_cost", "batch_second_batch_cost"] + + +def test_get_spend_logs_id_prefers_the_response_id_over_the_standard_logging_id(): + """An unredacted response keeps deciding its own row key, so cache-hit ids and every + other call type behave exactly as they did before.""" + assert ( + get_spend_logs_id( + "acompletion", + {"id": "chatcmpl-from-response"}, + {"litellm_call_id": "call-id-1", "standard_logging_object": {"id": "id-from-standard-payload"}}, + ) + == "chatcmpl-from-response" + ) + + +def test_batch_cost_row_does_not_collide_with_the_batch_creation_row(): + """Creating a batch writes a row keyed by the batch's own id, so keying the cost row + the same way makes the insert a duplicate of it. request_id is the primary key and the + flush skips duplicates, so the cost row is dropped with no error and the batch is + billed nothing. Observed against a live proxy: the poller computed and flushed the + cost, and the only row carrying that id was the acreate_batch row written when the + batch was submitted.""" + batch_id = "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDphYmM7bGxtX2JhdGNoX2lkOnh5eg" + + creation_row_id = get_spend_logs_id("acreate_batch", {"id": batch_id}, {"litellm_call_id": "call-create"}) + cost_row_id = get_spend_logs_id( + "aretrieve_batch", + dict(REDACTED_RESPONSE_PLACEHOLDER), + {"litellm_call_id": "call-poller", "standard_logging_object": {"id": batch_id}}, + ) + + assert creation_row_id == batch_id + assert cost_row_id != creation_row_id + assert cost_row_id == f"{batch_id}_batch_cost" + + +def test_batch_cost_row_id_is_stable_across_repeated_accounting(): + """The cost row stays keyed to the batch, so accounting the same batch twice collapses + to one row instead of billing it twice.""" + standard_logging_object = {"id": "batch_same"} + ids = [ + get_spend_logs_id( + "aretrieve_batch", + dict(REDACTED_RESPONSE_PLACEHOLDER), + {"litellm_call_id": f"call-{index}", "standard_logging_object": standard_logging_object}, + ) + for index in range(2) + ] + + assert ids[0] == ids[1] == "batch_same_batch_cost" From 363e3f3f03e959cb1242ced5f20dc051076fdd4e Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Thu, 13 Aug 2026 23:52:01 -0400 Subject: [PATCH 029/121] test(spend): annotate the batch cost row constants as Final --- basedpyright-code-budget.json | 4 ++-- ruff-strict-budget.json | 8 ++++---- .../proxy/spend_tracking/test_spend_tracking_utils.py | 6 +++--- type-discipline-budget.json | 4 ++-- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 521b4315e6e..d14a97f82f4 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 22947 + "limit": 22945 }, "reportArgumentType": { "limit": 2579 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 7312 + "limit": 7311 }, "reportFunctionMemberAccess": { "limit": 7 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 17c8f02dfdd..bd585bb2719 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1342 + "limit": 1341 }, "ASYNC230": { "limit": 11 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2924 + "limit": 2923 }, "C401": { "limit": 8 @@ -147,7 +147,7 @@ "limit": 3 }, "PLR1714": { - "limit": 257 + "limit": 256 }, "PLW0127": { "limit": 57 @@ -249,7 +249,7 @@ "limit": 113 }, "TRY300": { - "limit": 860 + "limit": 859 }, "UP028": { "limit": 2 diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 6e43998a671..0f6ac3f9b4f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -4,7 +4,7 @@ import json import os import sys from datetime import timezone -from typing import Any, cast +from typing import Any, Final, cast import pytest from fastapi.testclient import TestClient @@ -2962,8 +2962,8 @@ def test_user_traffic_carries_no_internal_call_origin(): assert metadata["internal_call_origin"] is None -REDACTED_RESPONSE_PLACEHOLDER = {"text": "redacted-by-litellm"} -CONSTANT_ID_FROM_HASHED_PLACEHOLDER = "00fcbef15a3b0097e14b0ca016ed30a0" +REDACTED_RESPONSE_PLACEHOLDER: Final = {"text": "redacted-by-litellm"} +CONSTANT_ID_FROM_HASHED_PLACEHOLDER: Final = "00fcbef15a3b0097e14b0ca016ed30a0" @pytest.mark.parametrize("call_type", ["aretrieve_batch", "acreate_file"]) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 894d99c92e0..fcd81c8e38e 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1074 + "limit": 1072 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16716 + "limit": 16715 }, "LIT011": { "limit": 5596 From c99a1ab0d7978a85724fb81c94ab66e704ded309 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Tue, 11 Aug 2026 22:10:17 -0400 Subject: [PATCH 030/121] fix(bedrock): resolve the managed-batch output bucket on the model-routed and cost-poller paths get_configured_s3_bucket_name accepts the output bucket only from the immutable _litellm_internal_model_credentials snapshot or AWS_S3_BUCKET_NAME. That refusal to read litellm_params is deliberate: the bucket is what validate_managed_cloud_file_id checks a file id against, so trusting a request-supplied value would let a caller redirect reads to a bucket of their choosing Two live entry points reach the Bedrock file-content transformation without ever building that snapshot. The managed-files pre-call hook sets data["model"] for any id carrying llm_output_file_id, which is every batch output, so get_file_content always takes the model-routed branch; that branch called llm_router.afile_content directly, and managed_files_obj.afile_content, the only caller that built the snapshot, is therefore unreachable for batch output. CheckBatchCost spread the deployment credentials as plain kwargs, and get_litellm_params does not carry s3_bucket_name across (gcs_bucket_name is listed for exactly this reason, its S3 counterpart is not), so the poller lost the bucket the same way The result was that every completed Bedrock managed batch failed files.content with "S3 bucket_name is required" and never had its cost tracked, leaving the row to be re-polled every cycle. Both paths now resolve the deployment credentials and pass the same MappingProxyType snapshot the managed-files hook already builds --- .../proxy/common_utils/check_batch_cost.py | 2 + .../openai_files_endpoints/files_endpoints.py | 8 ++ .../proxy_unit_tests/test_check_batch_cost.py | 96 +++++++++++++++++++ .../test_files_endpoint.py | 90 +++++++++++++++++ 4 files changed, 196 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 6fe37f0aacb..cfe60a79eed 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,6 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple from litellm._logging import verbose_proxy_logger @@ -537,6 +538,7 @@ class CheckBatchCost: credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} _file_content = await afile_content( file_id=raw_output_file_id, + _litellm_internal_model_credentials=MappingProxyType(dict(credentials)), **credentials, ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2432ea3729..1cbed2a68f9 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,6 +7,7 @@ import asyncio import traceback +from types import MappingProxyType from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -706,11 +707,18 @@ async def get_file_content( model: Final = cast(str | None, data.get("model")) if model: + deployment_credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model) + trusted_model_credentials: Final = ( + {"_litellm_internal_model_credentials": MappingProxyType(dict(deployment_credentials))} + if deployment_credentials is not None + else {} + ) response = await llm_router.afile_content( **{ "model": model, "file_id": file_id, **data, + **trusted_model_credentials, } ) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index fa274324fd6..20390159665 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -312,6 +312,102 @@ class TestCheckBatchCost: ), "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" + @pytest.mark.asyncio + async def test_output_fetch_passes_deployment_credentials_as_trusted_snapshot( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Bedrock resolves the output bucket ONLY from the immutable snapshot kwarg. + + Spreading the credentials as plain kwargs is not enough: get_litellm_params drops + s3_bucket_name, so without _litellm_internal_model_credentials the cost poller + cannot read the output file and every completed Bedrock batch stays unbilled. + """ + from types import MappingProxyType + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-bedrock-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "custom_llm_provider": "bedrock", + "s3_bucket_name": "configured-batch-bucket", + "aws_region_name": "us-east-1", + } + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "bedrock" + mock_deployment.litellm_params.model = "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"recordId":"req-1"}' + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"recordId": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("anthropic.claude-haiku-4-5-20251001-v1:0", "bedrock", None, None), + ), + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + + mock_afile_content.assert_awaited() + passed_kwargs = mock_afile_content.await_args[1] + snapshot = passed_kwargs.get("_litellm_internal_model_credentials") + assert snapshot is not None, "cost poller must pass the trusted credential snapshot" + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router 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 f27c8dfd2f4..0c26e5f7695 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 @@ -3149,6 +3149,96 @@ def test_require_managed_files_rejects_raw_provider_file_id( mock_call.assert_not_called() +def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkeypatch): + """A managed batch output id routes by model, and that branch must build the snapshot. + + The managed-files pre-call hook sets data["model"] for any id carrying + llm_output_file_id, so batch output retrieval always takes the model-routed branch + and never reaches managed_files_obj.afile_content. Bedrock resolves its output + bucket only from _litellm_internal_model_credentials, so without the snapshot every + Bedrock batch output retrieval fails with "S3 bucket_name is required". + """ + import base64 + from types import MappingProxyType + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.types.utils import SpecialEnums + + router = Router( + model_list=[ + { + "model_name": "anthropic.batch.claude-4.5-haiku", + "litellm_params": { + "model": "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "us-east-1", + "s3_bucket_name": "configured-batch-bucket", + }, + "model_info": {"id": "bedrock-batch-deployment-id"}, + } + ] + ) + + from unittest.mock import MagicMock + + managed_file_row = MagicMock() + managed_file_row.created_by = "test-user" + managed_file_row.team_id = None + managed_file_row.storage_backend = None + managed_file_row.storage_url = None + prisma_stub = MagicMock() + prisma_stub.db.litellm_managedfiletable.find_first = AsyncMock(return_value=managed_file_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_stub) + setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + + captured_kwargs: dict = {} + + async def _mock_router_afile_content(**kwargs): + captured_kwargs.update(kwargs) + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=b'{"recordId":"req-1"}', + headers={"content-type": "application/octet-stream"}, + ) + ) + + monkeypatch.setattr(router, "afile_content", _mock_router_afile_content) + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/jsonl", + "unified-output-id", + "anthropic.batch.claude-4.5-haiku", + "llm_output_file_id,s3://configured-batch-bucket/out/batch.jsonl", + "bedrock-batch-deployment-id", + ) + encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + try: + response = client.get( + f"/v1/files/{encoded_id}/content", + headers={"Authorization": "Bearer test-key", "custom-llm-provider": "bedrock"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + snapshot = captured_kwargs.get("_litellm_internal_model_credentials") + assert snapshot is not None, "model-routed branch must attach the trusted credential snapshot" + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert snapshot["s3_bucket_name"] == "configured-batch-bucket" + + def _unified_managed_file_id() -> str: import base64 From 460f0d29a95c25091fd375cd8f0f76297525ab78 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 12 Aug 2026 03:44:12 -0400 Subject: [PATCH 031/121] test(files): capture routed retrieval calls immutably The mock merged every call into one shared dict, so a second routed retrieval would overwrite the first and the assertions would still pass. Keep one frozen snapshot per call and assert exactly one call, which also makes an unintended second retrieval a failure rather than something the merge hides --- .../proxy/openai_files_endpoint/test_files_endpoint.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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 0c26e5f7695..e363a266688 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 @@ -3194,10 +3194,12 @@ def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkey monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) - captured_kwargs: dict = {} + # One frozen snapshot per call rather than one dict merged across calls, so a second + # invocation is visible instead of silently overwriting the first. + calls: list[MappingProxyType] = [] async def _mock_router_afile_content(**kwargs): - captured_kwargs.update(kwargs) + calls.append(MappingProxyType(dict(kwargs))) return HttpxBinaryResponseContent( response=httpx.Response( status_code=200, @@ -3231,7 +3233,8 @@ def test_get_file_content_model_routed_attaches_trusted_model_credentials(monkey app.dependency_overrides.pop(ps.user_api_key_auth, None) assert response.status_code == 200, response.text - snapshot = captured_kwargs.get("_litellm_internal_model_credentials") + assert len(calls) == 1, f"expected exactly one routed retrieval, got {len(calls)}" + snapshot = calls[0].get("_litellm_internal_model_credentials") assert snapshot is not None, "model-routed branch must attach the trusted credential snapshot" assert isinstance( snapshot, MappingProxyType From 60fe4e464cc56847735d9c3d3889717f51bee371 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 00:58:31 -0400 Subject: [PATCH 032/121] fix(bedrock): resolve the managed-batch output bucket on the inline accounting path too A third path reads a completed batch's output file, and it could not resolve the bucket either. When cost is accounted from the retrieve itself rather than from the poller, the batch success handler calls _handle_completed_batch, which fetches the output file through _extract_file_access_credentials. That helper forwarded a whitelist covering Azure and Vertex, gcs_bucket_name included, but nothing for Bedrock, and retrieve_batch built its litellm_params through get_litellm_params, whose fixed signature drops the trusted credential snapshot. So the snapshot never reached the file read and it failed with "S3 bucket_name is required" for a bucket the deployment had configured, leaving the batch's cost unrecorded. Adding s3_bucket_name to that whitelist would not have worked. The Bedrock file config deliberately resolves the bucket only from the immutable server-side snapshot or the environment, never from a request param, because the bucket is what managed file ids are validated against. The snapshot is therefore what has to flow, exactly as it already does for the model-routed and cost-poller paths. retrieve_batch now re-adds the snapshot after get_litellm_params, the same way the file operations already do, the whitelist forwards it, and the proxy attaches it for router-routed managed batches from the deployment behind the unified id. Verified against a live proxy reading a real completed Bedrock batch: the cost row appears within seconds of the retrieve carrying the batch's real spend and usage, where before the read raised and no row was written. Resolving those credentials is best effort. A batch whose deployment no longer resolves, which happens when a model group is removed while batches are in flight, still serves its status instead of failing the request on the lookup. This matters for the OSS and polling-disabled configurations, where the retrieve path is the only thing that accounts for a batch at all. --- litellm/batches/batch_utils.py | 1 + litellm/batches/main.py | 2 + litellm/proxy/batches_endpoints/endpoints.py | 8 +++ .../openai_files_endpoints/common_utils.py | 25 +++++++ .../test_litellm/batches/test_batch_utils.py | 14 ++++ tests/test_litellm/batches/test_main.py | 31 +++++++++ .../proxy/batches_endpoints/test_endpoints.py | 6 +- .../test_files_common_utils.py | 68 +++++++++++++++++++ 8 files changed, 154 insertions(+), 1 deletion(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index e73b887ae0a..f7aa6c50de8 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -309,6 +309,7 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict: "bucket_name", "timeout", "max_retries", + "_litellm_internal_model_credentials", ] for key in credential_keys: if key in litellm_params: diff --git a/litellm/batches/main.py b/litellm/batches/main.py index ce52c12818e..bb04d495555 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -22,6 +22,7 @@ from openai.types.batch import BatchRequestCounts import litellm from litellm._logging import verbose_logger +from litellm.files.main import _add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.llms.azure.batches.handler import AzureBatchesAPI @@ -527,6 +528,7 @@ def retrieve_batch( custom_llm_provider=custom_llm_provider, **kwargs, ) + _add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) if litellm_logging_obj is not None: litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e442cefa360..1301b9327ec 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, apply_team_provider_credentials, decode_model_from_file_id, + add_internal_model_credentials_for_batch, encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, @@ -537,6 +538,13 @@ async def retrieve_batch( detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) + if unified_batch_id: + add_internal_model_credentials_for_batch( + data=data, + llm_router=llm_router, + model_id=get_model_id_from_unified_batch_id(unified_batch_id), + ) + response = await llm_router.aretrieve_batch(**data) response._hidden_params["unified_batch_id"] = unified_batch_id if unified_batch_id: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 56e986c89cf..32676bd1d9f 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -465,6 +465,31 @@ def apply_team_provider_credentials( prepare_data_with_credentials(data=data, credentials=credentials) +def add_internal_model_credentials_for_batch( + data: dict, + llm_router: "Router", + model_id: str | None, +) -> None: + """ + Attach the deployment's immutable server-side credential snapshot to a router-routed + batch call (in-place). + + Cost accounting for a completed batch reads the batch's output file, and the Bedrock + file config resolves its bucket only from this snapshot, never from a request param, + because the bucket is what managed file ids are validated against. Without it that + read fails and the batch's cost is never recorded. + """ + if model_id is None: + return + try: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + except Exception: # noqa: BLE001 # the snapshot only enables cost accounting; a batch whose deployment no longer resolves must still be retrievable + return + if credentials is None: + return + data["_litellm_internal_model_credentials"] = MappingProxyType(dict(credentials)) + + def prepare_data_with_credentials( data: dict, credentials: dict, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 523b512e4cf..cacae3624f3 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable. import json import os import sys +from types import MappingProxyType import httpx import pytest @@ -1229,3 +1230,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200) assert models == ["claude-sonnet-4-5"] + + +def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): + """Bedrock resolves a batch's output bucket only from the immutable server-side + snapshot, never from a request param, so cost accounting on the retrieve path cannot + read the output file unless this key is forwarded. Without it the accounting raises + "S3 bucket_name is required" for a bucket the deployment has configured, and the + batch's cost is never recorded.""" + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"}) + + credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot}) + + assert credentials["_litellm_internal_model_credentials"] is snapshot diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index 1f7a91a5511..17e9ee29d4d 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -28,6 +28,7 @@ import sys from contextlib import ExitStack from dataclasses import dataclass from typing import Any, Dict +from types import MappingProxyType from unittest.mock import MagicMock, patch import pytest @@ -742,3 +743,33 @@ def test_resolve_timeout__httpx_timeout_returns_float_read(): resolved = bm._resolve_timeout(_params(timeout=t), {}, "openai") assert isinstance(resolved, float) assert resolved == 99.0 + + +def test_retrieve__forwards_trusted_model_credentials_into_litellm_params(seams): + """The batch's cost is computed by reading its output file after the retrieve, and + Bedrock resolves that bucket only from this immutable snapshot. get_litellm_params has + a fixed signature that drops it, so without re-adding it here the snapshot never + reaches the logging object and cost accounting fails on a bucket that is configured.""" + snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket"}) + logging_obj = MagicMock() + + bm.retrieve_batch( + batch_id="batch-1", + custom_llm_provider="openai", + litellm_logging_obj=logging_obj, + _litellm_internal_model_credentials=snapshot, + ) + + litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] + assert litellm_params["_litellm_internal_model_credentials"] is snapshot + + +def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): + """A retrieve with no snapshot must not invent an empty one, which would read as a + configured bucket of nothing.""" + logging_obj = MagicMock() + + bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="openai", litellm_logging_obj=logging_obj) + + litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] + assert "_litellm_internal_model_credentials" not in litellm_params diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a80c19f0708..aa5c63280b8 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1138,7 +1138,11 @@ async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness): # DISPATCH - router fired, direct litellm did not. assert retrieve_harness.router_aretrieve.call_count == 1 retrieve_harness.litellm_aretrieve.assert_not_called() - retrieve_harness.creds_resolver.assert_not_called() + + # Credentials are resolved for the deployment behind the unified id so the batch's + # output file can be read for cost accounting. This id resolves to nothing here, and + # the retrieve must still serve the batch rather than fail on the lookup. + retrieve_harness.creds_resolver.assert_called_once_with(model_id="gpt-4o-mini") # router receives the (still-encoded) batch id verbatim - this layer does # not decode it for the unified path. diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 4a021627c3e..a39f0c5f010 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -95,3 +95,71 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): "unified-out", "unified-err", ) + + +# =========================================================================== # +# add_internal_model_credentials_for_batch - the snapshot that lets a completed +# batch's output file be read, and therefore its cost be recorded +# =========================================================================== # + + +def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): + """Cost accounting for a completed batch reads its output file, and Bedrock resolves + that bucket only from this snapshot. It must be immutable so nothing downstream can + redirect the bucket that managed file ids are validated against.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials_for_batch, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"} + ) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-1") + + snapshot = data["_litellm_internal_model_credentials"] + assert snapshot["s3_bucket_name"] == "configured-bucket" + assert isinstance(snapshot, MappingProxyType) + with pytest.raises(TypeError): + snapshot["s3_bucket_name"] = "attacker-bucket" + router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="deployment-1") + + +@pytest.mark.parametrize( + "model_id, credentials", + [(None, {"s3_bucket_name": "b"}), ("deployment-1", None)], + ids=["no-model-id", "deployment-has-no-credentials"], +) +def test_add_internal_model_credentials_is_a_noop_without_a_resolvable_deployment(model_id, credentials): + """An unroutable batch must be left alone rather than given an empty snapshot, which + would look like a configured bucket of nothing.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials_for_batch, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock(return_value=credentials) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id=model_id) + + assert "_litellm_internal_model_credentials" not in data + + +def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): + """The snapshot only enables cost accounting, so a batch whose deployment no longer + resolves, which happens when a model group is removed while batches are in flight, + must still be retrievable rather than failing the request on the lookup.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + add_internal_model_credentials_for_batch, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock(side_effect=KeyError("deployment-gone")) + data = {"batch_id": "unified-batch-id"} + + add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-gone") + + assert data == {"batch_id": "unified-batch-id"} From d7afc1797cf2c0e2c326bf4f2b379991b92a9498 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 01:32:59 -0400 Subject: [PATCH 033/121] refactor(batches): share the trusted-credentials helper across both call paths The helper that carries the credential snapshot into litellm_params lived private in files/main.py, and the batch retrieve needed it too. It now sits beside get_litellm_params, which is what it augments, so neither caller reaches into the other's private surface. Typed as Mapping/MutableMapping of object rather than Any, which the strict import rules ban. The file-content route builds the snapshot through the same helper as the batch route instead of assembling a conditional mapping inline, which drops two mutable constructions and leaves one way to attach it. Its name loses the batch suffix now that both routes use it. --- litellm/batches/main.py | 4 ++-- litellm/files/main.py | 16 ++++------------ .../litellm_core_utils/get_litellm_params.py | 18 ++++++++++++++++++ litellm/proxy/batches_endpoints/endpoints.py | 4 ++-- .../openai_files_endpoints/common_utils.py | 2 +- .../openai_files_endpoints/files_endpoints.py | 10 ++-------- .../test_files_common_utils.py | 14 +++++++------- 7 files changed, 36 insertions(+), 32 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index bb04d495555..20d38bbb77f 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -22,7 +22,7 @@ from openai.types.batch import BatchRequestCounts import litellm from litellm._logging import verbose_logger -from litellm.files.main import _add_trusted_model_credentials_to_litellm_params +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.llms.azure.batches.handler import AzureBatchesAPI @@ -528,7 +528,7 @@ def retrieve_batch( custom_llm_provider=custom_llm_provider, **kwargs, ) - _add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) + add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs) if litellm_logging_obj is not None: litellm_logging_obj.update_from_kwargs( kwargs=kwargs, diff --git a/litellm/files/main.py b/litellm/files/main.py index 34421d13761..9a64c78552b 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -11,7 +11,6 @@ import time import uuid as uuid_module from collections.abc import Coroutine from functools import partial -from types import MappingProxyType from typing import Any, Final, Literal, cast import httpx @@ -34,6 +33,7 @@ import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse from litellm.files.types import FileContentProvider, FileContentStreamingResult +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.common_utils import get_azure_credentials @@ -85,14 +85,6 @@ bedrock_files_instance: Final = BedrockFilesHandler() ################################################# -def _add_trusted_model_credentials_to_litellm_params( - litellm_params_dict: dict[str, Any], kwargs: dict[str, Any] -) -> None: - trusted_model_credentials: Final = kwargs.get("_litellm_internal_model_credentials") - if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials - - @client async def acreate_file( file: FileTypes, @@ -372,7 +364,7 @@ def file_retrieve( ) if provider_config is not None: litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) @@ -494,7 +486,7 @@ def file_delete( pass optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) @@ -834,7 +826,7 @@ def file_content( try: optional_params: Final = GenericLiteLLMParams(**kwargs) litellm_params_dict: Final = get_litellm_params(**kwargs) - _add_trusted_model_credentials_to_litellm_params( + add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index f251ab4d74a..3eb8c163d5c 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping, MutableMapping +from types import MappingProxyType from typing import Final from litellm.llms.openai.data_residency import infer_openai_data_residency @@ -184,3 +186,19 @@ def get_litellm_params( litellm_params[key] = kwargs[key] return litellm_params + + +def add_trusted_model_credentials_to_litellm_params( + litellm_params_dict: MutableMapping[str, object], kwargs: Mapping[str, object] +) -> None: + """ + Carry the immutable server-side credential snapshot into litellm_params. + + get_litellm_params has a fixed signature, so callers that need the snapshot to + survive into the logging object and the downstream file read have to re-add it. Only + a MappingProxyType is accepted, since providers resolve trusted configuration such + as a Bedrock file bucket from it and must not read a request-supplied mapping. + """ + trusted_model_credentials: Final = kwargs.get("_litellm_internal_model_credentials") + if isinstance(trusted_model_credentials, MappingProxyType): + litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1301b9327ec..9a6bb054d1f 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,7 +25,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, apply_team_provider_credentials, decode_model_from_file_id, - add_internal_model_credentials_for_batch, + add_internal_model_credentials, encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, @@ -539,7 +539,7 @@ async def retrieve_batch( ) if unified_batch_id: - add_internal_model_credentials_for_batch( + add_internal_model_credentials( data=data, llm_router=llm_router, model_id=get_model_id_from_unified_batch_id(unified_batch_id), diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 32676bd1d9f..f2e6fb633e1 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -465,7 +465,7 @@ def apply_team_provider_credentials( prepare_data_with_credentials(data=data, credentials=credentials) -def add_internal_model_credentials_for_batch( +def add_internal_model_credentials( data: dict, llm_router: "Router", model_id: str | None, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 1cbed2a68f9..361b5b920e2 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,7 +7,6 @@ import asyncio import traceback -from types import MappingProxyType from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -44,6 +43,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + add_internal_model_credentials, apply_team_provider_credentials, encode_file_id_with_model, extract_file_creation_params, @@ -707,18 +707,12 @@ async def get_file_content( model: Final = cast(str | None, data.get("model")) if model: - deployment_credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model) - trusted_model_credentials: Final = ( - {"_litellm_internal_model_credentials": MappingProxyType(dict(deployment_credentials))} - if deployment_credentials is not None - else {} - ) + add_internal_model_credentials(data=data, llm_router=llm_router, model_id=model) response = await llm_router.afile_content( **{ "model": model, "file_id": file_id, **data, - **trusted_model_credentials, } ) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index a39f0c5f010..ad7f5e4725a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -98,7 +98,7 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): # =========================================================================== # -# add_internal_model_credentials_for_batch - the snapshot that lets a completed +# add_internal_model_credentials - the snapshot that lets a completed # batch's output file be read, and therefore its cost be recorded # =========================================================================== # @@ -108,7 +108,7 @@ def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): that bucket only from this snapshot. It must be immutable so nothing downstream can redirect the bucket that managed file ids are validated against.""" from litellm.proxy.openai_files_endpoints.common_utils import ( - add_internal_model_credentials_for_batch, + add_internal_model_credentials, ) router = MagicMock() @@ -117,7 +117,7 @@ def test_add_internal_model_credentials_attaches_an_immutable_snapshot(): ) data = {"batch_id": "unified-batch-id"} - add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-1") + add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-1") snapshot = data["_litellm_internal_model_credentials"] assert snapshot["s3_bucket_name"] == "configured-bucket" @@ -136,14 +136,14 @@ def test_add_internal_model_credentials_is_a_noop_without_a_resolvable_deploymen """An unroutable batch must be left alone rather than given an empty snapshot, which would look like a configured bucket of nothing.""" from litellm.proxy.openai_files_endpoints.common_utils import ( - add_internal_model_credentials_for_batch, + add_internal_model_credentials, ) router = MagicMock() router.get_deployment_credentials_with_provider = MagicMock(return_value=credentials) data = {"batch_id": "unified-batch-id"} - add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id=model_id) + add_internal_model_credentials(data=data, llm_router=router, model_id=model_id) assert "_litellm_internal_model_credentials" not in data @@ -153,13 +153,13 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): resolves, which happens when a model group is removed while batches are in flight, must still be retrievable rather than failing the request on the lookup.""" from litellm.proxy.openai_files_endpoints.common_utils import ( - add_internal_model_credentials_for_batch, + add_internal_model_credentials, ) router = MagicMock() router.get_deployment_credentials_with_provider = MagicMock(side_effect=KeyError("deployment-gone")) data = {"batch_id": "unified-batch-id"} - add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-gone") + add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-gone") assert data == {"batch_id": "unified-batch-id"} From 5649098e1b47a3ab6a341971ac7c79b987bda35c Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Thu, 13 Aug 2026 11:01:35 -0400 Subject: [PATCH 034/121] fix(batches): account a managed batch's cost exactly once Two components computed a managed batch's cost and each assumed it was the only one. Retrieving a batch computed it through the @client decorator's success callback, and CheckBatchCost computed it on its own schedule. Whichever observed completion first decided the outcome, so cost was either counted once per retrieve or not at all. The lockout is the worse half. Retrieving a batch that had reached completion set batch_processed=True, which is what takes a batch out of CheckBatchCost's queue, since it selects batch_processed=False. That write claimed the cost had been accounted for on behalf of a callback that had not run yet and was not awaited. When the callback then failed the cost was gone permanently, with the poller already retired and no retry left. Observed on a live proxy: two completed batches whose callbacks raised inside the logging worker, one on a provider output path that did not resolve and one on a batch whose output file id was still None, both left marked processed with no spend row and no way to recover them. Nothing logged at error level for the batches themselves. The over-count is the other half. Nothing suppressed recomputation, so each retrieve of an already-completed batch recorded that batch's full cost again. A caller polling its own batch to see whether it had finished inflated spend by however many times it looked. The flag now means what its name says, and only the component that actually recorded the cost sets it. When the poller is running it owns accounting, so retrieving a managed batch records no cost and leaves the flag alone; the poller computes once and sets it. When the poller cannot be relied on, either because polling is disabled by config or because the enterprise job never registered, the retrieve path is the only accountant and behaves exactly as before. Batches with no managed object row are untouched either way, since neither the flag nor the poller queue applies to them. --- litellm/proxy/batches_endpoints/endpoints.py | 7 ++ .../openai_files_endpoints/common_utils.py | 33 ++++-- .../proxy/batches_endpoints/test_endpoints.py | 44 +++++++ .../test_files_common_utils.py | 111 ++++++++++++++++++ 4 files changed, 186 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e442cefa360..d5bc4ac0116 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, apply_team_provider_credentials, + batch_cost_poller_is_active, decode_model_from_file_id, encode_batch_response_ids, encode_file_id_with_model, @@ -496,6 +497,12 @@ async def retrieve_batch( "Batch %s is in non-terminal state %s, syncing with provider", batch_id, response.status ) + if unified_batch_id and batch_cost_poller_is_active(): + data["litellm_metadata"] = { + **(data.get("litellm_metadata") or {}), + "batch_ignore_default_logging": True, + } + # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 56e986c89cf..b8c250e718b 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1231,6 +1231,29 @@ async def get_batch_from_database( return None, None +def batch_cost_poller_is_active() -> bool: + """ + Whether the CheckBatchCost poller is running and will therefore account for a + managed batch's cost itself. + + False whenever the poller cannot be relied on: polling disabled by config, or the + job absent from the scheduler because the enterprise import failed. + """ + from litellm.constants import PROXY_BATCH_POLLING_ENABLED + + if not PROXY_BATCH_POLLING_ENABLED: + return False + try: + import litellm.proxy.proxy_server as proxy_server_module + + scheduler = getattr(proxy_server_module, "scheduler", None) + if scheduler is None: + return False + return scheduler.get_job("check_batch_cost_job") is not None + except Exception: # noqa: BLE001 # scheduler backends raise varied types from get_job; an unreadable scheduler means the poller cannot be relied on + return False + + async def update_batch_in_database( batch_id: str, unified_batch_id: str | Literal[False], @@ -1304,15 +1327,7 @@ async def update_batch_in_database( "updated_at": litellm.utils.get_utc_datetime(), } - # When a batch reaches completion, also mark batch_processed=True. - # The cost callback is enqueued asynchronously during the - # aretrieve_batch call that detected completion (via the @client - # decorator). It is not awaited, so there is a theoretical window - # where the callback hasn't executed yet. In practice the callback - # completes reliably. Setting the flag here unblocks file deletion - # which queries batch_processed=False. CheckBatchCost acts as a - # safety net for the rare case where the callback fails. - if db_status == "complete": + if db_status == "complete" and not batch_cost_poller_is_active(): update_data["batch_processed"] = True try: diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a80c19f0708..6d00e56030f 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2406,3 +2406,47 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc await call_cancel(cancel_harness, _unified_batch_id()) assert cancel_harness.router_acancel.call_count == 1 + + +# =========================================================================== # +# Retrieve - who accounts for a managed batch's cost. Retrieving a batch and +# the CheckBatchCost poller both computed it, so whichever observed completion +# first won and the other either double counted or was locked out. +# =========================================================================== # + + +@pytest.mark.asyncio +async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness): + """With the poller running it is the single accountant, so the retrieve must not also + record cost. Without this the same batch is billed once per retrieve, and a caller + polling its own batch inflates spend by however many times it looked.""" + with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router.aretrieve_batch.await_count == 1 + metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {} + assert metadata.get("batch_ignore_default_logging") is True + + +@pytest.mark.asyncio +async def test_retrieve__managed_batch_still_accounts_inline_without_a_poller(retrieve_harness): + """No poller means nothing else will ever account for this batch, so the retrieve has + to keep doing it. Suppressing here unconditionally would lose batch cost entirely on + any proxy running with batch polling disabled.""" + with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=False)): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router.aretrieve_batch.await_count == 1 + metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {} + assert metadata.get("batch_ignore_default_logging") is None + + +@pytest.mark.asyncio +async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retrieve_harness): + """An unmanaged batch has no managed object row and so no poller queue entry. It must + keep accounting inline whatever the poller is doing.""" + with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): + await call_retrieve(retrieve_harness, "batch-raw-xyz") + + metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {} + assert metadata.get("batch_ignore_default_logging") is None diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 4a021627c3e..4268d9c5e3e 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -95,3 +95,114 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): "unified-out", "unified-err", ) + + +class _FakeScheduler: + def __init__(self, job): + self._job = job + + def get_job(self, job_id): + assert job_id == "check_batch_cost_job" + return self._job + + +@pytest.mark.parametrize( + "polling_enabled, job, expected", + [ + (True, object(), True), + (True, None, False), + (False, object(), False), + ], + ids=["poller-running", "job-absent-enterprise-import-failed", "polling-disabled-by-config"], +) +def test_batch_cost_poller_is_active(monkeypatch, polling_enabled, job, expected): + """The predicate must only claim the poller when it can actually be relied on, so a + proxy with polling switched off or without the enterprise job keeps accounting for + batch cost on the retrieve path.""" + import litellm.constants + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_cost_poller_is_active, + ) + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", polling_enabled, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", _FakeScheduler(job), raising=False) + + assert batch_cost_poller_is_active() is expected + + +def test_batch_cost_poller_is_active_is_false_when_no_scheduler_exists(monkeypatch): + import litellm.constants + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_cost_poller_is_active, + ) + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", None, raising=False) + + assert batch_cost_poller_is_active() is False + + +def _completed_batch() -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-done", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id="file-out", + ) + + +async def _run_update(monkeypatch, poller_active: bool) -> dict: + import litellm.proxy.openai_files_endpoints.common_utils as cu + + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: poller_active) + monkeypatch.setattr(cu, "ensure_batch_response_managed_file_ids", AsyncMock()) + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + + db_batch_object = MagicMock() + db_batch_object.status = "in_progress" + + await cu.update_batch_in_database( + batch_id="unified-batch-id", + unified_batch_id="unified-batch-id", + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + db_batch_object=db_batch_object, + operation="retrieve", + ) + + assert update_mock.await_count == 1 + return update_mock.await_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_retrieving_a_completed_batch_leaves_batch_processed_to_the_cost_poller(monkeypatch): + """batch_processed is what removes a batch from CheckBatchCost's queue, which selects + batch_processed=False. Retrieving a batch records no cost when the poller is active, so + setting the flag here retired the poller on behalf of work nobody had done: a cost + callback that then failed lost the batch's cost permanently with no retry left. The + status update must still happen so callers see the terminal state.""" + data = await _run_update(monkeypatch, poller_active=True) + + assert "batch_processed" not in data + assert data["status"] == "complete" + + +@pytest.mark.asyncio +async def test_retrieving_a_completed_batch_still_marks_processed_without_a_cost_poller(monkeypatch): + """With no poller to hand off to, this path is the only accountant, so it keeps setting + the flag. Otherwise a proxy with polling disabled would never unblock file deletion.""" + data = await _run_update(monkeypatch, poller_active=False) + + assert data["batch_processed"] is True + assert data["status"] == "complete" From ec52858865b8553fa2d7fad5cc2e701dc7ed8199 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 00:20:10 -0400 Subject: [PATCH 035/121] fix(batches): only hand accounting to the poller once it can mark batches done The handoff asked whether the poller was running, when what matters is whether it will actually account for the batch. Those differ on a schema without the batch_processed column: the poller cannot filter on it, so it falls back to a query that excludes complete and completed rows, and it cannot set it either. A caller retrieving a provider-completed batch before the poller saw it therefore suppressed inline accounting, then marked the row complete, and the fallback query could never find it again. Nobody accounted for that batch, so its cost escaped the caller's budget entirely. The poller now publishes batch_processed_support_confirmed, set only once a filtered query has actually succeeded, and the handoff requires it. Defaulting to unconfirmed keeps accounting on the retrieve path in exactly the cases the poller would drop the batch, including the window before the poller's first cycle. All four combinations account exactly once: unconfirmed leaves the retrieve accounting and setting the marker, whether or not the column exists, and confirmed is only reachable when the column is present, where the poller accounts and sets it. A scheduler that hands back something other than a bound method leaves no poller to interrogate, which reads as unconfirmed rather than as working. --- .../proxy/common_utils/check_batch_cost.py | 2 + litellm/proxy/batches_endpoints/endpoints.py | 9 +- .../openai_files_endpoints/common_utils.py | 19 ++- .../proxy_unit_tests/test_check_batch_cost.py | 6 + .../test_files_common_utils.py | 130 +++++++++++++++++- 5 files changed, 152 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 6fe37f0aacb..990964dc81f 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -51,6 +51,7 @@ class CheckBatchCost: # Cached after the first poll cycle. Once we know the column is absent we skip # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True + self.batch_processed_support_confirmed: bool = False async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: """ @@ -722,6 +723,7 @@ class CheckBatchCost: take=MAX_OBJECTS_PER_POLL_CYCLE, order={"created_at": "asc"}, ) + self.batch_processed_support_confirmed = True except Exception as query_err: if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower(): raise diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index d5bc4ac0116..563c380d34b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -498,10 +498,11 @@ async def retrieve_batch( ) if unified_batch_id and batch_cost_poller_is_active(): - data["litellm_metadata"] = { - **(data.get("litellm_metadata") or {}), - "batch_ignore_default_logging": True, - } + litellm_metadata = data.get("litellm_metadata") + if not isinstance(litellm_metadata, dict): + litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend + data["litellm_metadata"] = litellm_metadata + litellm_metadata["batch_ignore_default_logging"] = True # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b8c250e718b..012c0e6ea5b 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1233,11 +1233,16 @@ async def get_batch_from_database( def batch_cost_poller_is_active() -> bool: """ - Whether the CheckBatchCost poller is running and will therefore account for a - managed batch's cost itself. + Whether the CheckBatchCost poller will account for a managed batch's cost itself. - False whenever the poller cannot be relied on: polling disabled by config, or the - job absent from the scheduler because the enterprise import failed. + False whenever the poller cannot be relied on: polling disabled by config, the job + absent from the scheduler because the enterprise import failed, or the poller not + yet having confirmed that the batch_processed column exists. That last condition + matters because the poller needs the column both to find outstanding batches and to + mark them accounted; without it the poller falls back to a query that excludes + terminal statuses, so a batch the retrieve path has already marked complete becomes + invisible to it. Defaulting to False until the poller confirms support keeps the + retrieve path accounting in exactly the cases the poller would drop the batch. """ from litellm.constants import PROXY_BATCH_POLLING_ENABLED @@ -1249,7 +1254,11 @@ def batch_cost_poller_is_active() -> bool: scheduler = getattr(proxy_server_module, "scheduler", None) if scheduler is None: return False - return scheduler.get_job("check_batch_cost_job") is not None + job = scheduler.get_job("check_batch_cost_job") + if job is None: + return False + poller = getattr(getattr(job, "func", None), "__self__", None) + return getattr(poller, "batch_processed_support_confirmed", False) is True except Exception: # noqa: BLE001 # scheduler backends raise varied types from get_job; an unreadable scheduler means the poller cannot be relied on return False diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index fa274324fd6..ce03dd33f85 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -143,6 +143,11 @@ class TestCheckBatchCost: assert "complete" not in not_in assert "completed" not in not_in assert find_call[1]["where"]["batch_processed"] is False + # A successful filtered query is the only proof the column exists. The retrieve + # path reads this to decide whether handing accounting to the poller is safe: + # without the column the poller's fallback query excludes complete/completed, so + # a batch already marked complete would never be accounted by anyone. + assert check_batch_cost_instance.batch_processed_support_confirmed is True @pytest.mark.asyncio async def test_fallback_query_used_when_batch_processed_missing( @@ -171,6 +176,7 @@ class TestCheckBatchCost: assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE # Column absence is now cached — next call should go straight to fallback assert check_batch_cost_instance._has_batch_processed_column is False + assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio async def test_column_absence_cached_across_cycles( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 4268d9c5e3e..20858a2a15f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -106,19 +106,44 @@ class _FakeScheduler: return self._job +class _FakePoller: + def __init__(self, confirmed): + self.batch_processed_support_confirmed = confirmed + + def check_batch_cost(self): + return None + + +def _job_for(poller): + if poller is None: + return None + job = MagicMock() + job.func = poller.check_batch_cost + return job + + @pytest.mark.parametrize( "polling_enabled, job, expected", [ - (True, object(), True), + (True, _job_for(_FakePoller(confirmed=True)), True), + (True, _job_for(_FakePoller(confirmed=False)), False), (True, None, False), - (False, object(), False), + (False, _job_for(_FakePoller(confirmed=True)), False), + ], + ids=[ + "poller-running-and-column-confirmed", + "poller-running-but-column-unconfirmed", + "job-absent-enterprise-import-failed", + "polling-disabled-by-config", ], - ids=["poller-running", "job-absent-enterprise-import-failed", "polling-disabled-by-config"], ) def test_batch_cost_poller_is_active(monkeypatch, polling_enabled, job, expected): """The predicate must only claim the poller when it can actually be relied on, so a - proxy with polling switched off or without the enterprise job keeps accounting for - batch cost on the retrieve path.""" + proxy with polling switched off, without the enterprise job, or whose poller has not + confirmed batch_processed support keeps accounting for batch cost on the retrieve + path. The unconfirmed case is the one that matters for legacy schemas: without the + column the poller falls back to a query excluding terminal statuses, so a batch the + retrieve path already marked complete would never be accounted by anyone.""" import litellm.constants import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -206,3 +231,98 @@ async def test_retrieving_a_completed_batch_still_marks_processed_without_a_cost assert data["batch_processed"] is True assert data["status"] == "complete" + + +def test_batch_cost_poller_is_active_is_false_when_the_job_has_no_bound_poller(monkeypatch): + """A scheduler that hands back a plain function rather than a bound method leaves no + poller to interrogate, so the predicate stays conservative instead of assuming the + column is supported.""" + import litellm.constants + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_cost_poller_is_active, + ) + + def unbound_check_batch_cost(): + return None + + job = MagicMock() + job.func = unbound_check_batch_cost + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", _FakeScheduler(job), raising=False) + + assert batch_cost_poller_is_active() is False + + +def test_batch_cost_poller_is_active_is_false_when_get_job_raises(monkeypatch): + """Scheduler backends raise varied types; an unreadable scheduler must not be read as + a working poller.""" + import litellm.constants + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_cost_poller_is_active, + ) + + class _ExplodingScheduler: + def get_job(self, job_id): + raise RuntimeError("scheduler not started") + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", _ExplodingScheduler(), raising=False) + + assert batch_cost_poller_is_active() is False + + + +@pytest.mark.asyncio +async def test_retrieving_a_batch_whose_status_is_unchanged_writes_nothing(monkeypatch): + """A caller polling an already-complete batch must not write at all, so repeated polls + cannot flip batch_processed or disturb whichever component owns accounting.""" + import litellm.proxy.openai_files_endpoints.common_utils as cu + + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: False) + monkeypatch.setattr(cu, "ensure_batch_response_managed_file_ids", AsyncMock()) + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + + db_batch_object = MagicMock() + db_batch_object.status = "completed" + + await cu.update_batch_in_database( + batch_id="unified-batch-id", + unified_batch_id="unified-batch-id", + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + db_batch_object=db_batch_object, + operation="retrieve", + ) + + update_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_batch_in_database_is_a_noop_for_unmanaged_batches(monkeypatch): + """Batches with no managed object row have neither the flag nor a poller queue entry, so + this path must leave them alone entirely.""" + import litellm.proxy.openai_files_endpoints.common_utils as cu + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + + await cu.update_batch_in_database( + batch_id="batch-raw-xyz", + unified_batch_id=False, + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + operation="retrieve", + ) + + update_mock.assert_not_awaited() From c9e9c279fe6270132079db6da41b69c7d8809169 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 02:18:40 -0400 Subject: [PATCH 036/121] fix(batches): decide batch cost ownership once per retrieve The ownership question was asked twice for one retrieve: once before the provider call to decide whether to suppress inline accounting, and again afterwards to decide whether to mark the batch accounted. Between those two points the poller can complete its first successful filtered query and become usable, so the two answers disagree. The retrieve then accounts for the batch inline, having decided the poller was unusable, while the later check sees a usable poller and leaves the marker unset, so the poller accounts for the same batch again and its spend is counted twice. The retrieve now decides once and passes that decision to update_batch_in_database, which prefers it over re-deriving one. Callers that record no cost of their own leave it unset and keep deriving it as before, so the cancel path is unchanged. --- litellm/proxy/batches_endpoints/endpoints.py | 4 +- .../openai_files_endpoints/common_utils.py | 12 +++- .../test_files_common_utils.py | 69 +++++++++++++++++++ 3 files changed, 83 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 563c380d34b..554a286555b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -497,7 +497,8 @@ async def retrieve_batch( "Batch %s is in non-terminal state %s, syncing with provider", batch_id, response.status ) - if unified_batch_id and batch_cost_poller_is_active(): + poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active() + if poller_owns_accounting: litellm_metadata = data.get("litellm_metadata") if not isinstance(litellm_metadata, dict): litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend @@ -581,6 +582,7 @@ async def retrieve_batch( verbose_proxy_logger=verbose_proxy_logger, db_batch_object=db_batch_object, operation="retrieve", + poller_owns_accounting=poller_owns_accounting, ) ### CALL HOOKS ### - modify outgoing data diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 012c0e6ea5b..c24d7b6e9de 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1273,6 +1273,7 @@ async def update_batch_in_database( db_batch_object=None, operation: str = "update", user_api_key_dict=None, + poller_owns_accounting: bool | None = None, ): """ Update batch status and object in ManagedObjectTable. @@ -1287,6 +1288,12 @@ async def update_batch_in_database( db_batch_object: Optional existing database object; fetched by unified_object_id when omitted operation: Description of operation ("update", "cancel", etc.) user_api_key_dict: Optional auth context for creating managed file IDs + poller_owns_accounting: Whether the caller already decided that the cost poller + owns this batch's accounting. Callers that suppress their own inline + accounting must pass the same decision they acted on, because re-deciding + here can observe a poller that became usable in between and leave the batch + unmarked after it was already accounted for, billing it twice. Left None by + callers that record no cost themselves. """ import litellm.utils @@ -1336,7 +1343,10 @@ async def update_batch_in_database( "updated_at": litellm.utils.get_utc_datetime(), } - if db_status == "complete" and not batch_cost_poller_is_active(): + poller_owns: Final = ( + batch_cost_poller_is_active() if poller_owns_accounting is None else poller_owns_accounting + ) + if db_status == "complete" and not poller_owns: update_data["batch_processed"] = True try: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 20858a2a15f..c16450446a5 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -326,3 +326,72 @@ async def test_update_batch_in_database_is_a_noop_for_unmanaged_batches(monkeypa ) update_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_the_caller_s_accounting_decision_wins_over_a_later_poller_transition(monkeypatch): + """The ownership decision is made before the provider retrieval and acted on there, so + re-deciding afterwards can observe a poller that only just became usable. That split + left the retrieve accounting inline while the row stayed unmarked, so the poller + accounted for the same batch again and billed it twice. Passing the decision through + makes both halves agree even when the poller transitions mid-flight.""" + import litellm.proxy.openai_files_endpoints.common_utils as cu + + # The predicate now reports an active poller, i.e. it flipped during the retrieval. + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: True) + monkeypatch.setattr(cu, "ensure_batch_response_managed_file_ids", AsyncMock()) + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + db_batch_object = MagicMock() + db_batch_object.status = "in_progress" + + await cu.update_batch_in_database( + batch_id="unified-batch-id", + unified_batch_id="unified-batch-id", + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + db_batch_object=db_batch_object, + operation="retrieve", + poller_owns_accounting=False, + ) + + data = update_mock.await_args.kwargs["data"] + assert data["batch_processed"] is True + assert data["status"] == "complete" + + +@pytest.mark.asyncio +async def test_a_caller_that_handed_off_accounting_still_leaves_the_marker_alone(monkeypatch): + """The mirror case: a caller that suppressed its own accounting must leave the marker + for the poller even if the predicate has since stopped reporting one, otherwise the + batch is retired without anyone having accounted for it.""" + import litellm.proxy.openai_files_endpoints.common_utils as cu + + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: False) + monkeypatch.setattr(cu, "ensure_batch_response_managed_file_ids", AsyncMock()) + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + db_batch_object = MagicMock() + db_batch_object.status = "in_progress" + + await cu.update_batch_in_database( + batch_id="unified-batch-id", + unified_batch_id="unified-batch-id", + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + db_batch_object=db_batch_object, + operation="retrieve", + poller_owns_accounting=True, + ) + + data = update_mock.await_args.kwargs["data"] + assert "batch_processed" not in data + assert data["status"] == "complete" From b066ed3e31f52bc7e47c076002798e2b839b15ef Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 13:10:27 +0000 Subject: [PATCH 037/121] fix(model_prices): correct Gemini 2.5 shutdown dates and DeepSeek V4 max output tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 21 +++++++++++-------- model_prices_and_context_window.json | 21 +++++++++++-------- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b288269b0a2..33a23d3b9e1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19803,7 +19803,7 @@ "uses_embed_content": true }, "gemini/gemini-embedding-001": { - "deprecation_date": "2028-05-14", + "deprecation_date": "2026-07-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "gemini", "max_input_tokens": 2048, @@ -19979,6 +19979,7 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -20290,6 +20291,7 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_token_cost": 1e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -20586,6 +20588,7 @@ "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "input_cost_per_token_priority": 1.25e-06, @@ -47289,8 +47292,8 @@ "input_cost_per_token_cache_hit": 2.8e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -47315,8 +47318,8 @@ "input_cost_per_token_cache_hit": 3.625e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 8.7e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -47341,8 +47344,8 @@ "input_cost_per_token_cache_hit": 2.8e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -47367,8 +47370,8 @@ "input_cost_per_token_cache_hit": 3.625e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 8.7e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b288269b0a2..33a23d3b9e1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19803,7 +19803,7 @@ "uses_embed_content": true }, "gemini/gemini-embedding-001": { - "deprecation_date": "2028-05-14", + "deprecation_date": "2026-07-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "gemini", "max_input_tokens": 2048, @@ -19979,6 +19979,7 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -20290,6 +20291,7 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_token_cost": 1e-08, + "deprecation_date": "2026-10-16", "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -20586,6 +20588,7 @@ "gemini/gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.25e-06, "input_cost_per_token_above_200k_tokens": 2.5e-06, "input_cost_per_token_priority": 1.25e-06, @@ -47289,8 +47292,8 @@ "input_cost_per_token_cache_hit": 2.8e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -47315,8 +47318,8 @@ "input_cost_per_token_cache_hit": 3.625e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 8.7e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -47341,8 +47344,8 @@ "input_cost_per_token_cache_hit": 2.8e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", @@ -47367,8 +47370,8 @@ "input_cost_per_token_cache_hit": 3.625e-09, "litellm_provider": "deepseek", "max_input_tokens": 1000000, - "max_output_tokens": 8192, - "max_tokens": 8192, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "output_cost_per_token": 8.7e-07, "source": "https://api-docs.deepseek.com/quick_start/pricing", From ce66cbce0e1edc3bdbf040024017edf3beb861ac Mon Sep 17 00:00:00 2001 From: pokepoke81 <4258646+pokepoke81@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:45:21 -0400 Subject: [PATCH 038/121] fix(databricks): surface prompt-cache token counts in streaming usage chunk_parser built ModelResponseStream without passing usage, so the cache_read_input_tokens and cache_creation_input_tokens that Databricks returns for Anthropic models never reached the cost calculator. Every streamed request was billed at the full input rate even when served from cache. ModelResponseStream already coerces a usage dict into Usage, which maps those keys into prompt_tokens_details, so passing the chunk's usage through is sufficient. --- .../llms/databricks/chat/transformation.py | 1 + .../test_databricks_chat_transformation.py | 76 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 8b44ab4feaf..8a625569cfa 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -733,6 +733,7 @@ class DatabricksChatResponseIterator(BaseModelResponseIterator): created=chunk["created"], model=chunk["model"], choices=translated_choices, + usage=chunk.get("usage"), ) except KeyError as e: raise DatabricksException( diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 00f3e7a6faf..d6b8e1a3652 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -423,3 +423,79 @@ def test_databricks_config_probes_capabilities_under_databricks_namespace(): without this override they probed the ``anthropic`` cost-map namespace and ignored the exact ``databricks/databricks-claude-*`` entries.""" assert DatabricksConfig().custom_llm_provider == "databricks" + + +def _streaming_chunk(usage=None, choices=None): + base = { + "id": "chatcmpl-test", + "created": 1234567890, + "model": "databricks-claude-sonnet-5", + "choices": [{"delta": {"content": "hi"}}] if choices is None else choices, + } + return base if usage is None else {**base, "usage": usage} + + +@pytest.mark.parametrize( + "cache_read, cache_creation, expected_cached, expected_written", + [ + (12002, 0, 12002, 0), + (0, 12002, 0, 12002), + ], + ids=["warm_cache_read", "cold_cache_write"], +) +def test_chunk_parser_surfaces_prompt_cache_usage(cache_read, cache_creation, expected_cached, expected_written): + """Databricks returns Anthropic prompt-cache counts in the streaming usage object, + but chunk_parser dropped usage entirely, so cache-aware pricing never reached the + cost calculator and every streamed request was billed at the full input rate.""" + iterator = DatabricksChatResponseIterator(streaming_response=None, sync_stream=True) + + result = iterator.chunk_parser( + _streaming_chunk( + usage={ + "prompt_tokens": 12011, + "completion_tokens": 8, + "total_tokens": 12019, + "cache_read_input_tokens": cache_read, + "cache_creation_input_tokens": cache_creation, + } + ) + ) + + assert result.usage is not None + assert result.usage.prompt_tokens == 12011 + assert result.usage.completion_tokens == 8 + assert result.usage.prompt_tokens_details is not None + assert result.usage.prompt_tokens_details.cached_tokens == expected_cached + assert result.usage._cache_creation_input_tokens == expected_written + + +def test_chunk_parser_surfaces_usage_only_final_chunk(): + """stream_options={"include_usage": True} emits a trailing chunk whose choices + list is empty; usage must still reach the caller.""" + iterator = DatabricksChatResponseIterator(streaming_response=None, sync_stream=True) + + result = iterator.chunk_parser( + _streaming_chunk( + usage={ + "prompt_tokens": 100, + "completion_tokens": 5, + "total_tokens": 105, + "cache_read_input_tokens": 90, + }, + choices=[], + ) + ) + + assert result.choices == [] + assert result.usage is not None + assert result.usage.prompt_tokens_details.cached_tokens == 90 + + +def test_chunk_parser_without_usage_still_parses_content(): + iterator = DatabricksChatResponseIterator(streaming_response=None, sync_stream=True) + + result = iterator.chunk_parser(_streaming_chunk()) + + assert result.id == "chatcmpl-test" + assert result.model == "databricks-claude-sonnet-5" + assert result.choices[0]["delta"]["content"] == "hi" From 732e23a4f91482cebdf2b05aad572a7ec001e0b7 Mon Sep 17 00:00:00 2001 From: pokepoke81 <4258646+pokepoke81@users.noreply.github.com> Date: Fri, 14 Aug 2026 11:24:30 -0400 Subject: [PATCH 039/121] Remove comment about prompt-cache usage in test Remove outdated comment regarding prompt-cache counts in chunk_parser. --- .../databricks/chat/test_databricks_chat_transformation.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index d6b8e1a3652..165046a2298 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -444,9 +444,6 @@ def _streaming_chunk(usage=None, choices=None): ids=["warm_cache_read", "cold_cache_write"], ) def test_chunk_parser_surfaces_prompt_cache_usage(cache_read, cache_creation, expected_cached, expected_written): - """Databricks returns Anthropic prompt-cache counts in the streaming usage object, - but chunk_parser dropped usage entirely, so cache-aware pricing never reached the - cost calculator and every streamed request was billed at the full input rate.""" iterator = DatabricksChatResponseIterator(streaming_response=None, sync_stream=True) result = iterator.chunk_parser( From 4c49d03732dc107dc52ea7f469e848ec556fec85 Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Fri, 14 Aug 2026 14:56:06 -0400 Subject: [PATCH 040/121] fix(anthropic): preserve optional Responses tool properties Translating Anthropic tools left the outbound function-tool `strict` unset, which the Responses API does not read as non-strict. OpenAI's function-calling docs say strict mode requires every field in `properties` to be marked required, and with `strict` omitted the schema gets normalized to satisfy that instead of being rejected. What users see is a tool whose `required` lists every property, so models fill optional Anthropic tool arguments with empty values. Send `strict` explicitly so an unset value stays non-strict and an explicit `strict: true` still reaches the provider On the Chat Completions adapter, `strict` was also missing from `mapped_tool_params`, so a tool-level `strict` was merged into the OpenAI function `parameters` schema (mutating the caller's `input_schema` along the way) instead of being set on the function. Map it to `function.strict` and leave it unset when the caller omits it, since Chat Completions already defaults to non-strict --- .../adapters/transformation.py | 3 + .../responses_adapters/transformation.py | 8 ++- litellm/types/llms/anthropic.py | 3 +- ...al_pass_through_adapters_transformation.py | 47 ++++++++++++++++ .../test_responses_adapters_transformation.py | 55 +++++++++++++++++++ 5 files changed, 114 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 51f2b661421..ea0eebe0511 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -741,6 +741,7 @@ class LiteLLMAnthropicMessagesAdapter: "input_schema", "description", "cache_control", + "strict", "type", ] @@ -770,6 +771,8 @@ class LiteLLMAnthropicMessagesAdapter: function_chunk["parameters"] = tool["input_schema"] if "description" in tool: function_chunk["description"] = tool["description"] + if "strict" in tool: + function_chunk["strict"] = bool(tool["strict"]) for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index bf3f6153e7c..03e66388c91 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -231,7 +231,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if (isinstance(tool_type, str) and tool_type.startswith("web_search")) or tool_name == "web_search": result.append({"type": "web_search_preview"}) continue - func_tool: dict[str, Any] = {"type": "function", "name": tool_name} + # Responses turns strict mode on when `strict` is omitted, silently rewriting + # `required` to every property. Anthropic tools are non-strict unless asked. + func_tool: dict[str, Any] = { + "type": "function", + "name": tool_name, + "strict": bool(tool_dict.get("strict")), + } if "description" in tool_dict: func_tool["description"] = tool_dict["description"] if "input_schema" in tool_dict: diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 69d291eebd0..17ba78b0190 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict -from typing_extensions import NotRequired, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from .openai import ( ChatCompletionCachedContent, @@ -48,6 +48,7 @@ class AnthropicMessagesTool(TypedDict, total=False): name: Required[str] description: str input_schema: AnthropicInputSchema | None + strict: ReadOnly[bool] type: Literal["custom"] cache_control: dict | ChatCompletionCachedContent | None defer_loading: bool diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index fe6adade6a8..0c30d8a8322 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -3508,3 +3508,50 @@ def test_translate_anthropic_tools_to_openai_preserves_parameters_type(): params = new_tools[0]["function"]["parameters"] assert params["type"] == "object" assert new_tools[0]["type"] == "function" + + +def test_translate_anthropic_tools_to_openai_maps_strict_onto_function_not_parameters(): + """A tool-level `strict` lands on the OpenAI function, leaving the caller's `input_schema` untouched.""" + adapter = LiteLLMAnthropicMessagesAdapter() + input_schema = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + } + tools = [{"type": "custom", "name": "get_weather", "strict": True, "input_schema": input_schema}] + + new_tools, _ = adapter.translate_anthropic_tools_to_openai(tools=tools) + + function = new_tools[0]["function"] + assert function["strict"] is True + assert "strict" not in function["parameters"] + assert input_schema == { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, + } + + +def test_translate_anthropic_tools_to_openai_omits_unset_strict(): + """Chat Completions already defaults to non-strict, so an unset `strict` stays unset.""" + adapter = LiteLLMAnthropicMessagesAdapter() + tools = [ + { + "type": "custom", + "name": "search", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}, "cursor": {"type": "string"}}, + "required": ["query"], + }, + } + ] + + new_tools, _ = adapter.translate_anthropic_tools_to_openai(tools=tools) + + function = new_tools[0]["function"] + assert "strict" not in function + assert "strict" not in function["parameters"] + assert function["parameters"]["required"] == ["query"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index a736ca684aa..90733dc9134 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -605,6 +605,7 @@ class TestTranslateToolsToResponsesAPI: { "type": "function", "name": "get_weather", + "strict": False, "description": "Get current weather for a city.", "parameters": { "type": "object", @@ -614,6 +615,60 @@ class TestTranslateToolsToResponsesAPI: } ] + def test_tool_with_optional_properties_stays_non_strict(self): + """Regression: an unset Anthropic `strict` must not become the Responses strict default, + which would rewrite `required` to include every optional property.""" + tools = [ + { + "name": "search", + "input_schema": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "cursor": {"type": "string"}, + }, + "required": ["query"], + "additionalProperties": False, + }, + } + ] + + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + + assert result[0]["strict"] is False + assert result[0]["parameters"]["required"] == ["query"] + + def test_tool_forwards_explicit_strict_true(self): + """An explicit Anthropic `strict: True` still reaches Responses as True.""" + tools = [ + { + "name": "search", + "strict": True, + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + } + ] + + result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + + assert result == [ + { + "type": "function", + "name": "search", + "strict": True, + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + } + ] + def test_tool_without_description(self): """Tool without a description omits the description key.""" tools = [{"name": "ping", "input_schema": {"type": "object", "properties": {}}}] From 9858d021eef07fefb955d7dc9d4c8e1595afb495 Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Fri, 14 Aug 2026 15:13:24 -0400 Subject: [PATCH 041/121] fix(guardrails): record MCP tool guardrail evaluations and blocks in usage monitor MCP tool calls run their guardrails against a throwaway LLM-shaped dict built by `ProxyLogging._convert_mcp_to_llm_format`, not against the dict the tool call is logged from. `@log_guardrail_information` therefore appended `standard_logging_guardrail_information` to that throwaway dict's metadata bucket, where `get_standard_logging_object_payload` never saw it, so the Guardrails Monitor reported zero evaluations and zero blocks for all MCP traffic. Thread the request's `litellm_logging_obj` into `pre_call_tool_check` and `_create_during_hook_task` and bridge the guardrail records onto it: - Seed `data["litellm_logging_obj"]`, which unified guardrails read and pass into `apply_guardrail`. - Call `_sync_guardrail_info_to_logging_obj` in a `finally`, which is what native guardrails need and what makes the block path work: a blocked call raises straight out of `pre_call_tool_check`, so the record has to be attached before the exception leaves the frame. Only the guardrail evaluation records are copied. The synthetic request's messages and tool arguments are deliberately left behind -- they can carry end-user data and nothing in the monitor needs them. In `call_mcp_tool`, flush the failure handlers before `post_call_failure_hook` so the `status="failure"` standard logging object exists when `_ProxyDBLogger.async_post_call_failure_hook` writes the spend-log row the monitor's "Total Blocked" counts. Both handlers gate on `should_run_logging("sync_failure")` / `("async_failure")` and then mark it, so the `@client` wrapper's own post-raise logging is a no-op and nothing is double-counted -- the same pattern `_fire_mcp_tool_call_logging` already uses for `isError=True`. Threaded through every MCP tool entry point: the managed-server path, the local-OpenAPI registry path, the legacy registry fallback, and the Responses API's `_execute_tool_calls`. --- .../mcp_server/mcp_server_manager.py | 89 +++++- .../proxy/_experimental/mcp_server/server.py | 17 + .../mcp/litellm_proxy_mcp_handler.py | 1 + .../mcp_server/test_mcp_block_recording.py | 126 ++++++++ .../test_mcp_guardrail_usage_monitor.py | 301 ++++++++++++++++++ .../mcp/test_litellm_proxy_mcp_handler.py | 36 +++ 6 files changed, 560 insertions(+), 10 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a1adda2bc95..0ad372064cd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast from urllib.parse import ParseResult, urlparse @@ -46,6 +46,9 @@ from litellm.constants import ( ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth +from litellm.integrations.custom_guardrail import ( + _sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic +) from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -162,6 +165,7 @@ if TYPE_CHECKING: from mcp.types import CreateMessageRequestParams from litellm.caching.caching import InMemoryCache + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.mcp_server.mcp_toolset import MCPToolset try: @@ -1233,6 +1237,35 @@ def _create_elicitation_callback(): return _elicitation_callback +def _record_mcp_guardrail_evaluations( + synthetic_llm_data: dict[str, Any], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict + litellm_logging_obj: "LiteLLMLoggingObj | None", +) -> None: + """Bridge guardrail decision records off an MCP synthetic request onto the request's logger. + + MCP guardrails run against a throwaway LLM-shaped dict from + ``ProxyLogging._convert_mcp_to_llm_format``, so ``@log_guardrail_information`` + files ``standard_logging_guardrail_information`` in that dict's metadata bucket, + which ``get_standard_logging_object_payload`` never reads. Native (non-unified) + guardrails receive no ``logging_obj`` kwarg, so the decorator cannot bridge on + their behalf; this calls the same helper it would have. + + Only the decision records move. The synthetic request's messages and tool + arguments stay behind: they can carry end-user data, and the monitor needs none + of it. + """ + if litellm_logging_obj is None: + return + + try: + _sync_guardrail_info_to_logging_obj(synthetic_llm_data, litellm_logging_obj) + except Exception as e: # noqa: BLE001 # callers run this from a `finally` on the block path + # The breadth is the point. Narrowing to the knowable AttributeError/TypeError + # would let an unexpected type escape that ``finally`` and replace the guardrail's + # block with a bookkeeping error. + verbose_logger.warning("Failed to record MCP guardrail evaluation for logging: %s", e) + + class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") @@ -4543,6 +4576,7 @@ class MCPServerManager: proxy_logging_obj: ProxyLogging | None, server: MCPServer, raw_headers: dict[str, str] | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> dict[str, Any]: """ Run pre-call checks and guardrail hooks for an MCP tool call. @@ -4552,6 +4586,10 @@ class MCPServerManager: present. An absent logger must never be able to turn an authorization decision into a no-op. + ``litellm_logging_obj`` is the request's logger, and it is what lands a + ``pre_mcp_call`` evaluation (or a block) on the spend-log row the Guardrails + Monitor counts. It stays optional so callers that do no logging are unchanged. + Returns a dict that may contain: - "arguments": hook-modified tool arguments (only if changed) - "extra_headers": headers injected by pre_mcp_call guardrail hooks @@ -4610,8 +4648,13 @@ class MCPServerManager: # Create MCP request object for processing mcp_request_obj: Final = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs) - # Convert to LLM format for existing guardrail compatibility + # Convert to LLM format for existing guardrail compatibility. + # Unified guardrails read the seeded logger off the request dict and pass it + # into ``apply_guardrail``, so ``@log_guardrail_information`` bridges their + # evaluations itself; the ``finally`` below covers native guardrails, which + # never receive it. Same seeding the pass-through routes do. synthetic_llm_data: Final = proxy_logging_obj._convert_mcp_to_llm_format(mcp_request_obj, pre_hook_kwargs) + synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj try: # Use standard pre_call_hook @@ -4636,6 +4679,12 @@ class MCPServerManager: # Re-raise guardrail exceptions to properly fail the MCP call verbose_logger.error("Guardrail blocked MCP tool call pre call: %s", e) raise e + finally: + # ``finally`` rather than after the ``try``: a block raises straight out of + # here, and the failure spend-log row that "Total Blocked" counts is built + # from this logger further up the stack, so the record has to be attached + # before the exception leaves this frame. + _record_mcp_guardrail_evaluations(synthetic_llm_data, litellm_logging_obj) return hook_result @@ -4647,8 +4696,14 @@ class MCPServerManager: user_api_key_auth: UserAPIKeyAuth | None, proxy_logging_obj: ProxyLogging, start_time: datetime.datetime, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ): - """Create and return a during hook task for MCP tool calls.""" + """Create and return a during hook task for MCP tool calls. + + ``litellm_logging_obj`` is the request's logger; see ``pre_call_tool_check``. + The task is awaited before the tool call's success logging runs, so a + ``during_mcp_call`` evaluation recorded on it is serialized with that call. + """ from litellm.types.llms.base import HiddenParams from litellm.types.mcp import MCPDuringCallRequestObject @@ -4667,15 +4722,23 @@ class MCPServerManager: "user_api_key_auth": user_api_key_auth, } + # Seeded for the same reason as in ``pre_call_tool_check``. synthetic_llm_data: Final = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, during_hook_kwargs) + synthetic_llm_data["litellm_logging_obj"] = litellm_logging_obj - return asyncio.create_task( - proxy_logging_obj.during_call_hook( - user_api_key_dict=user_api_key_auth, - data=synthetic_llm_data, - call_type=CallTypes.call_mcp_tool.value, - ) - ) + # Wrapped so the bridge runs inside the task: the caller only holds the task and + # gathers it later, so there is no other point that still sees a block here. + async def _run_during_call_hook() -> Mapping[str, Any] | None: + try: + return await proxy_logging_obj.during_call_hook( + user_api_key_dict=user_api_key_auth, + data=synthetic_llm_data, + call_type=CallTypes.call_mcp_tool.value, + ) + finally: + _record_mcp_guardrail_evaluations(synthetic_llm_data, litellm_logging_obj) + + return asyncio.create_task(_run_during_call_hook()) def _get_call_semaphore(self, mcp_server: MCPServer) -> asyncio.Semaphore | None: limit: Final = mcp_server.max_concurrent_requests @@ -5204,6 +5267,7 @@ class MCPServerManager: oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, host_progress_callback: Callable | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> CallToolResult: """ Call a tool with the given name and arguments @@ -5216,6 +5280,9 @@ class MCPServerManager: mcp_auth_header: MCP auth header (deprecated) mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} proxy_logging_obj: Optional ProxyLogging object for hook integration + litellm_logging_obj: Optional request logger the guardrail hooks record + their evaluations onto, so MCP guardrail activity reaches the + Guardrails Monitor. See ``pre_call_tool_check`` Returns: @@ -5246,6 +5313,7 @@ class MCPServerManager: proxy_logging_obj=proxy_logging_obj, server=mcp_server, raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, ) if "arguments" in hook_result: arguments = hook_result["arguments"] @@ -5260,6 +5328,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, proxy_logging_obj=proxy_logging_obj, start_time=start_time, + litellm_logging_obj=litellm_logging_obj, ) tasks.append(during_hook_task) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f237529b319..17457c3362f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2824,6 +2824,7 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, server=mcp_server, raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, ) # `pre_call_tool_check` may return guardrail-modified # arguments; honor them on the local path too. @@ -2962,6 +2963,7 @@ if MCP_AVAILABLE: proxy_logging_obj=proxy_logging_obj, server=prefix_server, raw_headers=raw_headers, + litellm_logging_obj=litellm_logging_obj, ) if "arguments" in hook_result: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args @@ -3149,6 +3151,20 @@ if MCP_AVAILABLE: traceback_str: Final = traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG) from litellm.proxy.proxy_server import proxy_logging_obj + # Ordering is load-bearing. ``_ProxyDBLogger.async_post_call_failure_hook``, + # reached below, writes the failure spend-log row from this logger's + # ``standard_logging_object``, which only exists once the failure handlers + # have run. Flush them first or the row lands with + # ``guardrail_information=None`` and a guardrail block is never counted. + # + # Not double-logged: both handlers gate on ``should_run_logging`` and then + # mark it, so the ``@client`` wrapper's own post-raise logging no-ops on this + # logger, same as ``_fire_mcp_tool_call_logging`` does for ``isError=True``. + if litellm_logging_obj is not None: + end_time: Final = datetime.now() # noqa: DTZ005 # naive to match `start_time`, which it is subtracted from + litellm_logging_obj.failure_handler(e, traceback_str, start_time, end_time) + await litellm_logging_obj.async_failure_handler(e, traceback_str, start_time, end_time) + if proxy_logging_obj and user_api_key_auth: await proxy_logging_obj.post_call_failure_hook( request_data=kwargs, @@ -3326,6 +3342,7 @@ if MCP_AVAILABLE: raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, host_progress_callback=host_progress_callback, + litellm_logging_obj=litellm_logging_obj, ) verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result) return call_tool_result diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 56818717c09..0321034dffe 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -798,6 +798,7 @@ class LiteLLM_Proxy_MCP_Handler: oauth2_headers=oauth2_headers, raw_headers=raw_headers, proxy_logging_obj=proxy_logging_obj, + litellm_logging_obj=litellm_logging_obj, ) if proxy_logging_obj: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py new file mode 100644 index 00000000000..64d926bc5e3 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_block_recording.py @@ -0,0 +1,126 @@ +"""Tests for guardrail-block recording in +``litellm.proxy._experimental.mcp_server.server.call_mcp_tool``. + +A pre-call MCP guardrail block *raises* into ``call_mcp_tool``'s +``except Exception``. The failure spend-log row that the Guardrails Monitor's +"Total Blocked" counts is written by ``_ProxyDBLogger.async_post_call_failure_hook`` +(reached via ``proxy_logging_obj.post_call_failure_hook``), which reads +``standard_logging_object`` off the request's logging obj -- and that only exists +once ``failure_handler`` / ``async_failure_handler`` have run. So the failure +handlers must run *before* ``post_call_failure_hook``, otherwise the row persists +with ``guardrail_information=None`` and the block is never counted. These tests +pin that ordering. + +``call_mcp_tool`` is wrapped by ``@client`` (``litellm.utils.client``), which uses +``functools.wraps`` and therefore exposes the raw undecorated coroutine as +``__wrapped__``. The tests drive ``__wrapped__`` directly so the except-block +ordering is observed in isolation, without the wrapper's own post-raise logging +firing. Note that this means they do not exercise the wrapper's dedup path; that +dedup rests on ``should_run_logging("sync_failure")`` / ``("async_failure")``, +which has its own coverage in the logging tests. + +``proxy_logging_obj`` is imported lazily inside the except block via +``from litellm.proxy.proxy_server import proxy_logging_obj``; the real +``proxy_server`` module is heavy, so a fake module is injected into ``sys.modules`` +to satisfy that lazy import without loading it. +""" + +import contextlib +import sys +import types +from unittest import mock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._experimental.mcp_server import server + + +class _RecordingLoggingObj: + """Stands in for ``LiteLLMLoggingObj``, recording the failure flush the fix + makes so the test can assert it happens before ``post_call_failure_hook``.""" + + def __init__(self, order: list) -> None: + self._order = order + self.failure_calls = 0 + self.async_failure_calls = 0 + + def failure_handler(self, *_args, **_kwargs) -> None: + self.failure_calls += 1 + self._order.append("failure_handler") + + async def async_failure_handler(self, *_args, **_kwargs) -> None: + self.async_failure_calls += 1 + self._order.append("async_failure_handler") + + +async def _call_block(logging_obj, order: list, *, user_api_key_auth=mock.sentinel.auth): + """Drive ``call_mcp_tool`` into its except path via ``arguments=None``, which + raises ``HTTPException(400)`` before any server-manager call, and return once it + re-raises.""" + + async def _record_post_call_failure_hook(**_kwargs) -> None: + order.append("post_call_failure_hook") + + proxy_logging_obj = mock.MagicMock() + proxy_logging_obj.post_call_failure_hook.side_effect = _record_post_call_failure_hook + + fake_proxy_server = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy_server.proxy_logging_obj = proxy_logging_obj # pyright: ignore[reportAttributeAccessIssue] + + with mock.patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}): + with contextlib.suppress(HTTPException): + await server.call_mcp_tool.__wrapped__( + name="t", + arguments=None, + user_api_key_auth=user_api_key_auth, + litellm_logging_obj=logging_obj, + ) + + +@pytest.mark.asyncio +async def test_block_flushes_failure_before_post_call_failure_hook(): + order: list = [] + await _call_block(_RecordingLoggingObj(order), order) + + assert order == ["failure_handler", "async_failure_handler", "post_call_failure_hook"], order + + +@pytest.mark.asyncio +async def test_block_flushes_each_handler_exactly_once(): + """Each handler runs once, so the block yields exactly one counted row rather + than double-counting on the shared logging obj.""" + order: list = [] + obj = _RecordingLoggingObj(order) + await _call_block(obj, order) + + assert (obj.failure_calls, obj.async_failure_calls) == (1, 1) + + +@pytest.mark.asyncio +async def test_block_flushes_failure_for_anonymous_calls(): + """With no ``user_api_key_auth`` the failure handlers still run, so OTel and the + other failure sinks see the block. + + ``post_call_failure_hook`` stays gated on auth, matching the pre-existing + contract: SpendLogs rows are attributable billing/audit records and the + downstream DB logger dereferences authenticated key, budget, and route data. + Counting anonymous MCP blocks needs a counter that does not live in SpendLogs, + which is a separate design change, not part of this fix. + """ + order: list = [] + obj = _RecordingLoggingObj(order) + await _call_block(obj, order, user_api_key_auth=None) + + assert order == ["failure_handler", "async_failure_handler"], order + + +@pytest.mark.asyncio +async def test_absent_logging_obj_still_calls_hook_and_skips_flush(): + """Without a logging obj the flush is skipped (no crash) but + ``post_call_failure_hook`` still fires. Byte-equivalent to stock behavior for + that branch; its value is as a mutation-killer for the ``is not None`` guard.""" + order: list = [] + await _call_block(None, order) + + assert order == ["post_call_failure_hook"], order diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py new file mode 100644 index 00000000000..24e6d2de10d --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_guardrail_usage_monitor.py @@ -0,0 +1,301 @@ +"""Tests for MCP guardrail evaluations reaching the Guardrails Monitor. + +MCP tool calls run their guardrails against a throwaway LLM-shaped dict built by +``ProxyLogging._convert_mcp_to_llm_format``, not against the dict the tool call is +logged from. ``@log_guardrail_information`` therefore appends +``standard_logging_guardrail_information`` to that throwaway dict's metadata +bucket, where ``get_standard_logging_object_payload`` never sees it, so the +Guardrails Monitor reported zero evaluations and zero blocks for MCP traffic. + +``pre_call_tool_check`` and ``_create_during_hook_task`` now take the request's +``litellm_logging_obj`` and bridge those records onto it. These tests pin both the +seeding (which unified guardrails consume off ``data["litellm_logging_obj"]``) and +the bridge (which native guardrails depend on), including on the block path. +""" + +import asyncio +import datetime +from typing import Any +from unittest import mock + +import pytest + +from litellm.exceptions import GuardrailRaisedException +from litellm.proxy._experimental.mcp_server import mcp_server_manager as MOD + + +class _FakeLoggingObj: + """Minimal stand-in for ``LiteLLMLoggingObj``. + + ``_sync_guardrail_info_to_logging_obj`` reads exactly these two attributes, + and the spend-log payload is built from ``litellm_params["metadata"]``, so a + real ``Logging`` instance would add setup cost without adding coverage. + """ + + def __init__(self) -> None: + self.litellm_params: dict[str, Any] = {"metadata": {}} + self.model_call_details: dict[str, Any] = {"litellm_params": self.litellm_params} + + @property + def recorded_guardrails(self) -> list: + return self.litellm_params["metadata"].get("standard_logging_guardrail_information", []) + + +def _bare_manager() -> MOD.MCPServerManager: + """An ``MCPServerManager`` without running ``__init__``. + + The authorization/validation helpers on the path are stubbed out so the test + reaches the guardrail hooks; they have their own coverage elsewhere. + """ + mgr = MOD.MCPServerManager.__new__(MOD.MCPServerManager) + mgr.check_allowed_or_banned_tools = lambda name, server: True + mgr.validate_allowed_params = lambda tool_name, arguments, server: None + + async def _ok(*_args, **_kwargs) -> None: + return None + + mgr.check_tool_permission_for_key_team = _ok + return mgr + + +def _fake_proxy_logging(capture: dict, *, guardrail_effect=None): + """A ``proxy_logging_obj`` double whose hooks capture the data they receive. + + ``guardrail_effect`` stands in for a guardrail: it is handed the synthetic + request dict so it can append a guardrail record (and optionally raise, the + way a blocking guardrail does). + """ + plo = mock.MagicMock() + plo._create_mcp_request_object_from_kwargs.return_value = mock.MagicMock() + # Mirror the real conversion's metadata bucket so a test can prove it survives. + plo._convert_mcp_to_llm_format.side_effect = lambda *_a, **_k: { + "metadata": {"headers": {"x-forwarded-for": "1.2.3.4"}} + } + + async def _hook(*, user_api_key_dict, data, call_type) -> None: + del user_api_key_dict # captured shape is what matters, not the auth double + capture["data"] = data + capture["call_type"] = call_type + if guardrail_effect is not None: + guardrail_effect(data) + + plo.pre_call_hook.side_effect = _hook + plo.during_call_hook.side_effect = _hook + return plo + + +def _record_guardrail(status: str = "success"): + """Write a guardrail record the way ``@log_guardrail_information`` does.""" + + def _effect(data: dict) -> None: + data.setdefault("metadata", {}).setdefault("standard_logging_guardrail_information", []).append( + {"guardrail_name": "test-guardrail", "guardrail_status": status} + ) + + return _effect + + +def _blocking_guardrail(): + record = _record_guardrail(status="guardrail_intervened") + + def _effect(data: dict) -> None: + record(data) + raise GuardrailRaisedException(guardrail_name="test-guardrail", message="blocked") + + return _effect + + +async def _run_pre_call(mgr, plo, logging_obj) -> dict: + return await mgr.pre_call_tool_check( + name="t", + arguments={}, + server_name="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + server=mock.MagicMock(), + raw_headers={}, + litellm_logging_obj=logging_obj, + ) + + +@pytest.mark.asyncio +async def test_pre_call_seeds_request_logging_obj_for_unified_guardrails(): + """Unified guardrails read ``data["litellm_logging_obj"]`` and pass it into + ``apply_guardrail``, whose ``@log_guardrail_information`` wrapper bridges the + evaluation onto that logger itself. Drop the seed and that path records + nothing.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + await _run_pre_call(_bare_manager(), _fake_proxy_logging(capture), logging_obj) + + assert capture["data"]["litellm_logging_obj"] is logging_obj + + +@pytest.mark.asyncio +async def test_pre_call_keeps_synthetic_request_headers_metadata(): + """The seed must not clobber the metadata bucket ``_convert_mcp_to_llm_format`` + builds: guardrails such as ``MCPJWTSigner`` read ``metadata["headers"]`` off + it.""" + capture: dict = {} + await _run_pre_call(_bare_manager(), _fake_proxy_logging(capture), _FakeLoggingObj()) + + assert capture["data"]["metadata"]["headers"] == {"x-forwarded-for": "1.2.3.4"} + + +@pytest.mark.asyncio +async def test_pre_call_bridges_allowed_evaluation_onto_request_logger(): + """An allowed ``pre_mcp_call`` evaluation must land on the request logger, which + is what the monitor's "Total Evaluations" counts.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail()) + + await _run_pre_call(_bare_manager(), plo, logging_obj) + + assert logging_obj.recorded_guardrails == [{"guardrail_name": "test-guardrail", "guardrail_status": "success"}] + + +@pytest.mark.asyncio +async def test_pre_call_bridges_blocked_evaluation_before_reraising(): + """A block raises straight out of ``pre_call_tool_check``, and the failure + spend-log row that "Total Blocked" counts is built from this logger further up + the stack. So the record has to be attached before the exception leaves the + frame -- hence the bridge lives in a ``finally``.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail()) + + with pytest.raises(GuardrailRaisedException): + await _run_pre_call(_bare_manager(), plo, logging_obj) + + assert logging_obj.recorded_guardrails == [ + {"guardrail_name": "test-guardrail", "guardrail_status": "guardrail_intervened"} + ] + + +@pytest.mark.asyncio +async def test_pre_call_without_logging_obj_is_unchanged(): + """Callers that thread no logger are unaffected: the seed is an explicit + ``None`` (which every consumer reads via ``.get``) and nothing is bridged. + Guards against the bridge assuming a logger exists.""" + capture: dict = {} + plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail()) + mgr = _bare_manager() + + result = await mgr.pre_call_tool_check( + name="t", + arguments={}, + server_name="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + server=mock.MagicMock(), + raw_headers={}, + ) + + assert result == {} + assert capture["data"]["litellm_logging_obj"] is None + + +@pytest.mark.asyncio +async def test_during_hook_seeds_and_bridges_onto_request_logger(): + """``during_mcp_call`` evaluations need the same treatment. The task is awaited + before the tool call's success logging runs, so the record is serialized with + that call.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_record_guardrail()) + + await _bare_manager()._create_during_hook_task( + name="t", + arguments={}, + server_name_from_prefix="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + start_time=datetime.datetime(2026, 7, 14), + litellm_logging_obj=logging_obj, + ) + + assert capture["data"]["litellm_logging_obj"] is logging_obj + assert logging_obj.recorded_guardrails == [{"guardrail_name": "test-guardrail", "guardrail_status": "success"}] + + +@pytest.mark.asyncio +async def test_during_hook_bridges_even_when_hook_raises(): + """A during-call guardrail block must still be recorded before the task's + exception propagates to the ``asyncio.gather`` in ``call_tool``.""" + capture: dict = {} + logging_obj = _FakeLoggingObj() + plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail()) + + task = _bare_manager()._create_during_hook_task( + name="t", + arguments={}, + server_name_from_prefix="s", + user_api_key_auth=None, + proxy_logging_obj=plo, + start_time=datetime.datetime(2026, 7, 14), + litellm_logging_obj=logging_obj, + ) + with pytest.raises(GuardrailRaisedException): + await task + + assert logging_obj.recorded_guardrails == [ + {"guardrail_name": "test-guardrail", "guardrail_status": "guardrail_intervened"} + ] + + +@pytest.mark.asyncio +async def test_bridge_failure_does_not_mask_a_guardrail_block(): + """Recording is best-effort bookkeeping. If the bridge itself raises, the guardrail's + block must still be what the caller sees, not a bookkeeping error. + + The bridge is forced to fail by making the logger's ``model_call_details`` raise, and + the swallow is asserted (not just the surviving exception type) so the test cannot go + vacuous if a refactor stops the bridge from touching that attribute. + """ + capture: dict = {} + plo = _fake_proxy_logging(capture, guardrail_effect=_blocking_guardrail()) + + broken_logging_obj = mock.MagicMock() + type(broken_logging_obj).model_call_details = mock.PropertyMock(side_effect=RuntimeError("boom")) + + with mock.patch.object(MOD.verbose_logger, "warning") as warn: + with pytest.raises(GuardrailRaisedException): + await _run_pre_call(_bare_manager(), plo, broken_logging_obj) + + assert warn.call_count == 1, "the bridge did not actually fail, so this test proves nothing" + assert "boom" in str(warn.call_args) + + +@pytest.mark.asyncio +async def test_call_tool_threads_logging_obj_into_both_hooks(): + """``call_tool`` is the single entry point every MCP dispatch route funnels + through, so it must hand the logger to both guardrail hook sites.""" + mgr = _bare_manager() + logging_obj = _FakeLoggingObj() + seen: dict = {} + + async def _fake_pre_call_tool_check(**kwargs): + seen["pre_call"] = kwargs.get("litellm_logging_obj") + return {} + + def _fake_during_hook_task(**kwargs): + seen["during_call"] = kwargs.get("litellm_logging_obj") + return asyncio.get_running_loop().create_future() + + mgr.pre_call_tool_check = _fake_pre_call_tool_check + mgr._create_during_hook_task = _fake_during_hook_task + mgr._resolve_mcp_server_for_tool_call = lambda server_name, name: mock.MagicMock(spec_path=None) + mgr._resolve_oauth2_headers_for_tool_call = mock.AsyncMock(return_value=None) + mgr._call_regular_mcp_tool = mock.AsyncMock(return_value=mock.MagicMock()) + + with mock.patch.object(MOD, "_resolve_byok_mcp_auth_header", mock.AsyncMock(return_value=None)): + await mgr.call_tool( + server_name="s", + name="t", + arguments={}, + proxy_logging_obj=mock.MagicMock(), + litellm_logging_obj=logging_obj, + ) + + assert seen == {"pre_call": logging_obj, "during_call": logging_obj} diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 4981caa10c3..418c716b1a1 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -450,6 +450,42 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio assert captured.get("litellm_trace_id") == "tid" +@pytest.mark.asyncio +async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch): + """The Responses-API MCP path must hand the request's litellm_logging_obj to + global_mcp_server_manager.call_tool, otherwise pre_call_tool_check / + _create_during_hook_task get None and no guardrail evaluation is bridged onto + the request logger, so MCP tool calls made through the Responses API report zero + guardrail evaluations in the monitor. Drop the litellm_logging_obj kwarg on the + call_tool invocation and this fails.""" + _setup_proxy_logging(monkeypatch) + call_tool_mock = _setup_mcp_call_environment(monkeypatch) + + sentinel_logging_obj = MagicMock() + sentinel_logging_obj.async_post_mcp_tool_call_hook = AsyncMock() + sentinel_logging_obj.async_success_handler = AsyncMock() + + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") + monkeypatch.setattr( + handler_module, + "function_setup", + lambda *_args, **_kwargs: (sentinel_logging_obj, None), + ) + + tool_name = "deepwiki-read_wiki_structure" + tool_calls = [{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}] + + await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map={tool_name: "deepwiki"}, + tool_calls=tool_calls, + user_api_key_auth=None, + ) + + assert call_tool_mock.await_count == 1 + assert call_tool_mock.await_args is not None + assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj + + @pytest.mark.asyncio async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch): """ From a62798de63873f146efde812918fb24d30ed0621 Mon Sep 17 00:00:00 2001 From: Scott Wilson Date: Fri, 14 Aug 2026 17:56:44 -0400 Subject: [PATCH 042/121] test(anthropic): type the Responses tool fixtures instead of suppressing The two new `translate_tools_to_responses_api` calls carried `# type: ignore[arg-type]`, which CLAUDE.md bans as LIT009: pyrightconfig.json sets enableTypeIgnoreComments to false, so the comment silently does nothing and the reportArgumentType error stands. Annotating the fixtures as list[AllAnthropicToolsValues] makes both calls check clean with no suppression at all. --- .../test_responses_adapters_transformation.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 90733dc9134..8b34163e075 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -21,7 +21,10 @@ from litellm.constants import ( from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) -from litellm.types.llms.anthropic import AnthropicMessagesRequest +from litellm.types.llms.anthropic import ( + AllAnthropicToolsValues, + AnthropicMessagesRequest, +) from litellm.types.llms.openai import ResponseAPIUsage @@ -618,7 +621,7 @@ class TestTranslateToolsToResponsesAPI: def test_tool_with_optional_properties_stays_non_strict(self): """Regression: an unset Anthropic `strict` must not become the Responses strict default, which would rewrite `required` to include every optional property.""" - tools = [ + tools: List[AllAnthropicToolsValues] = [ { "name": "search", "input_schema": { @@ -633,14 +636,14 @@ class TestTranslateToolsToResponsesAPI: } ] - result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + result = _ADAPTER.translate_tools_to_responses_api(tools) assert result[0]["strict"] is False assert result[0]["parameters"]["required"] == ["query"] def test_tool_forwards_explicit_strict_true(self): """An explicit Anthropic `strict: True` still reaches Responses as True.""" - tools = [ + tools: List[AllAnthropicToolsValues] = [ { "name": "search", "strict": True, @@ -653,7 +656,7 @@ class TestTranslateToolsToResponsesAPI: } ] - result = _ADAPTER.translate_tools_to_responses_api(tools) # type: ignore[arg-type] + result = _ADAPTER.translate_tools_to_responses_api(tools) assert result == [ { From 48de8106efc0295578dc6b99036ae0cda6b963c1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 23:59:39 +0000 Subject: [PATCH 043/121] fix(router): stop get_router_model_info from wiping cached pricing Merge deployment model_info into a copy of the lru_cache'd get_model_info() dict and drop unset Nones, so Deployment's mirrored pricing defaults no longer overwrite built-in prices process-wide. Fixes #36980 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 16 +++++++++--- tests/test_litellm/test_router.py | 42 +++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index fb2af41dcf2..baa4724b8eb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -29,6 +29,7 @@ import anyio import httpx import openai from openai import AsyncOpenAI +from pydantic import BaseModel from typing_extensions import overload import litellm @@ -9080,12 +9081,19 @@ class Router: model_info: Final = litellm.get_model_info(model=model_info_name) ## CHECK USER SET MODEL INFO - user_model_info: Final = deployment.get("model_info") or {} + raw_user_model_info: Final = deployment.get("model_info") or {} + user_model_info: Final = ( + raw_user_model_info.model_dump(exclude_none=True) + if isinstance(raw_user_model_info, BaseModel) + else {key: value for key, value in raw_user_model_info.items() if value is not None} + ) - if model_info is not None: - model_info.update(cast(ModelInfo, user_model_info)) + if model_info is None: + return model_info - return model_info + # get_model_info() hands back an lru_cache'd dict; merging into a copy keeps + # deployment overrides out of the shared entry + return cast(ModelMapInfo, {**model_info, **user_model_info}) def get_model_info(self, id: str) -> dict | None: """ diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bdbf33fb0e1..b3c348a1221 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7954,3 +7954,45 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): finally: for cb in router.optional_callbacks or []: litellm.logging_callback_manager.remove_callback_from_all_lists(cb) + + +def test_get_router_model_info_does_not_wipe_cached_pricing(): + """A Deployment's model_info declares the mirrored pricing fields with None defaults; + merging it must not write those Nones into the lru_cache'd dict get_model_info() owns, + or /model/info loses built-in prices for every model a worker serves.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + litellm.get_model_info.cache_clear() + expected = copy.deepcopy(litellm.get_model_info(model="anthropic/claude-sonnet-4-5")) + + router = litellm.Router(model_list=[]) + merged = router.get_router_model_info( + deployment=Deployment( + model_name="sonnet", + litellm_params=LiteLLM_Params(model="claude-sonnet-4-5", custom_llm_provider="anthropic"), + model_info=ModelInfo(id="sonnet-1"), + ), + received_model_name="sonnet", + ) + + assert litellm.get_model_info(model="anthropic/claude-sonnet-4-5") == expected + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert merged[field] == expected[field] + + +def test_get_router_model_info_keeps_explicit_pricing_overrides(): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + litellm.get_model_info.cache_clear() + router = litellm.Router(model_list=[]) + merged = router.get_router_model_info( + deployment=Deployment( + model_name="sonnet", + litellm_params=LiteLLM_Params(model="claude-sonnet-4-5", custom_llm_provider="anthropic"), + model_info=ModelInfo(id="sonnet-1", input_cost_per_token=1e-08), + ), + received_model_name="sonnet", + ) + + assert merged["input_cost_per_token"] == 1e-08 + assert litellm.get_model_info(model="anthropic/claude-sonnet-4-5")["input_cost_per_token"] != 1e-08 From eafddaaa12a75ad70d9a1dc4ec90f3f611e8a130 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 17:22:32 -0700 Subject: [PATCH 044/121] feat(scripts): queue heavy gates behind a machine-wide slot lock --- Makefile | 21 +- scripts/gate_slot_lock.py | 173 ++++++++++++ scripts/pre_commit_lint.sh | 10 + scripts/ruff_strict_gate.py | 5 +- scripts/type_check_gate.py | 23 +- scripts/type_discipline_gate.py | 5 +- tests/test_litellm/test_gate_slot_lock.py | 311 +++++++++++++++++++++ tests/test_litellm/test_pre_commit_lint.py | 41 +++ 8 files changed, 573 insertions(+), 16 deletions(-) create mode 100644 scripts/gate_slot_lock.py create mode 100644 tests/test_litellm/test_gate_slot_lock.py diff --git a/Makefile b/Makefile index 94d8c875af5..7fe5d1f8045 100644 --- a/Makefile +++ b/Makefile @@ -8,8 +8,8 @@ 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 check pre-commit \ - lint-install lint-fetch-base bootstrap + install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \ + lint-install lint-fetch-base bootstrap bootstrap-inner # Default target help: @@ -52,10 +52,17 @@ help: @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" + @echo "" + @echo "Heavy targets (check, bootstrap, lint) queue for LITELLM_GATE_SLOTS machine-wide" + @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." UV := uv UV_RUN := $(UV) run --no-sync +# Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so +# it runs before any venv exists. See scripts/gate_slot_lock.py. +GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py + LINT_DEP_INSTALL ?= install-dev LINT_E2E_DEP_INSTALL ?= lint-install LINT_DEP_BASE ?= lint-fetch-base @@ -74,6 +81,9 @@ install-dev: $(UV) sync --inexact --frozen bootstrap: + @$(GATE_SLOT_LOCK) $(MAKE) bootstrap-inner + +bootstrap-inner: $(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev $(UV_RUN) python scripts/prisma_generate_if_needed.py cd ui/litellm-dashboard && ../../scripts/with_dashboard_node.sh npm install --no-audit --no-fund @@ -230,7 +240,7 @@ check-import-safety: $(LINT_DEP_INSTALL) # base fetch) runs once up front; the checks themselves are independent, so a sub-make # fans them out with -j and the fast ones finish under basedpyright's shadow. lint: lint-install lint-fetch-base - $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks + $(GATE_SLOT_LOCK) $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety @@ -244,7 +254,10 @@ lint-dev: lint-format-changed check-circular-imports check-import-safety # 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 aren't in scope. # Not auto-installed as a git hook so it never slows an unrelated human commit. -check: bootstrap +check: + @$(GATE_SLOT_LOCK) $(MAKE) check-inner + +check-inner: bootstrap ./scripts/pre_commit_lint.sh pre-commit: diff --git a/scripts/gate_slot_lock.py b/scripts/gate_slot_lock.py new file mode 100644 index 00000000000..e7bd945ad65 --- /dev/null +++ b/scripts/gate_slot_lock.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""Machine-wide slot lock for this repo's heavy entrypoints. + +`make check`, `make bootstrap`, `make lint`, and the standalone budget gates +(scripts/ruff_strict_gate.py, scripts/type_discipline_gate.py, +scripts/type_check_gate.py) each hold one of N machine-wide slots while they +run, so however many sessions and worktrees share one machine, at most N of +them execute a basedpyright/pytest/prettier storm at a time instead of all +thrashing it at once. Slots are fcntl.flock files (macOS ships no flock(1) +binary, hence python3 + stdlib only, runnable before any venv exists) under a +per-user cache directory shared by every worktree and session: +~/.cache/litellm/gate-slots by default, $LITELLM_GATE_SLOT_DIR to override. +A holder's lock dies with its process, so a crash leaves nothing to clean up. + +$LITELLM_GATE_SLOTS sets the slot count (default 2); 0 disables locking. +Waiting is a blocking flock on a turnstile file plus a slow poll of the slots, +so contenders queue roughly first-come-first-served without busy-spinning. +A process that acquired (or deliberately skipped) a slot exports +LITELLM_GATE_SLOT_HELD, and nested acquisitions under that marker are no-ops, +so `make check` invoking the gates internally can never deadlock against +itself. Any filesystem error fails open and the command runs unlocked: the +lock is a courtesy to the machine, never a gate that may break a build (CI +runs one job per machine, so there it only ever takes the instant path). + +CLI: python3 scripts/gate_slot_lock.py [args...] +""" + +from __future__ import annotations + +import contextlib +import fcntl +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import IO, TYPE_CHECKING, Final + +if TYPE_CHECKING: + from collections.abc import Iterator + +HELD_MARKER_ENV: Final = "LITELLM_GATE_SLOT_HELD" +SLOT_COUNT_ENV: Final = "LITELLM_GATE_SLOTS" +SLOT_DIR_ENV: Final = "LITELLM_GATE_SLOT_DIR" +DEFAULT_SLOT_COUNT: Final = 2 +POLL_SECONDS: Final = 2.0 + + +def _slot_dir() -> Path: + override: Final = os.environ.get(SLOT_DIR_ENV) + return Path(override) if override else Path.home() / ".cache" / "litellm" / "gate-slots" + + +def _slot_count() -> int: + raw: Final = os.environ.get(SLOT_COUNT_ENV) + if not raw: + return DEFAULT_SLOT_COUNT + try: + return int(raw) + except ValueError: + print( + f"gate_slot_lock: ignoring non-integer {SLOT_COUNT_ENV}={raw!r}; " + f"using {DEFAULT_SLOT_COUNT} slots", + file=sys.stderr, + ) + return DEFAULT_SLOT_COUNT + + +def _try_slot(directory: Path, index: int) -> IO[bytes] | None: + handle: Final = (directory / f"slot-{index}.lock").open("wb") + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + handle.close() + return None + except OSError: + handle.close() + raise + return handle + + +def _wait_for_slot(directory: Path, count: int) -> IO[bytes]: + print( + f"gate_slot_lock: all {count} machine-wide slots are busy; queueing " + f"(set {SLOT_COUNT_ENV}=0 to disable)", + file=sys.stderr, + flush=True, + ) + with (directory / "turnstile.lock").open("wb") as turnstile: + fcntl.flock(turnstile, fcntl.LOCK_EX) + while True: + for index in range(count): + held = _try_slot(directory, index) + if held is not None: + return held + time.sleep(POLL_SECONDS) + + +def _locked_handle(count: int) -> IO[bytes]: + directory: Final = _slot_dir() + directory.mkdir(parents=True, exist_ok=True) + for index in range(count): + immediate = _try_slot(directory, index) + if immediate is not None: + return immediate + return _wait_for_slot(directory, count) + + +def acquire_slot() -> IO[bytes] | None: + """Hold a machine-wide slot for the life of the returned handle. + + The caller must keep the handle referenced until the process exits; + dropping it closes the file and releases the slot. Returns None without + locking when this process already runs under a held slot, when locking is + disabled, or when the filesystem refuses to cooperate.""" + if os.environ.get(HELD_MARKER_ENV): + return None + count: Final = _slot_count() + if count <= 0: + os.environ[HELD_MARKER_ENV] = "1" + return None + try: + handle: Final = _locked_handle(count) + except (OSError, RuntimeError) as error: + print(f"gate_slot_lock: locking unavailable ({error}); running unlocked", file=sys.stderr) + os.environ[HELD_MARKER_ENV] = "1" + return None + os.environ[HELD_MARKER_ENV] = "1" + return handle + + +@contextlib.contextmanager +def held_slot() -> Iterator[None]: + """Run the with-block while holding a machine-wide slot (or its no-op forms).""" + prior_marker: Final = os.environ.get(HELD_MARKER_ENV) + handle: Final = acquire_slot() + try: + yield + finally: + if handle is not None: + handle.close() + if not prior_marker: + os.environ.pop(HELD_MARKER_ENV, None) + + +def _wait_ignoring_interrupts(process: subprocess.Popen[bytes]) -> int: + while True: + try: + return process.wait() + except KeyboardInterrupt: + continue + + +def main() -> int: + if len(sys.argv) < 2: + print("usage: gate_slot_lock.py [args...]", file=sys.stderr) + return 2 + try: + held: Final = acquire_slot() + except KeyboardInterrupt: + return 130 + try: + code: Final = _wait_ignoring_interrupts(subprocess.Popen(sys.argv[1:])) + except FileNotFoundError as error: + print(f"gate_slot_lock: {error}", file=sys.stderr) + return 127 + if held is not None: + held.close() + return code if code >= 0 else 128 - code + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index 82498ec10cd..0861172056e 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -24,6 +24,16 @@ set -eu +# Queue for one of the machine-wide heavy-work slots (see scripts/gate_slot_lock.py) +# before anything else, so N parallel `make check` runs across worktrees execute two +# at a time instead of thrashing the machine. The wrapper exports +# LITELLM_GATE_SLOT_HELD, so this re-exec happens exactly once and everything this +# script spawns (make lint, the budget gates) skips its own acquisition. +if [ -z "${LITELLM_GATE_SLOT_HELD:-}" ]; then + script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0") + exec python3 "$script_dir/gate_slot_lock.py" "$0" "$@" +fi + 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 diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py index 507077ddf25..bf070beeb0f 100644 --- a/scripts/ruff_strict_gate.py +++ b/scripts/ruff_strict_gate.py @@ -215,7 +215,10 @@ def main() -> None: parser.add_argument("--base", default=DEFAULT_BASE) parser.add_argument("--update", action="store_true") args = parser.parse_args() - cmd_update(args.base) if args.update else cmd_check(args.base) + from gate_slot_lock import held_slot + + with held_slot(): + cmd_update(args.base) if args.update else cmd_check(args.base) if __name__ == "__main__": diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index c9f774c6113..763835e6d2e 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -670,16 +670,19 @@ def main() -> None: parser.add_argument("--update", action="store_true") parser.add_argument("--emit-counts-dir", type=Path) args = parser.parse_args() - ensure_typecheck_env() - head = count_basedpyright(run_basedpyright()) - if args.emit_counts_dir is not None: - cmd_emit_counts( - head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip() - ) - elif args.update: - cmd_update(head, args.base) - else: - cmd_check(head, args.base) + from gate_slot_lock import held_slot + + with held_slot(): + ensure_typecheck_env() + head = count_basedpyright(run_basedpyright()) + if args.emit_counts_dir is not None: + cmd_emit_counts( + head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip() + ) + elif args.update: + cmd_update(head, args.base) + else: + cmd_check(head, args.base) if __name__ == "__main__": diff --git a/scripts/type_discipline_gate.py b/scripts/type_discipline_gate.py index f937283d972..5f6474f20bc 100644 --- a/scripts/type_discipline_gate.py +++ b/scripts/type_discipline_gate.py @@ -267,7 +267,10 @@ def main() -> None: parser.add_argument("--base", default=DEFAULT_BASE) parser.add_argument("--update", action="store_true") args = parser.parse_args() - cmd_update(args.base) if args.update else cmd_check(args.base) + from gate_slot_lock import held_slot + + with held_slot(): + cmd_update(args.base) if args.update else cmd_check(args.base) if __name__ == "__main__": diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/test_litellm/test_gate_slot_lock.py new file mode 100644 index 00000000000..1cf52ae89f6 --- /dev/null +++ b/tests/test_litellm/test_gate_slot_lock.py @@ -0,0 +1,311 @@ +import fcntl +import importlib.util +import os +import signal +import subprocess +import sys +import time +from collections.abc import Callable, Sequence +from contextlib import suppress +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] +HELPER = ROOT / "scripts" / "gate_slot_lock.py" + +_spec = importlib.util.spec_from_file_location("gate_slot_lock", HELPER) +assert _spec is not None and _spec.loader is not None +gate_slot_lock = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(gate_slot_lock) + +START_THEN_WAIT_FOR = ( + "import pathlib, sys, time\n" + "pathlib.Path(sys.argv[1]).touch()\n" + "deadline = time.monotonic() + 20\n" + "while not pathlib.Path(sys.argv[2]).exists():\n" + " if time.monotonic() > deadline:\n" + " sys.exit(3)\n" + " time.sleep(0.05)\n" +) + +TOUCH_TARGET = "import pathlib, sys\npathlib.Path(sys.argv[1]).touch()\n" + +RECORD_INTERVAL = ( + "import sys, time\n" + "with open(sys.argv[1], 'a') as events:\n" + " events.write(f'start {time.monotonic()}\\n')\n" + " events.flush()\n" + " time.sleep(0.6)\n" + " events.write(f'end {time.monotonic()}\\n')\n" + " events.flush()\n" +) + + +def _env(lock_dir: Path, slots: str) -> dict[str, str]: + return { + "PATH": os.environ["PATH"], + "HOME": str(lock_dir.parent), + "LITELLM_GATE_SLOT_DIR": str(lock_dir), + "LITELLM_GATE_SLOTS": slots, + } + + +def _wrapped(payload: Sequence[str]) -> list[str]: + return [sys.executable, str(HELPER), sys.executable, "-c", *payload] + + +def _wait_until(predicate: Callable[[], bool], timeout_seconds: float) -> bool: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return predicate() + + +def _terminate_group(process: subprocess.Popen[bytes]) -> None: + with suppress(ProcessLookupError, PermissionError): + os.killpg(process.pid, signal.SIGKILL) + + +def _reap(process: subprocess.Popen[bytes]) -> None: + with suppress(subprocess.TimeoutExpired): + process.wait(timeout=10) + if process.poll() is None: + process.kill() + process.wait(timeout=10) + + +def test_six_contenders_never_exceed_two_slots_and_all_complete(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + events_file = tmp_path / "events.log" + env = _env(lock_dir, "2") + procs = [ + subprocess.Popen( + _wrapped([RECORD_INTERVAL, str(events_file)]), + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + for _ in range(6) + ] + try: + assert [proc.wait(timeout=60) for proc in procs] == [0] * 6 + finally: + for proc in procs: + if proc.poll() is None: + proc.kill() + proc.wait(timeout=10) + events = sorted( + (float(stamp), 1 if kind == "start" else -1) + for kind, stamp in (line.split() for line in events_file.read_text().splitlines()) + ) + assert len(events) == 12 + concurrency_peaks = [] + running = 0 + for _, delta in events: + running += delta + concurrency_peaks.append(running) + assert max(concurrency_peaks) <= 2 + + +def test_two_slots_admit_two_holders_at_once(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + first_started = tmp_path / "first.started" + second_started = tmp_path / "second.started" + env = _env(lock_dir, "2") + first = subprocess.Popen( + _wrapped([START_THEN_WAIT_FOR, str(first_started), str(second_started)]), env=env + ) + second = subprocess.Popen( + _wrapped([START_THEN_WAIT_FOR, str(second_started), str(first_started)]), env=env + ) + assert first.wait(timeout=30) == 0 + assert second.wait(timeout=30) == 0 + + +def test_contender_beyond_capacity_queues_until_the_slot_frees(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + holder_started = tmp_path / "holder.started" + release = tmp_path / "release" + done = tmp_path / "done" + env = _env(lock_dir, "1") + holder = subprocess.Popen( + _wrapped([START_THEN_WAIT_FOR, str(holder_started), str(release)]), env=env + ) + try: + assert _wait_until(holder_started.exists, 10) + contender = subprocess.Popen( + _wrapped([TOUCH_TARGET, str(done)]), + env=env, + stderr=subprocess.PIPE, + ) + try: + time.sleep(1.5) + assert not done.exists() + release.touch() + assert holder.wait(timeout=10) == 0 + assert contender.wait(timeout=30) == 0 + assert done.exists() + assert contender.stderr is not None + assert b"queueing" in contender.stderr.read() + finally: + release.touch() + _reap(contender) + finally: + release.touch() + _reap(holder) + + +def test_nested_wrapping_reenters_instead_of_deadlocking(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + nested = [ + sys.executable, + str(HELPER), + sys.executable, + str(HELPER), + sys.executable, + "-c", + "print('nested ok')", + ] + proc = subprocess.Popen( + nested, + env=_env(lock_dir, "1"), + stdout=subprocess.PIPE, + start_new_session=True, + ) + try: + stdout, _ = proc.communicate(timeout=20) + except subprocess.TimeoutExpired: + _terminate_group(proc) + pytest.fail("nested gate_slot_lock invocations deadlocked") + assert proc.returncode == 0 + assert b"nested ok" in stdout + + +def test_wrapped_command_exit_code_is_propagated(tmp_path: Path) -> None: + proc = subprocess.run( + [sys.executable, str(HELPER), sys.executable, "-c", "raise SystemExit(7)"], + env=_env(tmp_path / "locks", "2"), + ) + assert proc.returncode == 7 + + +def test_missing_command_exits_127_and_no_command_exits_2(tmp_path: Path) -> None: + env = _env(tmp_path / "locks", "2") + missing = subprocess.run( + [sys.executable, str(HELPER), str(tmp_path / "no-such-binary")], + env=env, + capture_output=True, + ) + assert missing.returncode == 127 + bare = subprocess.run([sys.executable, str(HELPER)], env=env, capture_output=True) + assert bare.returncode == 2 + + +def test_wrapped_command_killed_by_signal_maps_to_128_plus_signal(tmp_path: Path) -> None: + proc = subprocess.run( + _wrapped(["import os, signal\nos.kill(os.getpid(), signal.SIGTERM)\n"]), + env=_env(tmp_path / "locks", "2"), + ) + assert proc.returncode == 128 + signal.SIGTERM + + +def test_unusable_lock_dir_fails_open_and_still_runs_the_command(tmp_path: Path) -> None: + blocker = tmp_path / "blocker" + blocker.write_text("") + done = tmp_path / "done" + proc = subprocess.run( + _wrapped([TOUCH_TARGET, str(done)]), + env=_env(blocker / "locks", "2"), + capture_output=True, + ) + assert proc.returncode == 0 + assert done.exists() + assert b"running unlocked" in proc.stderr + + +def test_zero_slots_disables_locking_entirely(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + done = tmp_path / "done" + proc = subprocess.run( + _wrapped([TOUCH_TARGET, str(done)]), + env=_env(lock_dir, "0"), + ) + assert proc.returncode == 0 + assert done.exists() + assert not lock_dir.exists() + + +def test_non_integer_slot_count_warns_and_falls_back_to_default(tmp_path: Path) -> None: + proc = subprocess.run( + [sys.executable, str(HELPER), sys.executable, "-c", "print('ran')"], + env=_env(tmp_path / "locks", "lots"), + capture_output=True, + ) + assert proc.returncode == 0 + assert b"ran" in proc.stdout + assert b"LITELLM_GATE_SLOTS" in proc.stderr + + +def test_killed_holder_releases_its_slot_for_the_next_contender(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + holder_started = tmp_path / "holder.started" + never = tmp_path / "never" + env = _env(lock_dir, "1") + holder = subprocess.Popen( + _wrapped([START_THEN_WAIT_FOR, str(holder_started), str(never)]), + env=env, + start_new_session=True, + ) + try: + assert _wait_until(holder_started.exists, 10) + finally: + _terminate_group(holder) + holder.wait(timeout=10) + after = subprocess.run( + [sys.executable, str(HELPER), sys.executable, "-c", "print('freed')"], + env=env, + capture_output=True, + timeout=20, + ) + assert after.returncode == 0 + assert b"freed" in after.stdout + + +def test_acquire_slot_holds_marks_and_releases_in_process( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + lock_dir = tmp_path / "locks" + monkeypatch.setenv("LITELLM_GATE_SLOT_HELD", "") + monkeypatch.setenv("LITELLM_GATE_SLOT_DIR", str(lock_dir)) + monkeypatch.setenv("LITELLM_GATE_SLOTS", "1") + handle = gate_slot_lock.acquire_slot() + assert handle is not None + assert os.environ["LITELLM_GATE_SLOT_HELD"] == "1" + assert gate_slot_lock.acquire_slot() is None + with (lock_dir / "slot-0.lock").open("wb") as probe: + with pytest.raises(BlockingIOError): + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + handle.close() + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + fcntl.flock(probe, fcntl.LOCK_UN) + + +def test_held_slot_context_manager_releases_on_exit( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + lock_dir = tmp_path / "locks" + monkeypatch.setenv("LITELLM_GATE_SLOT_HELD", "") + monkeypatch.setenv("LITELLM_GATE_SLOT_DIR", str(lock_dir)) + monkeypatch.setenv("LITELLM_GATE_SLOTS", "1") + with gate_slot_lock.held_slot(): + assert os.environ["LITELLM_GATE_SLOT_HELD"] == "1" + with (lock_dir / "slot-0.lock").open("wb") as probe: + with pytest.raises(BlockingIOError): + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + assert not os.environ.get("LITELLM_GATE_SLOT_HELD") + with (lock_dir / "slot-0.lock").open("wb") as probe: + fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) + fcntl.flock(probe, fcntl.LOCK_UN) diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 33baf7474ce..5ea0e79a196 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -1,4 +1,5 @@ import os +import shutil import signal import subprocess import time @@ -420,6 +421,46 @@ def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log +def test_run_queues_through_the_machine_wide_gate_slot_lock(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + lock_dir = tmp_path / "gate-locks" + proc = _run(repo, bin_dir, {"LITELLM_GATE_SLOT_DIR": str(lock_dir)}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert (lock_dir / "slot-0.lock").exists() + + +def test_run_under_a_held_slot_skips_reacquiring_the_gate_lock(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + lock_dir = tmp_path / "gate-locks" + proc = _run( + repo, + bin_dir, + {"LITELLM_GATE_SLOT_DIR": str(lock_dir), "LITELLM_GATE_SLOT_HELD": "1"}, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert not lock_dir.exists() + + +def test_hook_symlink_install_still_resolves_the_slot_lock_helper(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + scripts_dir = repo / "scripts" + scripts_dir.mkdir() + shutil.copy(SCRIPT, scripts_dir / "pre_commit_lint.sh") + shutil.copy(SCRIPT.parent / "gate_slot_lock.py", scripts_dir / "gate_slot_lock.py") + (repo / ".git" / "hooks" / "pre-commit").symlink_to(Path("../../scripts/pre_commit_lint.sh")) + lock_dir = tmp_path / "gate-locks" + proc = subprocess.run( + ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "hooked"], + cwd=repo, + capture_output=True, + text=True, + env=_env(repo, bin_dir, {"LITELLM_GATE_SLOT_DIR": str(lock_dir)}), + timeout=120, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert (lock_dir / "slot-0.lock").exists() + + def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) proc = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"}) From 87765fcc761e81a1f97c901a7d1f1e365cf7fccf Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 00:29:50 +0000 Subject: [PATCH 045/121] fix(router): merge model_info without new mutable constructions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index baa4724b8eb..ae91de83333 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9079,21 +9079,26 @@ class Router: model_info_name = model model_info: Final = litellm.get_model_info(model=model_info_name) - - ## CHECK USER SET MODEL INFO - raw_user_model_info: Final = deployment.get("model_info") or {} - user_model_info: Final = ( - raw_user_model_info.model_dump(exclude_none=True) - if isinstance(raw_user_model_info, BaseModel) - else {key: value for key, value in raw_user_model_info.items() if value is not None} - ) - if model_info is None: return model_info - # get_model_info() hands back an lru_cache'd dict; merging into a copy keeps - # deployment overrides out of the shared entry - return cast(ModelMapInfo, {**model_info, **user_model_info}) + ## CHECK USER SET MODEL INFO + raw_user_model_info: Final = deployment.get("model_info") + user_model_info: Final = ( + raw_user_model_info.model_dump(exclude_none=True) + if isinstance(raw_user_model_info, BaseModel) + else raw_user_model_info + ) + + # get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset + # values are skipped or Deployment's None pricing defaults would erase the map's + merged_model_info: Final = copy.copy(model_info) + if user_model_info: + for key, value in user_model_info.items(): + if value is not None: + merged_model_info[key] = value + + return merged_model_info def get_model_info(self, id: str) -> dict | None: """ From 2de319575729dd0a77a0b92561b3044d1ed967b7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 14 Aug 2026 17:31:44 -0700 Subject: [PATCH 046/121] refactor(ui): re-sync badge and skeleton onto the base-vega shadcn style components.json has declared "style": "base-vega" since cfe9e39e55, but badge and skeleton were added a few days earlier under new-york and never re-synced, so both still carried the previous style's classes. Badge's destructive variant rendered as solid red with white text instead of the tinted wash the rest of the dashboard uses, which is already the convention for Button Re-runs npx shadcn add for both and keeps the two local deltas the registry cannot supply: cva comes from @/lib/cva.config, since class-variance-authority is not a dependency here, and both stay wrapped in React.forwardRef, which the tripwire in tests/setupTests.ts requires until the React 19 upgrade Adds Badge to ref-forwarding.test.tsx. Nothing covered it before, even though two TooltipTrigger sites compose over it, so the wrapper could have been dropped by the next re-sync without a single test going red Retargets one assertion in LogDetailContent.test.tsx. It regex-matched the whole class string for "destructive" to prove a tag was not alarming red, which the restored aria-invalid classes now satisfy for every variant; it checks the variant attribute and red utility classes instead --- .../src/components/ui/badge.tsx | 35 ++++++++----------- .../src/components/ui/ref-forwarding.test.tsx | 7 ++++ .../src/components/ui/skeleton.tsx | 2 +- .../LogDetailContent.test.tsx | 3 +- 4 files changed, 24 insertions(+), 23 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ui/badge.tsx b/ui/litellm-dashboard/src/components/ui/badge.tsx index f64de004b52..bf17c76aa55 100644 --- a/ui/litellm-dashboard/src/components/ui/badge.tsx +++ b/ui/litellm-dashboard/src/components/ui/badge.tsx @@ -1,22 +1,21 @@ -"use client"; - import * as React from "react"; -import { type VariantProps } from "cva"; +import { mergeProps } from "@base-ui/react/merge-props"; import { useRender } from "@base-ui/react/use-render"; +import { type VariantProps } from "cva"; import { cn, cva } from "@/lib/cva.config"; const badgeVariants = cva({ - base: "inline-flex w-fit shrink-0 items-center justify-center gap-1 overflow-hidden rounded-full border border-transparent px-2 py-0.5 text-xs font-medium whitespace-nowrap transition-[color,box-shadow] focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50 [&>svg]:pointer-events-none [&>svg]:size-3", + base: "group/badge inline-flex h-5 w-fit shrink-0 items-center justify-center gap-1 overflow-hidden rounded-4xl border border-transparent px-2 py-0.5 text-xs font-medium whitespace-nowrap transition-all focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50 has-data-[icon=inline-end]:pr-1.5 has-data-[icon=inline-start]:pl-1.5 aria-invalid:border-destructive aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 [&>svg]:pointer-events-none [&>svg]:size-3!", variants: { variant: { - default: "bg-primary text-primary-foreground [a&]:hover:bg-primary/90", - secondary: "bg-secondary text-secondary-foreground [a&]:hover:bg-secondary/90", + default: "bg-primary text-primary-foreground [a]:hover:bg-primary/80", + secondary: "bg-secondary text-secondary-foreground [a]:hover:bg-secondary/80", destructive: - "bg-destructive text-white focus-visible:ring-destructive/20 dark:bg-destructive/60 dark:focus-visible:ring-destructive/40 [a&]:hover:bg-destructive/90", - outline: "border-border text-foreground [a&]:hover:bg-accent [a&]:hover:text-accent-foreground", - ghost: "[a&]:hover:bg-accent [a&]:hover:text-accent-foreground", - link: "text-primary underline-offset-4 [a&]:hover:underline", + "bg-destructive/10 text-destructive focus-visible:ring-destructive/20 dark:bg-destructive/20 dark:focus-visible:ring-destructive/40 [a]:hover:bg-destructive/20", + outline: "border-border text-foreground [a]:hover:bg-muted [a]:hover:text-muted-foreground", + ghost: "hover:bg-muted hover:text-muted-foreground dark:hover:bg-muted/50", + link: "text-primary underline-offset-4 hover:underline", }, }, defaultVariants: { @@ -24,22 +23,16 @@ const badgeVariants = cva({ }, }); -type BadgeProps = React.ComponentPropsWithoutRef<"span"> & - VariantProps & { - render?: useRender.RenderProp; - }; +type BadgeProps = useRender.ComponentProps<"span"> & VariantProps; const Badge = React.forwardRef( ({ className, variant = "default", render, ...props }, ref) => useRender({ - render: render ?? , + defaultTagName: "span", ref, - props: { - "data-slot": "badge", - "data-variant": variant, - className: cn(badgeVariants({ variant }), className), - ...props, - }, + props: mergeProps<"span">({ className: cn(badgeVariants({ variant }), className) }, props), + render, + state: { slot: "badge", variant }, }), ); Badge.displayName = "Badge"; diff --git a/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx b/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx index dde91a60de3..48b1e8be226 100644 --- a/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx +++ b/ui/litellm-dashboard/src/components/ui/ref-forwarding.test.tsx @@ -2,6 +2,7 @@ import { render } from "@testing-library/react"; import * as React from "react"; import { describe, expect, it } from "vitest"; +import { Badge } from "./badge"; import { Button } from "./button"; import { Card, CardAction, CardContent, CardDescription, CardFooter, CardHeader, CardTitle } from "./card"; import { ChartContainer } from "./chart"; @@ -13,6 +14,12 @@ import { Table, TableBody, TableCaption, TableCell, TableFooter, TableHead, Tabl import { UiLoadingSpinner } from "./ui-loading-spinner"; describe("ui primitives forward refs to their DOM node", () => { + it("Badge", () => { + const ref = React.createRef(); + render(ok); + expect(ref.current).toBeInstanceOf(HTMLSpanElement); + }); + it("Button", () => { const ref = React.createRef(); render(); diff --git a/ui/litellm-dashboard/src/components/ui/skeleton.tsx b/ui/litellm-dashboard/src/components/ui/skeleton.tsx index 69ff4891cec..1104379dde6 100644 --- a/ui/litellm-dashboard/src/components/ui/skeleton.tsx +++ b/ui/litellm-dashboard/src/components/ui/skeleton.tsx @@ -4,7 +4,7 @@ import { cn } from "@/lib/cva.config"; const Skeleton = React.forwardRef>( ({ className, ...props }, ref) => ( -
+
), ); Skeleton.displayName = "Skeleton"; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index 01b040b5c17..74d7760c6fb 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -289,7 +289,8 @@ describe("LogDetailContent", () => { expect(screen.getByText("34,462")).toBeInTheDocument(); expect(screen.getByText("Prompt Cache Creation Tokens")).toBeInTheDocument(); expect(screen.getByText("83")).toBeInTheDocument(); - expect(screen.getByText("Miss").className).not.toMatch(/red|destructive/); + expect(screen.getByText("Miss")).not.toHaveAttribute("data-variant", "destructive"); + expect(screen.getByText("Miss").className).not.toMatch(/\b(bg|text|border)-red/); expect(screen.queryByText("Cache Hit")).not.toBeInTheDocument(); }); From b1696b3edf98614eb2ef61aad1339556c1e1c27d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 14 Aug 2026 17:38:47 -0700 Subject: [PATCH 047/121] test(ui): assert cache and retry tags by text instead of class name Three assertions in LogDetailContent.test.tsx matched a regex against the rendered class string to prove a tag was green or was not red. That pins styling rather than behavior, and jsdom does not resolve the utilities anyway, so the checks only ever proved that a substring survived into the class attribute The badge re-sync exposed it: base-vega's base string carries aria-invalid variants of the destructive token, so a "not destructive" regex started matching every badge regardless of variant Each one now asserts the tag's text is present, which is what the surrounding cases already do and what the user actually observes --- .../LogDetailsDrawer/LogDetailContent.test.tsx | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index 74d7760c6fb..aab8d2b8cb9 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -262,14 +262,14 @@ describe("LogDetailContent", () => { expect(screen.getByText("2 masked")).toBeInTheDocument(); }); - it("should display a green Response Cache 'Hit' tag when the response cache served the request", () => { + it("should display a Response Cache 'Hit' tag when the response cache served the request", () => { render(); expect(screen.getByText("Response Cache")).toBeInTheDocument(); - expect(screen.getByText("Hit").className).toMatch(/green/); + expect(screen.getByText("Hit")).toBeInTheDocument(); }); - it("should show prompt cache tokens without an alarming red tag when only provider prompt caching occurred", () => { + it("should show prompt cache tokens and no response-cache hit when only provider prompt caching occurred", () => { render( { expect(screen.getByText("34,462")).toBeInTheDocument(); expect(screen.getByText("Prompt Cache Creation Tokens")).toBeInTheDocument(); expect(screen.getByText("83")).toBeInTheDocument(); - expect(screen.getByText("Miss")).not.toHaveAttribute("data-variant", "destructive"); - expect(screen.getByText("Miss").className).not.toMatch(/\b(bg|text|border)-red/); + expect(screen.getByText("Miss")).toBeInTheDocument(); expect(screen.queryByText("Cache Hit")).not.toBeInTheDocument(); }); @@ -394,11 +393,10 @@ describe("LogDetailContent", () => { expect(within(retriesItem()).getByText("2 / 3")).toBeInTheDocument(); }); - it("should display a green 'None' tag for Retries when attempted_retries is 0", () => { + it("should display a 'None' tag for Retries when attempted_retries is 0", () => { render(); - const noneTag = within(retriesItem()).getByText("None"); - expect(noneTag.className).toMatch(/green/); + expect(within(retriesItem()).getByText("None")).toBeInTheDocument(); }); it("should display '-' for Retries when attempted_retries is absent from metadata", () => { From 17f5c909f06b16ad554b6aa005458bff2a323b78 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 14 Aug 2026 21:17:14 -0700 Subject: [PATCH 048/121] fix(make): acquire the gate slot before lint setup deps --- Makefile | 9 ++-- tests/test_litellm/test_gate_slot_lock.py | 52 ++++++++++++++++------- 2 files changed, 43 insertions(+), 18 deletions(-) diff --git a/Makefile b/Makefile index 7fe5d1f8045..bdb643e3ec9 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \ test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \ test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \ - info lint lint-dev lint-checks format \ + info lint lint-inner lint-dev lint-checks format \ 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 \ @@ -239,8 +239,11 @@ check-import-safety: $(LINT_DEP_INSTALL) # does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client, # base fetch) runs once up front; the checks themselves are independent, so a sub-make # fans them out with -j and the fast ones finish under basedpyright's shadow. -lint: lint-install lint-fetch-base - $(GATE_SLOT_LOCK) $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks +lint: + @$(GATE_SLOT_LOCK) $(MAKE) lint-inner + +lint-inner: lint-install lint-fetch-base + $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/test_litellm/test_gate_slot_lock.py index 1cf52ae89f6..17fa8547ce7 100644 --- a/tests/test_litellm/test_gate_slot_lock.py +++ b/tests/test_litellm/test_gate_slot_lock.py @@ -115,12 +115,8 @@ def test_two_slots_admit_two_holders_at_once(tmp_path: Path) -> None: first_started = tmp_path / "first.started" second_started = tmp_path / "second.started" env = _env(lock_dir, "2") - first = subprocess.Popen( - _wrapped([START_THEN_WAIT_FOR, str(first_started), str(second_started)]), env=env - ) - second = subprocess.Popen( - _wrapped([START_THEN_WAIT_FOR, str(second_started), str(first_started)]), env=env - ) + first = subprocess.Popen(_wrapped([START_THEN_WAIT_FOR, str(first_started), str(second_started)]), env=env) + second = subprocess.Popen(_wrapped([START_THEN_WAIT_FOR, str(second_started), str(first_started)]), env=env) assert first.wait(timeout=30) == 0 assert second.wait(timeout=30) == 0 @@ -131,9 +127,7 @@ def test_contender_beyond_capacity_queues_until_the_slot_frees(tmp_path: Path) - release = tmp_path / "release" done = tmp_path / "done" env = _env(lock_dir, "1") - holder = subprocess.Popen( - _wrapped([START_THEN_WAIT_FOR, str(holder_started), str(release)]), env=env - ) + holder = subprocess.Popen(_wrapped([START_THEN_WAIT_FOR, str(holder_started), str(release)]), env=env) try: assert _wait_until(holder_started.exists, 10) contender = subprocess.Popen( @@ -274,9 +268,7 @@ def test_killed_holder_releases_its_slot_for_the_next_contender(tmp_path: Path) assert b"freed" in after.stdout -def test_acquire_slot_holds_marks_and_releases_in_process( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_acquire_slot_holds_marks_and_releases_in_process(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: lock_dir = tmp_path / "locks" monkeypatch.setenv("LITELLM_GATE_SLOT_HELD", "") monkeypatch.setenv("LITELLM_GATE_SLOT_DIR", str(lock_dir)) @@ -293,9 +285,7 @@ def test_acquire_slot_holds_marks_and_releases_in_process( fcntl.flock(probe, fcntl.LOCK_UN) -def test_held_slot_context_manager_releases_on_exit( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: +def test_held_slot_context_manager_releases_on_exit(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: lock_dir = tmp_path / "locks" monkeypatch.setenv("LITELLM_GATE_SLOT_HELD", "") monkeypatch.setenv("LITELLM_GATE_SLOT_DIR", str(lock_dir)) @@ -309,3 +299,35 @@ def test_held_slot_context_manager_releases_on_exit( with (lock_dir / "slot-0.lock").open("wb") as probe: fcntl.flock(probe, fcntl.LOCK_EX | fcntl.LOCK_NB) fcntl.flock(probe, fcntl.LOCK_UN) + + +def _make_rule(target: str) -> tuple[list[str], list[str]]: + database = subprocess.run( + ["make", "--dry-run", "--print-data-base", "info"], + cwd=ROOT, + capture_output=True, + text=True, + check=True, + ).stdout + lines = database.splitlines() + for index, line in enumerate(lines): + if line != f"{target}:" and not line.startswith(f"{target}: "): + continue + recipe: list[str] = [] + for follower in lines[index + 1 :]: + if follower.startswith("#"): + continue + if not follower.startswith("\t"): + break + recipe.append(follower.strip()) + return line.split(":", 1)[1].split(), recipe + raise AssertionError(f"target {target} not found in make database") + + +def test_direct_make_lint_takes_a_slot_before_any_setup() -> None: + lint_prerequisites, lint_recipe = _make_rule("lint") + assert lint_prerequisites == [] + assert any("$(GATE_SLOT_LOCK)" in line for line in lint_recipe) + inner_prerequisites, _ = _make_rule("lint-inner") + assert "lint-install" in inner_prerequisites + assert "lint-fetch-base" in inner_prerequisites From de25e199d446d12f1995235f739a6172a51ce51a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 15 Aug 2026 00:33:20 -0700 Subject: [PATCH 049/121] fix(ui): de-duplicate the reset budget option and polish shadcn surfaces Create Key offered two options labelled "Never resets" in the Reset Budget dropdown. BudgetDurationDropdown renders its unset item using the caller's placeholder, and create_key_button passed placeholder="Never resets" alongside showNeverResets, so the omit option and the explicit-null option looked identical while behaving differently. An omitted budget_duration picks up default_key_generate_params and the linked budget tier's schedule, whereas the "none" sentinel is converted to an explicit null and truly never resets. The unset item now reads "Not set", matching getBudgetDurationLabel, and the create-key test mock passes the placeholder through so a future collision fails the suite The rest is migration cleanup found during manual QA. The models and endpoints tab strip hides its scrollbar and fades at the right edge using a vendored copy of the shadcn scroll-fade utility, keeping the CLI package out of the build. globals.css neutralises the @tailwindcss/forms resting-state rules for combobox-chip-input, which lets twelve call sites drop the same copy-pasted className workaround. Guardrails moves to the line tab variant and stops clipping its textarea focus ring, the log drawer JSON tree takes the app background, the audit log empty state centres, the caching page selects no longer stretch to the row height, the usage page team filter shares its row with the Export button through a new filterSlot prop, and the cost optimization and vector store tab strips drop their leftover full-width divider --- ui/litellm-dashboard/eslint-suppressions.json | 5 ++ .../caching/_components/cache_dashboard.tsx | 6 +- .../_components/CostOptimizationView.tsx | 2 +- .../_components/GuardrailTestPanel.tsx | 2 +- .../_components/GuardrailsPanel.tsx | 2 +- .../custom_code/CustomCodeModal.tsx | 5 +- .../guardrails/_components/pii_components.tsx | 1 - .../(dashboard)/models-and-endpoints/page.tsx | 2 +- .../old-usage/_components/usage.tsx | 2 +- .../components/chat_ui/ChatComposer.tsx | 7 +- .../src/app/(dashboard)/playground/page.tsx | 13 ++- .../components/EntityUsage/EntityUsage.tsx | 11 +-- .../vector-stores/_components/index.tsx | 2 +- ui/litellm-dashboard/src/app/globals.css | 72 +++++++++++++++ .../EntityUsageExport/UsageExportHeader.tsx | 87 ++++++++++--------- .../components/ModelSelect/ModelSelect.tsx | 6 +- .../CommunityEngagementButtons.tsx | 80 +++++++++-------- .../src/components/TeamSSOSettings.tsx | 6 +- .../common_components/team_multi_select.tsx | 7 +- .../organisms/create_key_button.test.tsx | 25 +++++- .../organisms/create_key_button.tsx | 2 +- .../src/components/public_model_hub.tsx | 2 +- .../search_tools/SearchToolSelector.tsx | 7 +- .../components/shared/DataTable/DataTable.tsx | 5 +- .../src/components/shared/MultiSelect.tsx | 2 +- .../src/components/ui/button-group.tsx | 76 ++++++++++++++++ .../src/components/user_agent_activity.tsx | 6 +- .../view_logs/LogDetailsDrawer/JsonViewer.tsx | 6 +- 28 files changed, 295 insertions(+), 154 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/ui/button-group.tsx diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 2b903f763d9..4a6d7006893 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2974,6 +2974,11 @@ "count": 1 } }, + "src/components/ui/button-group.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/ui/button.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx index 6759a9af63f..c1a4968f58b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx @@ -190,7 +190,7 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole Metrics" on the Usage page or individual requests in the Logs page.

-
+
= ({ accessToken, token, userRole )) } - + No virtual keys found @@ -237,7 +237,7 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole )) } - + No models found diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx index 702bb5b8034..8122cfa7a9c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx @@ -63,7 +63,7 @@ const CostOptimizationView: React.FC = ({ accessToken
- + Overall diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx index 8297c1e1e3b..5c79a2292a2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx @@ -131,7 +131,7 @@ export function GuardrailTestPanel({
{/* Input Section */} -
+
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index 9df2fa04c04..80267f97a94 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -126,7 +126,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole return (
- + {isAdmin && ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx index 1497c789cb4..5015884de31 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx @@ -532,10 +532,7 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS {option.label} ))} - + No matching modes diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx index 6994f43772c..5f8e833af8d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx @@ -59,7 +59,6 @@ export const CategoryFilter: React.FC = ({ categories, sele ))} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 998f65629e0..94737d88d0f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -175,7 +175,7 @@ export default function ModelsAndEndpointsPage() { ) : (
-
+
{visibleSlugs.map((slug) => { const key = slug || BASE_TAB_KEY; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx index 7ddf2622c7a..a0835729f1b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx @@ -891,7 +891,7 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use )) } - + No tags found diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx index cc541c939ad..934e502492e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx @@ -54,15 +54,12 @@ export function ChatComposer({ return (
{showSuggestions && suggestions.length > 0 && ( -
+
{suggestions.map((suggestion) => ( +
) : null, })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx index f37acb3d85a..8f51177bafe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.tsx @@ -7,7 +7,7 @@ import { PageHeader } from "@/components/shared/PageHeader"; import { Button } from "@/components/ui/button"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; -import { AccessGroupCreateModal } from "./AccessGroupsModal/AccessGroupCreateModal"; +import { AccessGroupCreateDialog } from "./access-group-create/AccessGroupCreateDialog"; import { AccessGroupsTable } from "./AccessGroupsTable"; import { AccessGroup } from "./types"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -104,7 +104,7 @@ export function AccessGroupsPage() { onDeleteClick={setGroupToDelete} /> - setIsCreateModalVisible(false)} /> + ({ + __esModule: true, + default: { success: vi.fn(), fromBackend: vi.fn() }, +})); +vi.mock("@/components/ModelSelect/ModelSelect", () => ({ + ModelSelect: ({ onChange }: { onChange: (values: string[]) => void }) => ( + + ), +})); +vi.mock("@/app/(dashboard)/hooks/agents/useAgents", () => ({ + useAgents: () => ({ data: { agents: [{ agent_id: "agent-1", agent_name: "Support Agent" }] } }), +})); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "GitHub MCP" }] }), +})); + +import { AccessGroupCreateDialog } from "./AccessGroupCreateDialog"; + +const Harness = ({ createAccessGroup }: { createAccessGroup: (body: unknown) => Promise }) => { + const [open, setOpen] = React.useState(true); + return ( + <> + + + + ); +}; + +const renderDialog = (overrides?: { createAccessGroup?: ReturnType }) => { + const createAccessGroup = overrides?.createAccessGroup ?? vi.fn().mockResolvedValue({}); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + render( + + + , + ); + return { createAccessGroup }; +}; + +describe("AccessGroupCreateDialog", () => { + it("blocks submit and shows an error when the name is missing", async () => { + const user = userEvent.setup(); + const { createAccessGroup } = renderDialog(); + + await user.click(screen.getByRole("button", { name: "Create Group" })); + + expect(await screen.findByRole("alert")).toHaveTextContent("Please enter the access group name"); + expect(createAccessGroup).not.toHaveBeenCalled(); + }); + + it("returns to the General Info tab when submitting an invalid form from another tab", async () => { + const user = userEvent.setup(); + const { createAccessGroup } = renderDialog(); + + await user.click(screen.getByRole("tab", { name: "Models" })); + await waitFor(() => expect(screen.queryByLabelText("Group Name")).not.toBeInTheDocument()); + + await user.click(screen.getByRole("button", { name: "Create Group" })); + + expect(await screen.findByLabelText("Group Name")).toBeInTheDocument(); + expect(await screen.findByRole("alert")).toHaveTextContent("Please enter the access group name"); + expect(createAccessGroup).not.toHaveBeenCalled(); + }); + + it("sends only the group name for a minimal create and closes the dialog", async () => { + const user = userEvent.setup(); + const { createAccessGroup } = renderDialog(); + + await user.type(screen.getByLabelText("Group Name"), "prod-models"); + await user.click(screen.getByRole("button", { name: "Create Group" })); + + await waitFor(() => expect(createAccessGroup).toHaveBeenCalledTimes(1)); + expect(createAccessGroup.mock.calls[0][0]).toStrictEqual({ access_group_name: "prod-models" }); + await waitFor(() => expect(screen.queryByLabelText("Group Name")).not.toBeInTheDocument()); + }); + + it("maps the description and model selections into the create body", async () => { + const user = userEvent.setup(); + const { createAccessGroup } = renderDialog(); + + await user.type(screen.getByLabelText("Group Name"), "prod-models"); + await user.type(screen.getByLabelText("Description"), "engineering access"); + await user.click(screen.getByRole("tab", { name: "Models" })); + await user.click(screen.getByRole("button", { name: "set-models" })); + await user.click(screen.getByRole("button", { name: "Create Group" })); + + await waitFor(() => expect(createAccessGroup).toHaveBeenCalledTimes(1)); + expect(createAccessGroup.mock.calls[0][0]).toStrictEqual({ + access_group_name: "prod-models", + description: "engineering access", + access_model_names: ["gpt-5.2"], + }); + }); + + it("keeps the dialog open with the entered values when the create fails", async () => { + const user = userEvent.setup(); + const { createAccessGroup } = renderDialog({ + createAccessGroup: vi.fn().mockRejectedValue(new Error("boom")), + }); + + await user.type(screen.getByLabelText("Group Name"), "prod-models"); + await user.click(screen.getByRole("button", { name: "Create Group" })); + + await waitFor(() => expect(createAccessGroup).toHaveBeenCalledTimes(1)); + expect(screen.getByLabelText("Group Name")).toHaveValue("prod-models"); + }); + + it("resets the form when the dialog is cancelled and reopened", async () => { + const user = userEvent.setup(); + renderDialog(); + + await user.type(screen.getByLabelText("Group Name"), "abandoned"); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(screen.queryByLabelText("Group Name")).not.toBeInTheDocument()); + + await user.click(screen.getByRole("button", { name: "reopen" })); + expect(screen.getByLabelText("Group Name")).toHaveValue(""); + }); + + it("resets the form when the dialog is dismissed with Escape and reopened", async () => { + const user = userEvent.setup(); + renderDialog(); + + await user.type(screen.getByLabelText("Group Name"), "abandoned"); + await user.keyboard("{Escape}"); + await waitFor(() => expect(screen.queryByLabelText("Group Name")).not.toBeInTheDocument()); + + await user.click(screen.getByRole("button", { name: "reopen" })); + expect(screen.getByLabelText("Group Name")).toHaveValue(""); + }); + + it("cannot be dismissed while a create is pending, then closes once on success", async () => { + const user = userEvent.setup(); + let resolveCreate: (value: unknown) => void = () => {}; + const createAccessGroup = vi.fn().mockImplementation( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + renderDialog({ createAccessGroup }); + + await user.type(screen.getByLabelText("Group Name"), "prod-models"); + await user.keyboard("{Enter}"); + await waitFor(() => expect(createAccessGroup).toHaveBeenCalledTimes(1)); + + await user.keyboard("{Escape}"); + expect(screen.getByLabelText("Group Name")).toHaveValue("prod-models"); + + resolveCreate({}); + await waitFor(() => expect(screen.queryByLabelText("Group Name")).not.toBeInTheDocument()); + }); + + it("does not fire a second create while one is pending", async () => { + const user = userEvent.setup(); + let resolveCreate: (value: unknown) => void = () => {}; + const createAccessGroup = vi.fn().mockImplementation( + () => + new Promise((resolve) => { + resolveCreate = resolve; + }), + ); + renderDialog({ createAccessGroup }); + + await user.type(screen.getByLabelText("Group Name"), "prod-models"); + await user.keyboard("{Enter}"); + await waitFor(() => expect(createAccessGroup).toHaveBeenCalledTimes(1)); + await user.keyboard("{Enter}"); + + expect(createAccessGroup).toHaveBeenCalledTimes(1); + resolveCreate({}); + await waitFor(() => expect(screen.queryByLabelText("Group Name")).not.toBeInTheDocument()); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx new file mode 100644 index 00000000000..3f2205b4206 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx @@ -0,0 +1,244 @@ +"use client"; + +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { BotIcon, InfoIcon, LayersIcon, ServerIcon } from "lucide-react"; +import * as React from "react"; + +import { accessGroupKeys } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { FieldGroup } from "@/components/shared/form/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { Textarea } from "@/components/ui/textarea"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { fetchClient } from "@/lib/http/api"; + +import { buildAccessGroupCreateBody, emptyAccessGroupFormValues, type AccessGroupCreateBody } from "./mapper"; +import { accessGroupCreateSchema } from "./schema"; + +const GENERAL_TAB = "general"; + +interface MultiSelectOption { + value: string; + label: string; +} + +interface MultiSelectProps { + id: string; + value: string[]; + onChange: (value: string[]) => void; + options: MultiSelectOption[]; + placeholder: string; + "aria-invalid": true | undefined; + "aria-describedby": string | undefined; +} + +const MultiSelect = ({ + id, + value, + onChange, + options, + placeholder, + "aria-invalid": ariaInvalid, + "aria-describedby": ariaDescribedBy, +}: MultiSelectProps) => ( + +); + +const defaultCreateAccessGroup = async (body: AccessGroupCreateBody): Promise => { + const { data } = await fetchClient.POST("/v1/access_group", { body }); + return data; +}; + +interface AccessGroupCreateDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; + createAccessGroup?: (body: AccessGroupCreateBody) => Promise; +} + +export const AccessGroupCreateDialog = ({ + open, + onOpenChange, + createAccessGroup = defaultCreateAccessGroup, +}: AccessGroupCreateDialogProps) => { + const queryClient = useQueryClient(); + const form = useZodForm(accessGroupCreateSchema, { defaultValues: emptyAccessGroupFormValues }); + const [activeTab, setActiveTab] = React.useState(GENERAL_TAB); + + const { data: agentsData } = useAgents(); + const { data: mcpServersData } = useMCPServers(); + + const mcpServerOptions = (mcpServersData ?? []).map((server) => ({ + value: server.server_id, + label: server.server_name ?? server.server_id, + })); + const agentOptions = (agentsData?.agents ?? []).map((agent) => ({ + value: agent.agent_id, + label: agent.agent_name, + })); + + const closeAndReset = () => { + form.reset(emptyAccessGroupFormValues); + setActiveTab(GENERAL_TAB); + onOpenChange(false); + }; + + const mutation = useMutation({ + mutationFn: (body: AccessGroupCreateBody) => createAccessGroup(body), + onSuccess: () => { + NotificationsManager.success("Access group created successfully"); + queryClient.invalidateQueries({ queryKey: accessGroupKeys.all }); + closeAndReset(); + }, + onError: (error: unknown) => + NotificationsManager.fromBackend(error instanceof Error ? error.message : "Failed to create access group"), + }); + + const handleOpenChange = (nextOpen: boolean) => { + if (!nextOpen && mutation.isPending) return; + if (!nextOpen) { + form.reset(emptyAccessGroupFormValues); + setActiveTab(GENERAL_TAB); + } + onOpenChange(nextOpen); + }; + + const onSubmit = form.handleSubmit( + (values) => { + if (mutation.isPending) return; + mutation.mutate(buildAccessGroupCreateBody(values)); + }, + // the only validated field (name) lives on the General Info tab + () => setActiveTab(GENERAL_TAB), + ); + + return ( + + + + Create Access Group + + +
+ + + + + General Info + + + + Models + + + + MCP Servers + + + + Agents + + + + + + + {({ ref, ...field }) => } + + + {({ ref, ...field }) => ( +