litellm/tests/test_litellm/proxy/hooks
stuxf b80246971b
fix(batches): count non-chat tokens, validate batch-file model access (VERIA-39) (#27015)
* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-01 17:36:12 -07:00
..
test_async_post_call_streaming_iterator_hook.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_batch_file_validation.py fix(batches): count non-chat tokens, validate batch-file model access (VERIA-39) (#27015) 2026-05-01 17:36:12 -07:00
test_dynamic_rate_limiter_v3.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_image_generation_guardrails.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_key_management_event_hooks.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_max_budget_limiter.py [Fix] Proxy: Skip Personal Budget Hook When Reservation Covers Counter 2026-05-01 15:57:42 -07:00
test_max_budget_per_session_limiter.py Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849) 2026-03-07 19:12:42 -08:00
test_max_iterations_limiter.py Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849) 2026-03-07 19:12:42 -08:00
test_parallel_request_limiter_v3.py fix(proxy): include team membership budget in combined_view for RPM/TPM (#24925) 2026-05-01 17:26:45 +05:30
test_post_call_failure_hook_integration.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_post_call_response_headers_hook.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_post_call_streaming_hook_integration.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_post_call_success_hook_integration.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_prompt_injection_detection.py fix: prompt injection not working (#16701) 2025-11-17 20:04:57 -08:00
test_proxy_track_cost_callback.py address budget reservation review edges 2026-04-30 21:21:26 -07:00
test_rate_limiter_toctou.py fix(rate-limit): fail closed on unrecognized OVER_LIMIT descriptor 2026-05-01 12:19:43 -07:00
test_send_invite_email.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00