From b80246971b80369acdc2d492a396bf540ad1f1e2 Mon Sep 17 00:00:00 2001 From: stuxf <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 17:36:12 -0700 Subject: [PATCH] fix(batches): count non-chat tokens, validate batch-file model access (VERIA-39) (#27015) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(batches): count non-chat tokens and validate every model in batch file Two security control bypasses on POST /v1/batches: 1. `_get_batch_job_input_file_usage` only summed tokens for `body.messages` (chat completions). Embedding (`input`) and text completion (`prompt`) batches reported zero, letting massive non-chat workloads slip past TPM rate limits. Extend the counter to handle string and list shapes for both fields. 2. The batch input file was forwarded to the upstream provider without inspecting the models named inside the JSONL — only the outer `model` query parameter was checked against the caller's allowlist. A caller restricted to gpt-3.5 could submit a batch targeting gpt-4o and the upstream would execute it under the proxy's shared API key. Add `_get_models_from_batch_input_file_content` (returns the distinct `body.model` values) and call it from `_enforce_batch_file_model_access` in the pre-call hook, which runs each model through `can_key_call_model` so the same allowlist semantics (wildcards, access groups, all-proxy-models, team aliases) the proxy enforces on `/chat/completions` apply here too. Any unauthorized model raises a 403 before the file is forwarded. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(batches): count pre-tokenized prompt/input shapes, classify 403 logs Two follow-ups from the Greptile review on the batch validation PR: 1. P1 TPM bypass via integer token arrays. The OpenAI batch schema accepts ``prompt`` and ``input`` as ``list[int]`` (a single pre-tokenized prompt) or ``list[list[int]]`` (multiple) in addition to the string and ``list[str]`` shapes. Pre-fix only the string shapes were counted, so a caller could submit a batch with hundreds of millions of pre-tokenized tokens and the rate limiter would record zero. Extract the per-field logic into ``_count_prompt_or_input_tokens`` and count each int as one token. 2. P2 access-denial logs were indistinguishable from I/O failures. ``count_input_file_usage`` caught every exception under a generic "Error counting input file usage" message, so an intentional 403 from ``_enforce_batch_file_model_access`` looked the same in the logs as a missing file or a Prisma timeout. Catch ``HTTPException`` separately and log 403s at WARNING level with a security-relevant message before re-raising. Tests cover the new shapes: single ``list[int]``, ``list[list[int]]`` (the worst-case bypass vector), and embeddings ``input`` with pre-tokenized arrays. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: Claude Opus 4.7 (1M context) --- litellm/batches/batch_utils.py | 78 ++++- litellm/proxy/hooks/batch_rate_limiter.py | 69 +++++ .../proxy/hooks/test_batch_file_validation.py | 285 ++++++++++++++++++ 3 files changed, 429 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_batch_file_validation.py diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4b965d4e63..aaf083e75d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -387,6 +387,27 @@ def _get_batch_job_total_usage_from_file_content( ) +def _get_models_from_batch_input_file_content( + file_content_dictionary: List[dict], +) -> List[str]: + """Extract the distinct ``body.model`` values from a batch *input* file. + + Used by the proxy's batch pre-call hook to enforce that the caller is + authorized for every model named inside the JSONL — not just the one + on the outer request — so the proxy's per-key model allowlist isn't + bypassed by smuggling expensive models into the batch file. + """ + models: List[str] = [] + seen: set = set() + for _item in file_content_dictionary: + body = _item.get("body") or {} + model = body.get("model") + if model and model not in seen: + seen.add(model) + models.append(model) + return models + + def _get_batch_job_input_file_usage( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", @@ -403,11 +424,25 @@ def _get_batch_job_input_file_usage( for _item in file_content_dictionary: body = _item.get("body", {}) model = body.get("model", model_name or "") - messages = body.get("messages", []) + # Chat completion payloads. + messages = body.get("messages") if messages: - item_tokens = token_counter(model=model, messages=messages) - prompt_tokens += item_tokens + prompt_tokens += token_counter(model=model, messages=messages) + continue + + # Text completion payloads (`prompt`). + prompt = body.get("prompt") + if prompt: + prompt_tokens += _count_prompt_or_input_tokens(model=model, value=prompt) + continue + + # Embedding payloads (`input`). + input_data = body.get("input") + if input_data: + prompt_tokens += _count_prompt_or_input_tokens( + model=model, value=input_data + ) return Usage( total_tokens=prompt_tokens + completion_tokens, @@ -416,6 +451,43 @@ def _get_batch_job_input_file_usage( ) +def _count_prompt_or_input_tokens(model: str, value: Any) -> int: + """Token-count a ``prompt`` / ``input`` field that the OpenAI batch + schema allows in four shapes: + + - ``str``: a single text prompt. + - ``list[str]``: multiple text prompts. + - ``list[int]``: a pre-tokenized prompt (each int counts as 1 token). + - ``list[list[int]]``: multiple pre-tokenized prompts. + + Pre-fix only the string shapes were counted, so a caller could send + a large ``list[list[int]]`` payload and slip past TPM rate limits + with a recorded cost of zero tokens. + """ + if isinstance(value, str): + return token_counter(model=model, text=value) + if isinstance(value, list): + total = 0 + for chunk in value: + if isinstance(chunk, str): + total += token_counter(model=model, text=chunk) + elif isinstance(chunk, int): + # Single pre-tokenized prompt at the top level: each + # int counts as one token. + total += 1 + elif isinstance(chunk, list): + # Nested pre-tokenized prompt: every int contributes a + # token. Mixed string/int items still count. + total += sum(1 if isinstance(t, int) else 0 for t in chunk) + total += sum( + token_counter(model=model, text=t) + for t in chunk + if isinstance(t, str) + ) + return total + return 0 + + def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage: """ Get the tokens of a batch job from the response body diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 2ee0588f19..f740d5dd40 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import ( _get_batch_job_input_file_usage, _get_file_content_as_dictionary, + _get_models_from_batch_input_file_content, ) from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -246,6 +247,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content_as_dict = _get_file_content_as_dictionary(file_content.content) + # Validate every model named in the batch JSONL against the + # caller's per-key model allowlist. Without this, a caller + # could smuggle restricted/expensive models inside the file + # and the upstream provider would execute the batch under + # the proxy's shared API key. + if user_api_key_dict is not None: + await self._enforce_batch_file_model_access( + user_api_key_dict=user_api_key_dict, + file_content_as_dict=file_content_as_dict, + ) + input_file_usage = _get_batch_job_input_file_usage( file_content_dictionary=file_content_as_dict, custom_llm_provider=custom_llm_provider, @@ -256,12 +268,69 @@ class _PROXY_BatchRateLimiter(CustomLogger): request_count=request_count, ) + except HTTPException as e: + # Distinguish intentional 403s from `_enforce_batch_file_model_access` + # from genuine I/O failures so security-relevant rejections show up + # in the access log instead of getting buried in error noise. + if e.status_code == 403: + verbose_proxy_logger.warning( + f"Batch rejected: caller not authorized for a model named in {file_id}: {e.detail}" + ) + else: + verbose_proxy_logger.error( + f"Batch input file rejected for {file_id}: status={e.status_code} detail={e.detail}" + ) + raise except Exception as e: verbose_proxy_logger.error( f"Error counting input file usage for {file_id}: {str(e)}" ) raise + async def _enforce_batch_file_model_access( + self, + user_api_key_dict: UserAPIKeyAuth, + file_content_as_dict: List[dict], + ) -> None: + """Reject the batch if the caller is not authorized for every + ``body.model`` named inside the JSONL. + + Reuses ``can_key_call_model`` so the same allowlist semantics + (wildcards, access groups, ``all-proxy-models``, team aliases) + the proxy enforces on `/chat/completions` apply here. + """ + from litellm.proxy.auth.auth_checks import can_key_call_model + from litellm.proxy.proxy_server import llm_router + + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return + + llm_model_list = llm_router.model_list if llm_router is not None else None + for model in models: + try: + await can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + except HTTPException: + raise + except Exception as e: + # `can_key_call_model` raises ProxyException on denial; + # re-shape to a 403 so the batch endpoint returns a + # consistent rejection without leaking internal types. + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Batch input file references a model the caller is " + f"not authorized to use: model={model}, reason={str(e)}" + ) + }, + ) + async def _fetch_managed_file_content( self, file_id: str, diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py new file mode 100644 index 0000000000..7f1006543b --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -0,0 +1,285 @@ +""" +VERIA-39 regression tests: + +- The batch input-file token counter must measure embeddings (`input`) + and text-completion (`prompt`) payloads, not only chat (`messages`). +- The batch rate-limiter pre-call hook must reject batch files that name + models the caller is not authorized to use. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +# --------------------------------------------------------------------------- +# Token counter — covers all three batch payload shapes +# --------------------------------------------------------------------------- + + +def test_token_counter_counts_chat_messages(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_text_completion_prompt(): + """Pre-fix this returned 0 tokens (the function only inspected + `messages`), letting `prompt`-style batches slip past TPM limits.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_embedding_input_string(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + {"body": {"model": "text-embedding-3-small", "input": "hello world"}} + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_embedding_input_list(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "text-embedding-3-small", + "input": ["hello", "world"], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_text_completion_prompt_list(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": ["alpha", "beta"], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_pre_tokenized_prompt_int_list(): + """OpenAI's text-completion API accepts a single pre-tokenized prompt as + a list of ints. Each int is one token; pre-fix this shape was silently + counted as zero, leaving a TPM bypass.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [1, 2, 3, 4, 5], + } + } + ] + ) + assert usage.prompt_tokens == 5 + + +def test_token_counter_counts_pre_tokenized_prompt_list_of_int_lists(): + """Multiple pre-tokenized prompts (`list[list[int]]`) — the most + important bypass shape. A 1000-token batch must report 1000 tokens, + not zero.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [[1] * 250, [2] * 250, [3] * 500], + } + } + ] + ) + assert usage.prompt_tokens == 1000 + + +def test_token_counter_counts_pre_tokenized_input_for_embeddings(): + """Same shape applies to embeddings (`input`).""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "text-embedding-3-small", + "input": [[1, 2, 3], [4, 5, 6]], + } + } + ] + ) + assert usage.prompt_tokens == 6 + + +# --------------------------------------------------------------------------- +# Model extractor +# --------------------------------------------------------------------------- + + +def test_model_extractor_returns_distinct_models(): + from litellm.batches.batch_utils import _get_models_from_batch_input_file_content + + models = _get_models_from_batch_input_file_content( + [ + {"body": {"model": "gpt-4o", "messages": []}}, + {"body": {"model": "gpt-4o", "messages": []}}, # duplicate + {"body": {"model": "gpt-4o-mini", "messages": []}}, + {"body": {}}, # missing model + ] + ) + assert models == ["gpt-4o", "gpt-4o-mini"] + + +# --------------------------------------------------------------------------- +# Pre-call hook model validation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pre_call_rejects_unauthorized_model_in_batch_file(): + """Pre-fix the hook only validated the outer `model` parameter and + forwarded the file as-is. With this fix, a model named inside the + JSONL that the caller cannot use must trigger a 403.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + # Simulated decoded batch file: caller is restricted to gpt-3.5 + # but the JSONL points at gpt-4o. + file_dict = [ + {"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "x"}]}} + ] + + user = UserAPIKeyAuth( + api_key="sk-restricted", + user_id="alice", + models=["gpt-3.5-turbo"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + # `can_key_call_model` raises a ProxyException for non-allowed models. + async def _raise_unauthorized(**kwargs): + raise Exception( + f"Key not allowed to access model. This key only has access to models={kwargs['valid_token'].models}" + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=AsyncMock(side_effect=_raise_unauthorized), + ), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + ) + + assert exc.value.status_code == 403 + assert "gpt-4o" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_pre_call_allows_authorized_model_in_batch_file(): + """If every model in the JSONL is on the caller's allowlist, the hook + must not raise.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + file_dict = [ + { + "body": { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "x"}], + } + } + ] + + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=["gpt-3.5-turbo"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=AsyncMock(return_value=True), + ), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + # Should not raise + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + ) + + +@pytest.mark.asyncio +async def test_pre_call_skips_check_when_no_models_present(): + """Files without any `body.model` (corrupt or empty) must not 500; + the rate limiter logs a warning elsewhere and proceeds.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice") + + # Should not raise even though `can_key_call_model` is the default + # (would fail). The early-return on empty models keeps the call out + # entirely. + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=[], + ) + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=[{"body": {}}], + )