Commit graph

27 commits

Author SHA1 Message Date
yuneng-jiang
29ac88ebc6
fix(batches): register ownership for every batch create path (#39810)
Some checks failed
Unit Tests / caching-local (push) Waiting to run
Unit Tests / core-utils (push) Waiting to run
Unit Tests / enterprise-package (push) Waiting to run
Unit Tests / proxy-auth (push) Waiting to run
Unit Tests / proxy-endpoints (push) Waiting to run
Unit Tests / proxy-extras (push) Waiting to run
Unit Tests / proxy-infra (push) Waiting to run
Unit Tests / proxy-server (push) Waiting to run
Unit Tests / responses-caching-types (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests / enterprise-routing (push) Waiting to run
Unit Tests / integrations (push) Waiting to run
Unit Tests / All Other Providers (push) Waiting to run
Unit Tests / Vertex AI (push) Waiting to run
Unit Tests / misc (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
* fix(batches): register ownership for every batch create path

Since the team isolation change, the managed files hook decided whether a
response came from a create by looking for the managed input file id on it,
which only the unified input path sets. Batches created from a model-encoded
input file id, a model param, or a raw provider id with ?provider= never got
an ownership row, so they vanished from GET /v1/batches for the key that
created them.

The create endpoint now stamps a create marker on the response before the
hooks run, and the hook keys ownership registration and the batch-created
metric on that marker instead of on the input id format.

* test(batches): assert ownership registration through the managed files hook

The endpoint tests asserted the private create marker, which is wiring, not
behaviour. They now run the create through the real managed files hook and
assert the ownership row is written for the creating key on every create
path, with the unified path driven by a genuine encoded input file id
instead of patched decoders.
2026-09-04 23:59:51 -07:00
Mateo Wang
75613bf22f
test: add regression coverage for twelve closed issues (#37974)
* test: add regression coverage for twelve closed issues

Adds targeted regression tests for behavior that was fixed but left ungated,
so the fixes cannot silently regress:

- #33772 openai cache_write_tokens cost
- #34309 Responses API cache cost_breakdown
- #35363 /v1/responses batch spend
- #36619 auto-router api_base/api_key leak on a shared model name
- #35359 batch fallbacks within the owning model group
- #36523 passthrough streamed Responses spend log
- #36646 passthrough embeddings spend log
- #37147 non-object metadata on create_batch is a 400
- #35362 unscoped list files reads the managed-file store
- #33221 gpt-5.6 bridges to Responses on function tools alone
- #34487 LLM complexity classifier runs for every caller metadata shape
- #35124 streamed /v1/messages emits success logging on both bridges

Cost assertions read rates from litellm.model_cost rather than hardcoding
dollar amounts, so they do not drift on repricing.

* fix: stop the new regression tests polluting and tripping over shared global state

Two shard failures, both from global state the new tests share with their
neighbours rather than from the behaviour under test.

test_main.py's local_cost_map pinned litellm.model_cost but left the
get_model_info lru_cache warm, so completion_cost billed at whatever prices
were cached earlier in the process while the assertions read the pinned map.
Clear the cache on both sides of the fixture, matching the local_model_cost_map
fixture in tests/test_litellm/conftest.py.

The anthropic messages streaming tests called GLOBAL_LOGGING_WORKER.flush()
on whatever queue happened to be around. A queue left non-empty by an earlier
test is still bound to that test's loop, so join() either hangs or raises
"bound to a different event loop". Rebind to the running loop before the call
and wait for the captured payload instead of a fixed sleep.
2026-08-22 22:24:05 +00:00
yuneng-jiang
6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802)
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00
ryan-crabbe-berri
680bcfd8aa
test(lint): ban blind pytest.raises(Exception) with ruff B017 (#37731)
* test(lint): ban blind pytest.raises(Exception) with ruff B017

A bare pytest.raises(Exception) accepts whatever the body throws. The TypeError
a refactor introduces satisfies it exactly as well as the rejection the test was
written for, so the crash reads as a pass and the test never goes red.

All 111 existing sites are narrowed here. A runtime probe recorded the concrete
exception each one actually catches, and each site now names that type. Where
the code under test genuinely raises a bare Exception, the site pins a stable
slice of the message with match= instead.

Two sites tell on themselves. The shared responses-API cancel test raises
"custom_llm_provider is required but passed as None" rather than talking to a
provider at all, because cancel_responses takes a provider, not a model. And
test_bedrock_guardrails_with_streaming was the only test in its file still
passing without AWS credentials, because the NoCredentialsError boto3 raised
long before the guardrail ran satisfied the blind raises.

* fix(test): widen the openai batch-dispatch assertion to OpenAIError

The narrowed NotFoundError only holds where OPENAI_API_KEY is set. Without one
the SDK raises OpenAIError while building the client, long before any 404, so CI
went red. OpenAIError covers both and still rejects a TypeError from a refactor.
2026-08-20 18:09:42 -07:00
Mateo Wang
77b7c6c40c
Merge pull request #37198 from BerriAI/litellm_lit5660_batches_limit_400
fix(proxy): reject out-of-range limit on GET /v1/batches with OpenAI-parity 400
2026-08-17 15:53:46 -07:00
Mateo Wang
a6de0736e1
Merge pull request #37201 from BerriAI/litellm_fix_batches_404
fix(proxy): return 404 instead of 500 for unresolvable batch and file ids on /v1/batches
2026-08-17 15:43:49 -07:00
mateo-berri
09c8d1f1f5 Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_lit5660_batches_limit_400 2026-08-17 12:46:15 -07:00
mateo-berri
4a7dfd75fc fix(proxy): return 404 instead of 500 for unresolvable batch and file ids on /v1/batches 2026-08-17 12:02:15 -07:00
mateo-berri
a9bb09905d test(batches): drop redundant section banner 2026-08-17 12:00:40 -07:00
mateo-berri
e4ce526900 fix(proxy): return 400 naming the missing required param on POST /v1/batches 2026-08-17 11:55:06 -07:00
mateo-berri
f55a193628 fix(proxy): reject out-of-range limit on GET /v1/batches with OpenAI-parity 400 2026-08-17 11:54:34 -07:00
mateo-berri
bb1c3366cf Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_batch_cost_accounted_once
# Conflicts:
#	tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py
2026-08-15 15:56:11 -07:00
mateo-berri
d9e377f129 fix(batches): confirm poller batch_processed support at startup so no retrieve accounts inline before the first poll cycle
Probe the column before the scheduler registers CheckBatchCost, closing the window where a retrieve that decided the poller was inactive billed a batch the first poll cycle then billed again. Also drop narration docstrings and section banners from the new tests.
2026-08-15 12:56:53 -07:00
Marty Sullivan
5649098e1b 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.
2026-08-14 01:36:52 -04:00
Marty Sullivan
60fe4e464c 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.
2026-08-14 01:17:56 -04:00
Anas Khan
7fcca523aa
fix(proxy/batches): stop forwarding custom_llm_provider twice in list and cancel (#32813)
* fix(proxy/batches): stop forwarding custom_llm_provider twice in list and cancel

The model-routing branches of list_batches and cancel_batch passed
custom_llm_provider as an explicit kwarg while also leaving it inside the dict
they splat, so every such call raised "got multiple values for keyword argument
'custom_llm_provider'" and returned a 500.

list_batches SCENARIO 2 called data.update(credentials) but never removed
custom_llm_provider before litellm.alist_batches(custom_llm_provider=..., **data);
it now uses prepare_data_with_credentials, the same helper the create and
retrieve branches already use, which pops it out.

cancel_batch SCENARIO 3 resolved the provider with
`provider or data.pop("custom_llm_provider", None) or ...`, so when the path
param provider was set the pop short-circuited and a body custom_llm_provider
stayed in data and collided with the explicit kwarg. The body value is now
popped unconditionally before the fallback chain, so the path param wins cleanly
and data no longer carries a duplicate.

Both paths already had strict-xfail regression tests documented "remove when
fixed"; those markers are dropped so the tests now guard the fix.

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

* fix(proxy/files): avoid duplicate custom_llm_provider in list

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

---------

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>
2026-08-13 12:50:32 -07:00
mateo-berri
8c0556abf6 fix(proxy): authenticate managed ids before routing 2026-08-08 02:42:34 -07:00
mateo-berri
5f7a663005 fix(proxy): enforce require_managed_files on every raw provider id route
require_managed_files was only checked on upload, so raw provider ids still
reached the batch, fine-tuning and vector store file routes. Ownership rows
exist only for managed ids, so those requests were forwarded under shared
credentials with no tenant check: knowing another tenant's id was enough to
read, run against, cancel or delete their object.

Generalise the file-id guard to validate_managed_id_requirement(resource_id,
resource_kind) and call it on batch create/retrieve/cancel, fine-tuning
create/retrieve/cancel (training_file and validation_file both) and the shared
vector store file id resolver. Behaviour is unchanged when the setting is off.
2026-08-08 01:33:32 -07:00
Mateo Wang
0c3017e1de
Merge pull request #35371 from rimysore/fix-managed-batch-cross-provider-fallback
fix(batches): prevent managed file fallbacks
2026-08-06 11:40:38 -07:00
mateo-berri
5339ec50e7 fix(batches): persist managed file ids for cancelled/failed/expired batches
When the batch cost poller found a batch in a terminal failed, expired, or
cancelled state it wrote the provider response straight to the managed object
table, so the stored blob kept raw provider file ids and a raw batch id. Since
the row is final after batch_processed=True and the read paths only resolve
existing managed ids, every later GET /batches/{id} and GET /batches leaked
raw provider output and error file ids that clients cannot fetch through the
proxy. The terminal branch now normalizes the response with
ensure_batch_response_managed_file_ids before persisting, minting managed ids
under the batch owner's identity

POST /batches/{id}/cancel had the same gap: it called update_batch_in_database
without the caller's auth context, so a cancel response that already carried
provider file ids could never mint managed ids. The endpoint now forwards
user_api_key_dict
2026-08-05 20:42:20 -07:00
Souravrajvi0
388943ac17
fix(proxy): register managed batch output files on terminal retrieve (#34092)
Terminal batch retrieve could return the raw provider output_file_id, which
skips managed-file ownership checks on /v1/files/{id}/content and lets any key
on the proxy download another user's batch output.

Retrieve now registers the missing managed-file row before responding, and
attributes ownership to the durable batch owner rather than the retrieving
caller, so output and error ids always come back as unified managed ids.

Fixes #33989
2026-08-05 16:38:49 -07:00
Rithvik Mysore Suresh
55726fc09e fix(batches): override existing fallback flag
Build one kwargs mapping so managed-file ownership always disables cross-model fallback without duplicating a request-enriched key.
2026-07-31 10:28:52 -04:00
Rithvik Mysore Suresh
f0ffc6507e fix(batches): keep managed files on owner
Managed files and batches are provider-owned. Cross-model fallbacks can dispatch creation with credentials that cannot access the input file and replace the owning provider's validation error.\n\nCloses #35359
2026-07-31 10:16:35 -04:00
mateo-berri
17ce2c4e92 fix(proxy): resolve named credentials on provider-only batch and files calls 2026-07-28 18:57:15 -07:00
yucheng-berri
fc5ab31fba
refactor(proxy/batches): make managed-file resolution purely additive, fall back like base (#34584)
Restore the original PR's behavior on every path that did not already
resolve: no database, a lookup error, a missing managed-file row, or a
row without a storage_url all fall back to dispatching the original id,
which the managed-files deployment hook still maps. This drops the 404
and 503 fail-closed responses I had added, which were the only behaviors
that diverged from litellm_internal_staging.

The change is now strictly additive: when a managed-file row with a
storage_url exists, the unified batch branch substitutes it so providers
like Vertex receive a real gs:// path instead of the opaque token; every
other path behaves exactly as before. Verified live that non-managed,
managed-owner, multi-model load-balanced, and missing-row requests are
byte-identical to base
2026-07-24 18:42:48 -07:00
yucheng-berri
5677bc237c
fix(proxy/batches): resolve managed unified input_file_id to storage_url with ownership check before dispatch (#34474)
* fix: resolve unified_file_id to real storage_url before dispatching batch create

litellm.create_batch() against a Vertex AI-backed model crashes with an
opaque error when the input file was uploaded as a LiteLLM-managed
'unified file' (multi-model file upload). The base64-encoded
unified_file_id token is a LiteLLM-internal identifier, not a real
provider-side file reference, but the batches_endpoints create_batch
handler forwards it unchanged to llm_router.acreate_batch() /
litellm.acreate_batch() for the unified_file_id branch. Provider-specific
code that expects a real file location (e.g. Vertex AI's batch
transformation, which parses a 'publishers/' segment out of the GCS URI)
then fails on the opaque token.

Resolve the unified_file_id to its real backend location
(LiteLLM_ManagedFileTable.storage_url) before dispatch, mirroring the
same lookup already used by the files retrieve/download endpoints for
managed files. Falls back to the previous (unchanged) behavior if no
managed-file record exists.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(proxy/batches): null-guard await on find_first for sync MagicMock test harnesses

* fix(proxy/batches): enforce ownership and correct lookup key when resolving managed input_file_id

The adopted resolution queried LiteLLM_ManagedFileTable with the decoded
litellm_proxy string, but the unified_file_id column stores the raw base64
file id (see schema.prisma and the enterprise managed-files hook), so the
lookup never matched in production and silently fell back to the opaque id.
Query with the raw id instead and lock the key with a regression test.

Move the resolution above the dispatch branches so the load-balanced router
path receives the resolved storage_url too, enforce managed-file ownership
with the same can_access_resource semantics the files retrieve and download
endpoints use (404 on denial), and downgrade database failures to a logged
fallback instead of aborting batch creation. Unresolved ids still dispatch
unchanged because the managed-files deployment hook can map them via
model_file_id_mapping

* fix(proxy/batches): fail closed when the managed file ownership lookup errors

A lookup exception previously fell back to dispatching the original
unified id with the ownership gate unexecuted; the managed-files
deployment hook maps unified ids from cache without re-checking
ownership, so a database outage let a caller dispatch another tenant's
file. Raise a clear 503 instead and lock the behavior with a regression
test. No-database and no-row cases still fall back unchanged

* test(proxy/batches): default harness prisma_client to None

The batch routing harness left proxy_server.prisma_client at its module
global, which a sibling test in the same shard can leave as a MagicMock.
The unified-file rows that do not opt into managed-file resolution then
entered the resolver and awaited a non-awaitable mock, surfacing as a
503. Patch prisma_client to None by default so those rows stay a no-op;
resolution tests still override it explicitly

* fix(proxy/batches): keep unified resolution in its own branch and fail closed on missing row

Cursor flagged that hoisting the storage_url substitution above the
load-balanced dispatch branch broke two things on that path: the
model_file_id_mapping deployment filter keys on the original unified id,
and the response returned the internal storage_url instead of the
unified id. Move the resolution back inside the unified branch and
exclude unified ids from the load-balanced branch so a managed file
always takes the resolving path (which restores input_file_id and the
unified_file_id hidden param on the response), and a load-balanced batch
keeps the original id for deployment filtering.

Also fail closed with a 404 when a unified id has no managed-file row
while a database is present: the id cannot be ownership-verified, and
dispatching it would both bypass the gate and hit the Vertex
publishers-segment IndexError. Owned rows without a storage_url (legacy)
still dispatch the original id

* fix(proxy/batches): do not divert unified files off the load-balanced branch

Excluding unified ids from the load-balanced branch (and not
unified_file_id) regressed a path that works on the base revision: a
multi-model managed file dispatched with an explicit router model under
load balancing was routed into the unified branch, which raises a 400
for anything other than exactly one target model. Verified live against
base (200, managed-files deployment hook remaps the unified id per
model) versus the guarded branch (400 Expected 1 model, got 2).

Restore the original three-condition load-balanced branch so that path
keeps working unchanged. Unified-file storage_url resolution and the
ownership 404 still apply on the non-load-balanced unified branch, which
is the common managed-batch flow; the load-balanced managed path retains
its existing behavior and its pre-existing enterprise-hook ownership gap,
unchanged from base

* refactor(proxy/batches): scope managed-file handling to resolution, drop ownership check

Narrow this PR to its one problem: resolving a managed unified input_file_id
to its backend storage_url so provider batch handlers (Vertex parses a
publishers/ segment) receive a real location instead of the opaque token,
and failing closed with a 404 when the token has no backing row so it is
never dispatched into the provider crash.

Remove the cross-tenant ownership check (can_access_resource) added earlier.
Batch-create had no ownership enforcement before this PR, and the gap spans
every managed-file call type, so it belongs in the enterprise managed-files
pre-call hook (its acreate_batch branch) where files, batches and
fine-tuning are covered uniformly, not partially in this one endpoint. Filed
as a follow-up. This also removes the load-balanced-path ownership
inconsistency the bots flagged, since there is no ownership branch to skip.

Drop the inline comments flagged against the no-comments rule; behavior is
documented in the helper docstring and the test docstrings

* fix(proxy/batches): fail closed with 503 when the managed-file lookup errors

A lookup exception previously fell back to dispatching the unresolved
unified token, which defeats the fail-closed guarantee: the token still
reaches the provider and can hit the same publishers-segment IndexError
the resolution prevents. Treat a lookup error like the missing-row case
and fail closed, but with a retryable 503 since the condition is
transient. No-database and no-storage_url rows still fall back unchanged

---------

Co-authored-by: htourinho-clgx <htourinho@cotality.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-24 17:29:38 -07:00
Sameer Kankute
2cf565ae28
test(batches): add 1:1 test file scaffold for batches component paths (#30529)
* test(batches): add 1:1 test file scaffold for batches component paths

Co-authored-by: Cursor <cursoragent@cursor.com>

* Add harness test for create batch endpoint

* Add retrieve endpoint harness tests

* Add list  endpoint harness tests

* Add cancel endpoint harness tests

* Add cancel endpoint harness tests

* Add test for litellm/batches/main.py

* Add test for litellm/tests/test_litellm/batches/test_batch_utils.py

* Add handler and transformation tests for all providers

* Fix: run batches tests in cicd

* fix(tests): remove azure/__init__.py that shadowed azure namespace package

Adding __init__.py to tests/test_litellm/llms/azure/ caused pytest to
insert tests/test_litellm/llms/ into sys.path[0], making our empty
azure/ dir shadow the real azure-identity namespace package. Any test
that patched azure.identity.* would then fail with AttributeError.

* style(tests): apply ruff format to test_batch_utils.py

Base migrated the formatter from black to ruff format (#31317); reformat the
batches scaffold test file to match.

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-06-29 09:22:58 +05:30