fix(check_batch_cost): delegate file-id reconciliation to ensure_batch_response_managed_file_ids (LIT-3386)

[oss-agent-shin]

What

When the user has uploaded the batch input file through the proxy (so a LiteLLM_ManagedFileTable row exists for the raw provider input_file_id), the CheckBatchCost background poller currently persists the raw provider input_file_id in LiteLLM_ManagedObjectTable.file_object instead of the managed unified ID. GET /v1/batches/{id} then returns a raw input_file_id that the user cannot correlate back to their managed upload, and listing endpoints can't round-trip it.

update_batch_in_database (used by the retrieve/cancel HTTP path) already does this reconciliation via the central helper ensure_batch_response_managed_file_ids. CheckBatchCost had its own partial reconciliation block from #27984 that only covered output_file_id / error_file_id and didn't consult existing managed_file rows — missing both input_file_id and dedup behavior.

Closes the remaining gap called out at the bottom of #28294 ("PR does not touch check_batch_cost.py") for LiteLLM ticket LIT-3386.

Change

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py: replace the partial in-place output/error registration with a single call to ensure_batch_response_managed_file_ids (the same helper update_batch_in_database uses on the retrieve/cancel HTTP path). The helper:

1. resolves the response's raw input_file_id to the existing managed unified ID (the gap the legacy block missed),
2. resolves raw output_file_id / error_file_id to existing managed rows when one already exists (avoiding duplicate writes by parallel paths),
3. only registers a brand-new managed_file row when no existing row maps to a raw ID.

Auth context (UserAPIKeyAuth.user_id = job.created_by) is unchanged — still anchored on the batch creator so managed rows are owned by the real user, not default_user_id. response._hidden_params["model_id"] is populated defensively from the deployment we just resolved so the helper has all the metadata it needs to register new rows when the provider transformation didn't set it.

Files

• enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py — swap manual reconciliation block for ensure_batch_response_managed_file_ids.
• tests/proxy_unit_tests/test_check_batch_cost.py — update the existing raw-output coverage to the new code path; add 2 new tests for input-file-id resolution and reuse-dedup behavior.

Files pushed via the GitHub Contents API (current token lacks repo+workflow scopes) — the diff itself is exactly the 2 files above.

Evidence

Runtime evidence harness drives the real CheckBatchCost.check_batch_cost() against a mock prisma + router that mirrors the production data path (the poller has no HTTP trigger by design, so this is the closest practical reproduction). Captures the JSON file_object payload that CheckBatchCost persists to litellm_managedobjecttable.update(). The setup pre-populates a managed_file row for the raw input file id (simulating a user who uploaded their batch input through the proxy first).

BEFORE — clean upstream main HEAD 5699a06:

--- BEFORE (clean main 5699a06) ---
update() call count = 1
response.input_file_id  AFTER CheckBatchCost: 'file-bedrock-input-bbb111'
response.output_file_id AFTER CheckBatchCost: 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1vdXRwdXQtYWFhMjIyOjpvbjo6bW9kZWwtYWJjMTIz'
response.error_file_id  AFTER CheckBatchCost: 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1lcnJvci1jY2MzMzM6Om9uOjptb2RlbC1hYmMxMjM'
persisted file_object.input_file_id  = 'file-bedrock-input-bbb111'
persisted file_object.output_file_id = 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1vdXRwdXQtYWFhMjIyOjpvbjo6bW9kZWwtYWJjMTIz'
persisted file_object.error_file_id  = 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1lcnJvci1jY2MzMzM6Om9uOjptb2RlbC1hYmMxMjM'
store_unified_file_id call count     = 2
  store_unified_file_id: file_id=bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZp... mappings={'model-abc123': 'file-bedrock-output-aaa222'} user_id=shin-real-user-id
  store_unified_file_id: file_id=bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZp... mappings={'model-abc123': 'file-bedrock-error-ccc333'} user_id=shin-real-user-id
