litellm/tests/test_litellm/proxy
mubashir1osmani 5339c2d783
fix(vertex/files): stream OpenAI->Vertex batch JSONL uploads (#31036)
* fix(vertex/files): stream OpenAI->Vertex batch JSONL uploads to fix OOM on large files

Large (1GB+) batch JSONL uploads to Vertex AI / GCS caused OOM or killed the worker
because the request body was buffered and multiplied 2-3x in size. The create-file
path is now streaming end-to-end: transform_create_file_request returns a
ResumableChunkedUploadConfig carrying a lazy _OpenAIToVertexBatchUploadStream, and the
HTTP handler opens a GCS resumable session and PUTs the body in bounded 8 MiB chunks
(Content-Range, 308 between chunks) so the transformed payload is never held in full.
The proxy /v1/files endpoint streams from Starlette's spooled upload handle instead of
reading the whole body, and batch rate limiting counts tokens and models in a single
streaming pass.

Only gcs_bucket_name is supported for the GCS target; the legacy bucket_name key is
intentionally not read.

Also removes the unreachable VertexAIFilesHandler create path and everything only it
kept alive (VertexAIJsonlFilesTransformation, _stream_openai_jsonl_to_vertex, the legacy
transform helpers), plus the orphaned batch_utils helpers the streaming rewrite replaced.

* fix(batches): return original JSONL on unparseable row to avoid silent batch truncation

The streaming rewrite of replace_model_in_jsonl accumulated physical lines and
skipped a row on JSONDecodeError to support multi-line objects, but a genuinely
malformed or truncated row never completes: it poisons the buffer, swallows every
following row, and the function still returned the partial rewrite (the rows before
the bad one, already model-rewritten) as if the batch were complete. That turned the
pre-rewrite behavior of returning the original file unchanged (so the provider rejects
the bad batch loudly) into a silent partial submission.

Restore the original-content fallback: when an unparseable remainder is left after the
loop, return the original file_content (rewinding a consumed seekable source) instead of
the truncated output. The multi-line happy path is unchanged.

* test(batches): mock resumable GCS upload in vertex batch prediction test

The vertex batch file-create path now streams to a GCS resumable session via
_aresumable_chunked_upload (httpx send) instead of AsyncHTTPHandler.post, so the
existing test's post mock no longer intercepted the upload and a real request hit
GCS (401). Mock _aresumable_chunked_upload to return the GCS object response; the
resumable protocol itself is covered in test_vertex_ai_files_streaming.py.

* fix(batches): resilient per-row token accounting; no hard-block on count failure

The batch input-file pass iterated a generator whose json.loads raised on a
malformed line; the outer except caught it and stopped the loop, so any body.model
on rows after a bad line was never collected and the model allowlist check ran
against a partial set. It also hard-blocked the batch with a 400 whenever token
counting raised, a backwards-incompatible change from the prior swallow-and-proceed
behavior that breaks legitimate rows the token counter cannot measure (e.g. some
multimodal content).

Iterate the JSONL line-by-line and account each row independently. A malformed line
is skipped (its request cannot run upstream anyway) and a row the counter cannot
measure falls back to a conservative size-based estimate. The loop never aborts, so
the allowlist check always sees every parseable model, and the token total is never
zeroed, so a crafted uncountable row still cannot evade the TPM limit, without
hard-rejecting a legitimate batch.

* perf(vertex/files): unblock async upload; drop empty finalize; widen batch MIME types

Three review follow-ups on the resumable batch upload:
- _aresumable_chunked_upload pulled chunks from a synchronous generator that runs
  the per-row transform inline on the event loop thread, blocking other requests
  between PUTs on large uploads. Each chunk is now produced via asyncio.to_thread.
- _iter_resumable_chunks no longer yields a trailing empty chunk, so an exactly
  chunk-aligned upload finalizes on its last data chunk instead of an extra
  zero-byte PUT; a 0-byte stream still finalizes via the caller's empty request.
- valid_content_type now accepts the MIME types clients label .jsonl batch uploads
  with (text/plain, application/json, ndjson, ...), so such a batch file no longer
  silently bypasses the streaming path into the buffered media upload.

* fix(vertex/files): keep legacy bucket_name as GCS bucket fallback

The rename to gcs_bucket_name dropped the legacy bucket_name key entirely, so an SDK caller passing bucket_name to a Vertex AI file create/retrieve/content call with GCS_BUCKET_NAME unset got ValueError("GCS bucket_name is required") where it previously resolved the bucket. _get_configured_bucket_name now reads gcs_bucket_name, then bucket_name, then the env var, and bucket_name is restored to OPTIONAL_KWARGS_KEYS so it survives get_litellm_params on the retrieve and content paths. gcs_bucket_name keeps precedence when both are present

* style: sort imports in llm_http_handler to satisfy I001 budget

---------

Co-authored-by: Yuneng Jiang <yuneng@berri.ai>
(cherry picked from commit 56825926af)
2026-06-24 18:10:04 -07:00
..
_experimental/mcp_server feat(mcp): scope a key to zero MCP servers with no-mcp-servers sentinel (#31029) 2026-06-23 20:00:04 -07:00
a2a feat(a2a): well-known agent-card discovery + LangGraph Platform mode (#28860) 2026-05-29 20:50:42 -07:00
agent_endpoints Support OAuth M2M for Databricks Apps A2A agents (#29586) 2026-06-04 23:03:37 -07:00
anthropic_endpoints test: isolate proxy master_key/prisma_client module globals between tests 2026-04-23 15:31:16 -07:00
auth fix(passthrough,streaming): recover cost on interrupted and agentic Anthropic streams (#31035) 2026-06-24 17:45:20 -07:00
client fix: scope CLI stored token to base_url to prevent cross-domain credential leakage (#26945) 2026-05-01 12:11:32 -07:00
common_utils fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter (#29358) 2026-05-30 17:48:16 -07:00
db fix(proxy): return 5xx on DB infra errors during auth; reserve 401 for genuine auth failures (#29986) 2026-06-15 18:22:38 -07:00
discovery_endpoints style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
experimental/mcp_server
google_endpoints Gemini managed agents support (#28270) 2026-05-19 16:02:03 -07:00
guardrails fix(guardrails): stop re-initializing DB guardrails on every poll (#30542) 2026-06-20 11:49:30 -07:00
health_endpoints feat(proxy): native /health/drain preStop hook for graceful shutdown (#29439) 2026-06-02 16:30:44 -07:00
hooks fix(vertex/files): stream OpenAI->Vertex batch JSONL uploads (#31036) 2026-06-24 18:10:04 -07:00
image_endpoints style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
management_endpoints Merge pull request #29528 from aanchal22/litellm_byok-alias-merge 2026-06-15 18:32:16 -07:00
management_helpers feat(mcp): scope a key to zero MCP servers with no-mcp-servers sentinel (#31029) 2026-06-23 20:00:04 -07:00
memory Litellm memory improvements v2 (#26541) 2026-04-25 19:03:43 -07:00
middleware fix(proxy): point /metrics 401 at the opt-out flag 2026-05-08 18:09:14 -07:00
openai_files_endpoint fix(vertex/files): stream OpenAI->Vertex batch JSONL uploads (#31036) 2026-06-24 18:10:04 -07:00
pass_through_endpoints fix(passthrough,streaming): recover cost on interrupted and agentic Anthropic streams (#31035) 2026-06-24 17:45:20 -07:00
policy_engine fix(guardrails): return HTTP 400 for litellm content filter blocks (#28418) 2026-05-30 14:28:46 -07:00
prompts style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
proxy_server fix(proxy): list public team model name in /v1/models (#30588) 2026-06-17 18:34:07 -07:00
public_endpoints fix(mcp): gate /public/mcp_hub strictly on litellm.public_mcp_servers (#27764) 2026-06-04 17:26:59 +05:30
rag_endpoints fix: harden /key/update authorization checks (#27878) 2026-05-14 04:16:04 +00:00
realtime_endpoints test: isolate proxy master_key/prisma_client module globals between tests 2026-04-23 15:31:16 -07:00
response_api_endpoints Litellm websocket improvements (#29563) 2026-06-03 11:48:35 -07:00
shutdown feat(proxy): native /health/drain preStop hook for graceful shutdown (#29439) 2026-06-02 16:30:44 -07:00
spend_tracking fix(proxy): record partial spend on the failure row for interrupted streams (#30788) 2026-06-24 17:41:54 -07:00
test_configs
types_utils chore(proxy): also scrub guardrail callbacks / module paths from DB overlay 2026-05-14 01:24:51 +00:00
ui_crud_endpoints chore(ci): merge dev branch (#28657) 2026-05-25 13:48:47 -07:00
utils fix(proxy): return deprecated-key lookup result directly in get_data combined view (#30327) 2026-06-15 18:37:37 -07:00
vector_store_endpoints Title: fix(proxy): resolve vector store file list credentials from team deployments (#29739) 2026-06-06 12:36:05 -07:00
__init__.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
conftest.py feat(proxy): native /health/drain preStop hook for graceful shutdown (#29439) 2026-06-02 16:30:44 -07:00
test_aiohttp_cleanup_closed.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_aiohttp_session_recovery.py fix: use AsyncMock for concurrent test consistency 2026-03-18 00:54:23 +00:00
test_api_key_masking_in_errors.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_audio_speech_prometheus_hooks.py fix req changes 2026-02-28 21:32:57 +05:30
test_batch_expiry.py fix(proxy): strip LiteLLM policy tracking from OpenAI batch metadata (#28425) 2026-05-26 11:35:42 -07:00
test_batch_metadata_none_fix.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_batch_retrieve_bedrock.py fix(proxy/batches): forward model to retrieve_batch for bedrock 2026-04-29 22:48:03 +02:00
test_budget_reservation.py fix(proxy): gate image-gen reservation strictly on model mode 2026-05-09 09:16:27 -07:00
test_caching_routes.py Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
test_chat_completion_metadata.py fix: propagate JWT auth metadata to OTEL spans (#19627) 2026-01-23 21:21:23 -08:00
test_common_request_processing.py fix(guardrails): return 400 not 500 when AIM blocks a request (#30573) 2026-06-20 11:58:38 -07:00
test_component_allowlists.py fix(anthropic/adapter): emit thinking block for reasoning_content-only streaming chunks (#29600) 2026-06-03 11:26:13 -07:00
test_cors_config.py refactor: extract _get_cors_config() for testability, fix no-op CORS tests 2026-04-11 22:24:04 +05:30
test_custom_proxy.py fix(ui/): fix routing for custom server root path (#15701) 2025-10-23 13:59:29 -07:00
test_dynamic_mcp_route.py Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
test_empty_model_list.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_enforce_user_param.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_fallback_management_endpoints.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_fastapi_offline_routes.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_filter_models_by_team_access_group.py fix: resolve access group names in _filter_models_by_team_id (#25224) 2026-05-01 11:55:36 +05:30
test_health_check_functions.py feat(proxy): skip disable_background_health_check models on GET /health when flag set (#27716) 2026-05-13 09:49:05 -07:00
test_health_check_max_tokens.py Litellm agent oss staging 05 11 2026 (#27733) 2026-05-13 14:09:12 -07:00
test_langfuse_passthrough_security.py chore(callbacks): guard dynamic integration hosts 2026-04-30 14:27:19 -07:00
test_lazy_openapi_snapshot.py [Fix] Tests: Move Misplaced Import in Lazy OpenAPI Snapshot Test 2026-05-01 14:24:53 -07:00
test_litellm_pre_call_utils.py fix(proxy): enforce tag budgets for key-level tags (#29108) 2026-05-28 11:39:02 -07:00
test_max_budget_env_var.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_mcp_asgi_response.py fix(mcp): surface upstream 401 for token-forwarding MCP servers (#27847) 2026-05-13 12:03:36 -07:00
test_model_dump_with_preserved_fields.py Fix_mapped tests part 2 2026-02-26 12:43:39 +05:30
test_model_id_header_propagation.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_model_info_default_limits.py fix(proxy): align /v1/model/info with router deployments (#30025) 2026-06-15 18:22:23 -07:00
test_model_level_guardrails.py fix(guardrails): run pre_call hook once for model-level guardrails (#30543) 2026-06-20 11:49:30 -07:00
test_model_list_healthy_only.py feat: add opt-in healthy_only filter to GET /v1/models (#30130) 2026-06-17 18:33:56 -07:00
test_openapi_schema_validation.py fix(proxy): preserve HTTP operations when injecting WebSocket stubs into OpenAPI schema 2026-05-06 00:28:42 +02:00
test_pricing_field_strip.py fix(proxy): move pricing strip below the litellm_metadata JSON-string parse 2026-05-03 02:04:05 +00:00
test_prometheus_cleanup.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_provider_url_destination_guard.py chore(proxy): move URL-valued model/file_id guard from SDK to proxy 2026-05-02 02:17:46 +00:00
test_proxy_cli.py feat(proxy): add option to disable server-side prepared statements for DB lookups (#29984) 2026-06-15 18:22:23 -07:00
test_proxy_logging_hook_detection.py perf: reduce per-request and per-chunk overhead across Anthropic streaming hot paths (#28289) 2026-05-23 12:15:59 -07:00
test_proxy_server.py fix(proxy): align /v1/model/info with router deployments (#30025) 2026-06-15 18:22:23 -07:00
test_proxy_types.py fix(proxy): add default=None to LiteLLM_TeamMembership.litellm_budget_table (#29684) 2026-06-04 12:13:11 -07:00
test_proxy_utils.py fix(proxy): record partial spend on the failure row for interrupted streams (#30788) 2026-06-24 17:41:54 -07:00
test_pyroscope.py Implement normalize_nonempty_secret_str function to trim whitespace from secrets and treat empty values as unset. Update proxy_server to use this function for Grafana credentials. Enhance tests to validate the new normalization behavior. 2026-05-04 18:17:31 +00:00
test_redis_auth_cache_flag.py refactor(proxy): replace DualCache with UserApiKeyCache for user API key management 2026-04-28 19:15:03 +00:00
test_response_model_sanitization.py perf: eliminate per-request callback scanning on proxy hot path (#27858) 2026-05-14 09:28:31 -07:00
test_route_a2a_models.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_route_llm_request.py feat(proxy): fix vector store retrieve/list/update/delete without model (#27929) 2026-05-15 04:37:59 -07:00
test_sensitive_route_auth.py chore(proxy): guard sensitive public endpoints 2026-04-30 11:52:47 -07:00
test_shared_health_check.py feat(proxy): skip disable_background_health_check models on GET /health when flag set (#27716) 2026-05-13 09:49:05 -07:00
test_spend_log_cleanup.py fix(proxy): keep spend log cleanup running after batch failures and surface DB errors (#27303) 2026-05-06 18:39:15 +00:00
test_swagger_chat_completions.py fix(proxy): avoid misleading multi-method operation ids 2026-04-30 20:44:14 -07:00
test_team_member_update.py
test_team_org_move.py fix(team_endpoints): auto-add SSO team members to org on move (proxy admin only) (#26377) 2026-04-24 08:36:25 -07:00
test_tools_allowlist_enforcement.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_update_llm_router_resilience.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00