Commit graph

67 commits

Author SHA1 Message Date
Cursor Agent
63579f1e35
fix(spend-tracking): keep batch spend keys joinable after v1.99 provenance gate
Batch cost attribution and the legacy queue endpoint already store the
VerificationToken hash in user_api_key, but omitted user_api_key_hash.
Since v1.99 the spend-log writer re-hashes any key without that provenance
flag, so DailyUserSpend.api_key no longer joins VerificationToken and Usage
shows key-hash-... rows with null api_key_alias / user_email.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
2026-09-03 14:12:57 +00:00
mateo-berri
048e82ec65 Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_decrease_anys_opus5_r2
# Conflicts:
#	basedpyright-code-budget.json
#	enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
#	ruff-strict-budget.json
#	type-discipline-budget.json
2026-08-29 15:11:45 -07:00
Mateo Wang
128d7e5278 refactor(batches): make count_error_file_failed_requests public for the poller import 2026-08-29 14:09:18 -07:00
Mateo Wang
4ef012627d fix: count error-file failures in the batch cost poller path 2026-08-29 12:06:42 -07:00
mateo-berri
14484d67fd refactor(types): replace Any with real types across 54 more backend files
Second pass over the highest-Any-density modules that the first pass left
untouched: guardrail hooks, the gemini and anthropic transformation layers,
the proxy spend-tracking and pass-through endpoints, and the caching clients.

Untyped `response.json()` bodies and `dict[str, Any]` request payloads are
described once at their boundary with a TypedDict or Protocol, so the fields
read downstream resolve to real types instead of Any. No cast, no type: ignore,
no noqa, and no new Any annotations.
2026-08-29 19:00:43 +00:00
Mateo Wang
c3edb95e8d Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_managed_batches_observability
# Conflicts:
#	tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
2026-08-29 10:56:02 -07:00
devin-ai-integration[bot]
4bf40c4e8d
fix(logging): stop billing and logging response reads as LLM calls (#36890)
* fix(logging): stop billing and logging response reads as LLM calls

Retrieving, deleting or cancelling a stored response, and vector store management calls, run through the same logging lifecycle as inference. A retrieved response replays the usage of the call that created it, so every read priced it again and wrote a second spend log row for the same tokens. Non-inference calls now cost 0, report no usage, log no placeholder chat message, and get a litellm.responses_management operation name instead of reading as chat.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(responses): keep billing background response jobs after the poll

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(logging): use an empty list for read-call messages

A tuple matches no branch in the loggers that walk this value, so lunary's
parse_messages falls through to clean_message and raises AttributeError on the
success hook. An empty list reads as no messages everywhere: it satisfies the
isinstance(list) checks in newrelic, mlflow and datadog, iterates zero times in
traceloop and helicone, and is what StandardLoggingPayload.messages is typed to
hold. None would be type-legal too but is not iterable, so it trades one crash
for another in mlflow and traceloop.

* fix(otel): stop the legacy emitter reporting replayed tokens on response reads

The zeroing so far lands in the standard logging payload, which the legacy
OpenTelemetry emitter does not read for usage: it takes prompt, completion and
total tokens straight off the response object, so a retrieval span still carried
the token counts of the call that produced the response, and the token usage
histogram still recorded them. That emitter is the default, so the spend row said
zero while the trace said otherwise. The background cost poller keeps its counts,
the same exemption the pricing path already makes.

* fix(logging): keep billing a background response when its retrieval is read

A response created with background=true comes back queued and carries no usage, so
its create bills nothing. The retrieval that first sees the finished job is the only
place that job's tokens are ever visible, and pricing every read at zero therefore
loses the spend outright rather than deduplicating it. On a proxy without the
enterprise cost poller a background job ended up costing $0 end to end.

is_unbilled_non_inference_call now takes the response it is deciding about and treats
a background response the same way it already treats the poller's own read, which is
the same exemption seen from the other side. The legacy OpenTelemetry emitter's time
per output token metric picks up the read gate it was missing, so it stops dividing a
read's latency by the replayed completion token count.

* test(proxy): pass the read response to the non-inference predicate

The poller test called is_unbilled_non_inference_call with the pre-background signature, so it broke when the predicate gained the response it classifies. It now hands the predicate a foreground read, and asserts that the same read is free without the origin stamp, so the stamp is what the test proves.

* fix(otel): stop the v2 metrics recorder reporting replayed tokens on response reads

The v2 span builder sources usage from the standard logging payload, so the
earlier fix already zeroes it there. The metrics recorder reads response_obj
directly, so a responses-management read still recorded the original
generation's tokens into gen_ai.client.token.usage and divided generation time
by them for gen_ai.server.time_per_output_token.

The read still records operation and response duration, under the
litellm.responses_management operation, so it stays observable.

* fix(proxy): keep the response-cost headers on calls priced at zero

Pricing responses reads and vector-store management routes at zero dropped the whole
x-litellm-response-cost family off those replies. The header build reads a falsy zero as
a cost this response never recorded and filters it out, and a call that returns before
pricing stores no cost breakdown for the component headers to read, so a client parsing
the cost off a read got a KeyError where it had previously been handed a number.

Those calls now advertise the family at zero. Retrieving a background response, and the
cost poller's read of one, still report their real cost.

The params-taking form of the predicate moves from opentelemetry into
internal_call_metadata so the proxy header build and the OTEL recorders share one copy.

* fix(proxy): report a zero cost split only under a zero cost total

The component headers were filled from call-type membership alone, while the
total they sit beside keeps its real value when the read priced normally, so a
breakdown that had not landed by the time headers were built could advertise a
real total next to an all-zero split. The split is now reported as zero only
when the total agrees with it, and is otherwise left absent.

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Yucheng Zhu <yucheng@berri.ai>
2026-08-26 18:34:17 -07:00
mateo-berri
9dabd72f2d refactor(repositories): type prisma table access with one generic protocol
Every repository handed its `.table` back untyped, so a dozen modules had
each grown a private `_PrismaTableActions` Protocol to paper over it. They
had drifted: some declared `update` as returning the row, others the row or
None, and none agreed on whether `find_many` was covariant

Replace all of them with a single `TableActions[RowT_co]` in
`litellm/repositories/prisma_protocols.py`, keyed to the prisma row each
repository is bound to. Query inputs stay `Mapping[str, object]` so callers
keep passing plain dicts, and `find_many` returns `Sequence` so the row type
stays covariant

Typing the nullable returns honestly surfaced paths that were already
crashing. A team admin could never edit or delete a memory entry owned by
their team: the write-auth check fed a raw prisma row to a helper that
expects the domain model, so `members_with_roles` arrived as plain dicts and
the request died as a 500 instead of applying the edit. Non-admin members hit
the same 500 in place of the 403 they were owed, so refusal and breakage were
indistinguishable. `/v2/model/info?user_models_only=true` dereferenced a
missing user row rather than returning the 400 the route already had, three
team routes dereferenced a team deleted between the read and the write, and
the agent registry dereferenced a missing agent instead of naming it

basedpyright drops 2,132 errors, 1,454 of them reportAny and 73
reportExplicitAny. The dashboard's generated types pick up `string[]` where
they had `unknown[]` for a team's members, admins and models
2026-08-25 12:14:17 +00:00
mubashir1osmani
65af77c43b merge(litellm_internal_staging): reconcile batch observability with per-line resilience
Staging split batch output-line costing into _safe_output_line_stats /
_compute_output_line_stats / _output_line_cost so one uncostable line can no
longer zero a whole batch, and added _provider_output_file_id so model-encoded
output file ids decode before the fetch. This branch's pass/fail counting was
written against the pre-split shape, where every None line meant a provider
failure.

Keep staging's structure and layer the counts on a three-way classification: a
provider-reported failure yields PROVIDER_FAILED, a provider-successful line
litellm cannot price yields UNCOSTABLE and stays in successful_requests billed
at $0. Without that split a litellm-side pricing gap would be reported to the
customer as a failed request and the counts would stop reconciling with the
provider's own request_counts.

Route the error-file fetch through _provider_output_file_id too, and carry the
new dataclass return through the callers staging added after this branch
forked.
2026-08-24 19:08:40 -04:00
mateo-berri
af18f77db6 fix(check_batch_cost): leave a lagging-output completed batch for the next poll cycle 2026-08-22 12:48:39 -07:00
devin-ai-integration[bot]
b2aff8be0f
fix(proxy): claim batch cost rows atomically so multi-pod polling can't double-bill (#37685)
Every pod and uvicorn worker schedules its own CheckBatchCost poller against the
shared managed-object table, so two of them can select the same completed batch in
one polling window and both write an aretrieve_batch spend log for it, counting
that batch's cost twice.

Claim the row with a compare-and-swap on batch_processed, and skip the batch when
another pod already holds it. The claim sits immediately before the spend log is
written rather than before the results fetch, because batch_processed is also what
blocks deletion of the files the fetch reads and what keeps an unbilled row
selectable by later poll cycles, so claiming up front would strand the spend of any
worker that died mid-fetch. A failed spend log write hands the row back.

Co-authored-by: Yassin Kortam <yassin@berri.ai>
2026-08-20 16:21:04 -07:00
mubashir1osmani
24a6d4de94 Merge remote-tracking branch 'berri/litellm_internal_staging' into litellm_managed_batches_observability
# Conflicts:
#	enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
#	litellm/batches/batch_utils.py
2026-08-18 19:48:44 -04:00
mubashir1osmani
4401b85855 Merge remote-tracking branch 'berri/litellm_internal_staging' into litellm_managed_batches_observability 2026-08-17 19:08:12 -04:00
mateo-berri
5a11fe141e fix(batches): price poller-tracked batches from the deployment's registered rates 2026-08-17 15:36:39 -07:00
mateo-berri
2f9d331b4c fix(proxy): retire terminal batches whose advertised output file 404s instead of retrying 2026-08-17 12:36:39 -07:00
mateo-berri
21984101e5 fix(proxy): bill cancelled and failed batches that still produced an output file 2026-08-17 12:23:59 -07:00
mubashir1osmani
141ada1118 feat(batches): aggregate reasoning tokens and per-line pass/fail counts
Batch retrieval already computed cost/usage on completion, but silently
dropped reasoning tokens and never counted per-line success/failure.
Adds BatchCostUsageResult (replacing bare cost/usage/models tuples) with
successful_requests/failed_requests, and threads reasoning_tokens through
the aggregated Usage. Both surface on SpendLogs the same way batch_models
already does.
2026-08-17 11:25:55 -04: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
mateo-berri
a10669b28c Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_batch_cost_accounted_once 2026-08-15 12:17:48 -07:00
mateo-berri
f93098068e Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_do_36634
# Conflicts:
#	litellm/batches/batch_utils.py
2026-08-15 12:12:47 -07:00
Mateo Wang
dc92749c07
Merge pull request #35360 from BerriAI/devin_ai_fix_batch_cost_completed_no_output
fix(batches): mark terminal batch with no output file as processed in CheckBatchCost
2026-08-14 17:33:07 -07:00
Marty Sullivan
ec52858865 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.
2026-08-14 01:46:56 -04:00
Marty Sullivan
c99a1ab0d7 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
2026-08-14 01:17:56 -04:00
mateo-berri
eacea13a25 fix(batches): persist real terminal status when billing expired batches 2026-08-13 21:24:54 -07:00
Devin AI
19184694f5 fix(batches): mark terminal batch with no output file as processed in CheckBatchCost
A managed batch whose request lines all failed can reach a terminal provider
status (completed) with output_file_id=None and only an error_file_id. Such a
row matched neither the completed-with-output billing branch nor the
failed/expired/cancelled branch, so batch_processed stayed False and the poller
re-selected it on every cycle for the lifetime of the deployment; output/error
file deletion is also gated on batch_processed, so those files could never be
deleted.

Broaden the terminal handling so a completed/complete/expired batch with an
output file is billed, and any terminal batch with nothing to bill
(failed/cancelled, or completed/expired with no output) is marked terminal
exactly once. Non-terminal statuses (validating/in_progress) are still left for
the next poll, and an expired batch that did produce output is now billed.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-08-13 20:26:38 -07:00
mateo
3b09484344 refactor(batches): decode unified ids through the public helper
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-08-13 19:48:16 +00:00
mateo
8947008fd2 fix(batches): only retire on a 404 that names the batch
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-08-13 04:22:43 +00:00
mateo-berri
da84142288 fix(batches): only trust a 404 from the batch's own deployment 2026-08-12 21:19:37 -07:00
mateo
c11ebbed27 fix(batches): stop uncostable batches from starving the cost poll page
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-08-12 23:45:36 +00:00
yucheng-berri
efc4e6f28c
fix(batches): keep batch state in sync on a poll without claiming attribution (#34456)
A poll of a Vertex passthrough batch wrote nothing to the managed-object row,
so status and file_object stayed frozen at the create-time snapshot and
GET /v1/batches served a stale status and an empty output file id for the life
of the batch. Only the create may claim a batch, but every observation of one
may refresh its state.

store_unified_object_id takes create_if_missing, which the poll clears: it
refreshes status and file_object through update_many, and leaves a row that is
absent absent rather than creating one owned by the observer, since created_by
and team_id are written by whoever reaches the create branch. The update payload
is now shared with the upsert so it cannot drift into writing api_key,
request_tags, created_by or team_id.

The passthrough identity re-assertion that was previously part of this PR ships
separately in #36121, so this PR keeps only the batch attribution work.

The creating key owns user_api_key_alias only when it actually has one. Guarding
the overwrite on the presence of a key rather than on a resolved alias nulled the
field out for every key generated without key_alias, and for any key rotated or
deleted before its batch finished, losing the creating user's alias that the spend
row previously carried. The guard now matches the team-alias line below it.
2026-08-08 16:01:47 -07:00
Mateo Wang
73ea5e5602
Merge pull request #36048 from BerriAI/litellm_cancelled_batch_unified_output_ids
fix(batches): persist managed file ids for cancelled/failed/expired batches
2026-08-06 10:40:07 -07:00
Mateo Wang
b66d4e6965
Merge pull request #35137 from BerriAI/litellm_fix_responses_cost_router_35131
fix(proxy): fetch background responses through the router in CheckResponsesCost
2026-08-06 03:26:36 -07:00
mateo-berri
e0c4c7cee0 Merge branch 'litellm_internal_staging' into bugfix/managed-batch-cost-not-logged 2026-08-05 23:42:18 -07:00
Devin AI
55c392bda1 merge litellm_internal_staging 2026-08-06 06:15:17 +00:00
mateo-berri
4f7d1fce3a fix(proxy): fall back to the SDK when a queued response's deployment is missing 2026-08-05 23:13:15 -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
mateo-berri
f3bfa19ce5 fix(managed_files): resolve model_name identically across all output file registration paths so full unified ids converge 2026-08-05 17:36:22 -07:00
elinacse
833670f7db fix(batch): track cost for managed batches with no attributable key/user/team
LiteLLM_ManagedObjectTable only stores created_by (user_id) and team_id,
never the raw API key hash. A batch created with the master key or a
team-less key has both null, so CheckBatchCost's synthetic logging_obj
for the completed batch carried no attributable key/user/team/end-user.
_should_track_cost_callback silently skipped the DB write in that case
(by design, to avoid tracking truly anonymous requests), with no error
or warning: batch_processed still became true, but no LiteLLM_SpendLogs
row was ever written despite real, already-incurred provider cost.

Extend the same allowance already made for unauthenticated pass-through
requests to aretrieve_batch's cost event, and pass job.team_id through
so a batch's team gets real attribution when one exists.
2026-08-02 12:20:46 +05:30
Devin AI
6062ed7ff1 fix(batches): encode public model group on background-created output file ids
CheckBatchCost built unified output file ids with the provider model name, so key model-access checks resolved the file to e.g. gpt-5.5 and every GET /v1/files/{output_file_id}/content failed. Resolve the model group from the batch's managed input file, falling back to the deployment's model_name.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-07-31 19:57:12 +00:00
Devin AI
4429742e83 fix(proxy): fetch background responses through the router in CheckResponsesCost
Closes #35131
2026-07-29 21:11:44 +00:00
Sameer Kankute
c2141b1113
feat(batches): track cost for unmanaged Bedrock batches, generalize the flag (#32315)
* feat(batches): track cost for unmanaged Bedrock batches, generalize the flag

CheckBatchCost skipped Bedrock batches whose unified_object_id is a raw
model-invocation-job ARN, the same root cause previously fixed for
unmanaged Vertex batches. Bedrock batches embed the model name in their
s3:// input file name instead (litellm-bedrock-files-{model}-{uuid}.jsonl),
so the same routing mechanism now derives the model from that layout and
matches it to a configured bedrock deployment.

track_unmanaged_vertex_batch_cost is renamed to track_unmanaged_batch_cost
since two providers now share this mechanism.

* fix(batches): parse Bedrock batch output and price with deployment model name

Bedrock model-invocation-job results use modelOutput/error rows and short
internal model ids that are not in the cost map, so unmanaged batch cost
tracking logged tokens but $0 spend. Use deployment model name for pricing
and add regression tests.

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

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-13 10:33:13 +05:30
Mateo Wang
9076c33347
fix(batches): price anthropic passthrough message batches correctly in batch cost job (#32307)
* fix(batches): price anthropic passthrough message batches correctly in batch cost job

Anthropic message batches created via the /anthropic passthrough were never
cost tracked. The CheckBatchCost job fetched batch results from the Files API
(POST /v1/files/msgbatch_.../content), which Anthropic rejects with "File id
must have file_ prefix"; the error response was silently wrapped as file
content, parsed as zero successful rows, logged as a $0 aretrieve_batch spend
row, and the job was marked batch_processed=true so the $0 was permanent.

Route msgbatch_ file ids to GET /v1/messages/batches/{id}/results in the
anthropic files transformation, raise on HTTP error status in
retrieve_file_content instead of returning the error body as content, parse
Anthropic's results JSONL shape (result.type == "succeeded",
result.message.usage with cache creation/read tokens) in batch_utils, price
cache creation tokens at cache_creation_input_token_cost in the batch cost
fallback (50% batch discount preserved for base input, cache reads, cache
writes, and output), and leave the managed object row unprocessed when cost
tracking fails so a later poll retries instead of permanently recording $0.

* fix(batches): carry cache token details into aggregated anthropic batch usage
2026-07-06 20:33:57 -07:00
Sameer Kankute
6d796d0f1f
feat(proxy): track cost for unmanaged Vertex AI batch jobs (#31442)
* feat(proxy): track cost for unmanaged Vertex AI batch jobs

CheckBatchCost previously skipped Vertex batches created via the raw GCS
input_file_id path, since their unified_object_id is a raw provider job id
that fails the base64 managed-id check. Behind the opt-in general_settings
flag track_unmanaged_vertex_batch_cost, the poller now derives the model
from the gs:// input_file_id, maps it to a configured vertex_ai deployment,
polls the batch, computes cost, and marks batch_processed=True.

* Update tracking for failed", "expired", "cancelled"

* fix(proxy): apply ruff format to proxy_server.py

* address greptile review feedback (greploop iteration 1)

Filter unmanaged Vertex batch deployments by vertex_ai provider so a
shared model group name can't route to a wrong-provider deployment.
Move gs:// URI parsing into VertexAIBatchTransformation. Add test
coverage for the failed/expired/cancelled terminal-status DB update.

* fix: route unmanaged vertex batches to matching deployment

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
2026-07-02 21:19:54 +05:30
Sameer Kankute
cbdc70d544
fix(managed_batches): convert raw output_file_id to managed ID in CheckBatchCost poller (#27984)
* fix(managed_batches): convert raw output_file_id to managed ID in CheckBatchCost poller

CheckBatchCost bypasses async_post_call_success_hook, causing raw provider
output_file_ids to be persisted in LiteLLM_ManagedObjectTable. This fix converts
output_file_id and error_file_id to managed base64 IDs before the DB write.

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

* fix(check_batch_cost): persist managed file before mutating response and propagate team_id

- Move setattr after store_unified_file_id so the response only receives the
  managed ID once the DB record is successfully written. Avoids serializing
  an orphaned managed ID into file_object when the store call fails.
- Populate team_id on the minimal UserAPIKeyAuth from job.team_id so the
  managed file record is created with the correct team ownership, allowing
  other team members to access the batch output file via /files/{id}/content.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* test(managed_batches): extend test to cover error_file_id conversion

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

* fix managed file test

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
2026-05-15 04:41:38 -07:00
ishaan-berri
7a9a9f0c79
fix: batch-limit stale managed object cleanup to prevent 300K row UPD… (#25258)
* fix: batch-limit stale managed object cleanup to prevent 300K row UPDATE (#25257)

* Add STALE_OBJECT_CLEANUP_BATCH_SIZE constant

Configurable batch limit (default 1000) for stale managed object cleanup,
preventing unbounded UPDATE queries from hitting 300K+ rows at once.

* Batch-limit stale managed object cleanup with single bounded SQL query

Two fixes to _cleanup_stale_managed_objects:

1. Replace unbounded update_many with a single execute_raw using a
   subquery LIMIT, capping each poll cycle to STALE_OBJECT_CLEANUP_BATCH_SIZE
   rows. Zero rows loaded into Python memory — everything stays in Postgres.
   Uses the same PostgreSQL raw-SQL pattern as spend_log_cleanup.py
   (the proxy requires PostgreSQL per schema.prisma).

2. Extract _expire_stale_rows as a separate method for testability.

Keeps the file_purpose='response' filter to avoid incorrectly expiring
long-running batch or fine-tune jobs that legitimately exceed the
staleness cutoff.

* docs: add STALE_OBJECT_CLEANUP_BATCH_SIZE to env vars reference

* test: remove deprecated embed-english-v2.0 cohere embedding tests
2026-04-06 19:11:55 -07:00
Sameer Kankute
bbd8ca3b3d
feat(prometheus): add metrics for managed batch lifecycle
- Add Prometheus metrics for managed batch and file operations
- Track batch creation, file size, duration, and deletion events
- Add CheckBatchCost polling metrics (jobs polled/processed, errors)
- Record metrics in managed_files hook and check_batch_cost utility
- Metrics include labels for model, provider, user, and status

Made-with: Cursor
2026-03-27 20:30:09 +05:30
yuneng-jiang
4fc0975d22 Fix flaky e2e batch test: set batch_processed=True on completion in retrieve_batch
The retrieve_batch endpoint sets batch status to "complete" but never set
batch_processed=True, permanently blocking file deletion. CheckBatchCost
(the safety net) also excluded completed batches from its primary query,
so batch_processed was never set by either path.

Three fixes:
1. update_batch_in_database sets batch_processed=True when status reaches
   "complete", with old-schema fallback retry
2. CheckBatchCost primary query no longer excludes complete/completed
   (batch_processed=False filter prevents reprocessing)
3. retrieve_batch early-return now includes "complete" (DB-normalized
   spelling) to avoid unnecessary provider re-polls

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 18:18:32 -07:00
Ishaan Jaff
1b96064600
fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll (#23472)
* fix(proxy): cap managed-object poll size + expire stale rows + kill-switch flag to prevent OOM/Prisma connection loss

* fix(constants): simplify PROXY_BATCH_POLLING_ENABLED readability

* docs+test: document new polling env vars, add pagination+stale-cleanup tests

* fix: exclude stale_expired from batch poll queries; fix update_many assertions in tests

* fix: scope stale cleanup to file_purpose, fix file_object mocks, add CheckBatchCost tests

* fix: avoid duplicate cost logging in fallback path; guard integer constants against zero/negative values

* fix: cache _has_batch_processed_column; guard cleanup from aborting poll; narrow fallback except

* fix: add complete/completed to primary query not_in; fix vacuous test assertion

- Primary find_many was missing "complete" and "completed" in its not_in
  filter, creating asymmetry with the fallback query. A job whose status
  was set to "complete" but whose batch_processed flag update failed would
  be silently re-fetched and re-processed every cycle, emitting duplicate
  cost logs.

- test_fallback_completion_update_omits_batch_processed patched
  _is_base64_encoded_unified_file_id to return None, causing an immediate
  continue — so update() was never called and the assertion looped over an
  empty list (vacuously true). Rewrote the test to mock the full
  completion pipeline, verify update() is called exactly once, and assert
  batch_processed is absent from the update data.

- Added symmetric test (primary path) proving batch_processed IS included
  when the column exists.

Made-with: Cursor
2026-03-13 11:01:40 -07:00
Ryan Crabbe
c4db53a98a Address review feedback: remove dead code, add error handling, strengthen test assertions
- Remove unused `completed_jobs` list (dead code after per-job update refactor)
- Wrap DB update in try/except to prevent one failed update from aborting remaining jobs
- Add test assertions verifying batch_processed, status, and file_object are written to DB
2026-03-06 09:25:50 -08:00