ensure_batch_response_managed_file_ids invocations = 0
FLAG_INPUT_IS_RAW=True
FLAG_INPUT_IS_MANAGED=False
FLAG_OUTPUT_IS_RAW=False
FLAG_ERROR_IS_RAW=False

AFTER — with this fix:

--- AFTER (fix applied — calls ensure_batch_response_managed_file_ids) ---
update() call count = 1
response.input_file_id  AFTER CheckBatchCost: 'bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29ubDt1bmlmaWVkX2lkLGFhYTExMTExLWJiYmItMjIyMjt0YXJnZXRfbW9kZWxfbmFtZXMsY2xhdWRlLXNvbm5ldA'
response.output_file_id AFTER CheckBatchCost: 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1vdXRwdXQtYWFhMjIyOjpvbjo6bW9kZWwtYWJjMTIz'
response.error_file_id  AFTER CheckBatchCost: 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1lcnJvci1jY2MzMzM6Om9uOjptb2RlbC1hYmMxMjM'
persisted file_object.input_file_id  = 'bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9qc29ubDt1bmlmaWVkX2lkLGFhYTExMTExLWJiYmItMjIyMjt0YXJnZXRfbW9kZWxfbmFtZXMsY2xhdWRlLXNvbm5ldA'
persisted file_object.output_file_id = 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1vdXRwdXQtYWFhMjIyOjpvbjo6bW9kZWwtYWJjMTIz'
persisted file_object.error_file_id  = 'bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZpbGUtYmVkcm9jay1lcnJvci1jY2MzMzM6Om9uOjptb2RlbC1hYmMxMjM'
store_unified_file_id call count     = 2
  store_unified_file_id: file_id=bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZp... mappings={'model-abc123': 'file-bedrock-output-aaa222'} user_id=shin-real-user-id
  store_unified_file_id: file_id=bGl0ZWxsbV9wcm94eTo6bWFuYWdlZF9mb3I6OmZp... mappings={'model-abc123': 'file-bedrock-error-ccc333'} user_id=shin-real-user-id
ensure_batch_response_managed_file_ids invocations = 1
  ensure(): auth.user_id=shin-real-user-id
FLAG_INPUT_IS_RAW=False
FLAG_INPUT_IS_MANAGED=True
FLAG_OUTPUT_IS_RAW=False
FLAG_ERROR_IS_RAW=False

Key delta (BEFORE -> AFTER):

field	BEFORE	AFTER
persisted file_object.input_file_id	'file-bedrock-input-bbb111' (raw provider value) -> FLAG_INPUT_IS_RAW=True	base64 managed unified id -> FLAG_INPUT_IS_MANAGED=True
ensure_batch_response_managed_file_ids invocations	0 (legacy partial block ran)	1 (delegated to central helper)
response.output_file_id / error_file_id	already managed (legacy block)	still managed (helper) — no regression
auth context on managed_file write	user_id=shin-real-user-id	user_id=shin-real-user-id (unchanged)

Unit tests

9/9 pass in tests/proxy_unit_tests/test_check_batch_cost.py (6 pre-existing kept as-is, 1 updated to drive through the helper, 2 new tests added):

$ python3 -m pytest tests/proxy_unit_tests/test_check_batch_cost.py --no-header
============================== 9 passed in 0.36s ==============================

Related coverage (tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py, 8/8) and broader batch/managed_file tests (pytest -k "batch or managed_file" in proxy_unit_tests, 14/14) also pass.

New tests:

• test_input_file_id_resolved_to_existing_managed_id — mocks litellm_managedfiletable.find_first to return a managed row for the raw input file id; asserts the persisted file_object.input_file_id is the managed unified id (was the raw value on main).
• test_existing_output_managed_row_is_reused_not_duplicated — mocks an existing managed row for the raw output_file_id; asserts store_unified_file_id / get_unified_output_file_id are NOT called (no duplicate write) and response.output_file_id is the existing managed id.

Updated test:

• test_raw_output_file_id_converted_to_managed_id — same scope (covers the no-existing-managed-row path) but now drives through the helper instead of the legacy block; uses a real base64-encoded unified_object_id and mocks find_first -> None.

Linear

LIT-3386 — https://linear.app/litellm-ai/issue/LIT-3386/fix-remaining-batch-id-transformation-in-checkbatchcost

---

Verification (ship-pr)

• PR base branch: litellm_oss_agent_shin_daily_branch ✓
• Customer name leak: none (LIT-3386 has no Linear project/customer)
• Greptile: 5/5 on commit de102db — "Safe to merge; the change narrows scope by delegating to an already-tested central helper, and the defensive copy-on-read guards against cross-iteration state bleed." Two P2 suggestions from Greptile's first pass addressed in de102db.
• Veria AI: pass — "No security issues found"
• GitHub Actions: 44/44 green (pytest, lint, code-quality, semgrep, secret-scan, validate-model-prices-json, all proxy-* suites, all-other-providers, vertex, integrations, enterprise-routing, build-ui, docs)
• Mergeable state: clean
• Unit tests local: 9/9 in tests/proxy_unit_tests/test_check_batch_cost.py; related test_batch_update_db_managed_output_file_id.py 8/8; broader -k "batch or managed_file" 14/14
• Runtime evidence: BEFORE (clean main 5699a06) + AFTER terminal output inlined above (persisted file_object.input_file_id flips from raw provider value to managed unified id; helper invocation count flips from 0 → 1)
• Diff scope: 2 files (check_batch_cost.py + the test file)

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR fixes CheckBatchCost's batch file ID reconciliation by delegating to ensure_batch_response_managed_file_ids — the same helper used by the HTTP retrieve/cancel path — instead of maintaining a separate partial reconciliation block that only handled output_file_id/error_file_id and missed input_file_id.

• check_batch_cost.py: Replaces the manual output/error file-ID registration loop with a single call to ensure_batch_response_managed_file_ids, adds copy-on-read for _hidden_params to avoid class-level dict aliasing across poll iterations, and emits a warning when _hidden_params assignment is not possible (e.g. __slots__).
• test_check_batch_cost.py: Updates the raw-output test to drive through the real helper path (removing fragile cross-module patches), and adds two new tests for input_file_id resolution and dedup behavior when a managed row already exists.

Confidence Score: 5/5

Safe to merge; the change narrows scope by delegating to an already-tested central helper, and the defensive copy-on-read guards against cross-iteration state bleed.

The diff replaces a narrow manual reconciliation block with a call to the same helper used on the HTTP path, so correctness follows from the helper's own test coverage. The two previous concerns (class-level dict mutation and silent degraded mode) are both addressed: _hidden_params is copied before mutation and a verbose_proxy_logger.warning is emitted on assignment failure. New tests drive real LiteLLMBatch objects through the full code path with stronger assertions than the previous mock-heavy approach.

No files require special attention.

Important Files Changed

Filename	Overview
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py	Replaces partial manual file-ID reconciliation with the central helper; correctly copies _hidden_params before mutation and warns on assignment failure.
tests/proxy_unit_tests/test_check_batch_cost.py	Existing test upgraded to use real LiteLLMBatch objects and stronger assertions; two new tests added for input_file_id resolution and managed-row dedup; no real network calls.

_{Reviews (2): Last reviewed commit: "fix(check_batch_cost): copy-on-read _hid..." | Re-trigger Greptile}

[oss-agent-shin]

Addressed both Greptile P2 suggestions in de102db:

1. _hidden_params now copy-on-read (dict(getattr(response, "_hidden_params", None) or {}) + reassign) — defensively protects against the class-level default-dict aliasing concern, so per-iteration model_id/model_name injection never bleeds into another LiteLLMBatch.
2. When the response disallows attribute assignment (e.g. __slots__), a verbose_proxy_logger.warning is now emitted so the degraded fallback is observable.

All 9 tests still pass (pytest tests/proxy_unit_tests/test_check_batch_cost.py → 9 passed in 0.39s). @greptile-apps re-review please.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs.