fix(batches): keep the endpoint a batch was created for on retrieve, cancel and cost polling

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -662,6 +662,11 @@ class CheckBatchCost:
 
     @staticmethod
     def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
+        stored_batch: Final = CheckBatchCost._stored_batch(job)
+        return stored_batch.input_file_id if stored_batch is not None else None
+
+    @staticmethod
+    def _stored_batch(job: "ManagedObjectRow") -> Optional["LiteLLMBatch"]:
         from litellm.types.utils import LiteLLMBatch
 
         file_object: Final = job.file_object
@@ -674,7 +679,7 @@ class CheckBatchCost:
         if not isinstance(parsed_file_object, dict):
             return None
         try:
-            return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
+            return LiteLLMBatch.model_validate(parsed_file_object)
         except Exception:
             return None
 
@@ -1138,6 +1143,8 @@ class CheckBatchCost:
         - if not, return False
         - if so, return True
         """
+        from litellm.proxy.openai_files_endpoints.common_utils import with_created_batch_endpoint
+
         try:
             from litellm.integrations.prometheus import PrometheusLogger
 
@@ -1199,7 +1206,7 @@ class CheckBatchCost:
             )
 
             try:
-                response = await self.llm_router.aretrieve_batch(
+                provider_response = await self.llm_router.aretrieve_batch(
                     model=model_id,
                     batch_id=batch_id,
                     litellm_metadata={
@@ -1218,6 +1225,7 @@ class CheckBatchCost:
                 else:
                     await self._expire_if_old_and_unfinished(job)
                 continue
+            response = with_created_batch_endpoint(provider_response, self._stored_batch(job))
             if self._is_completed_with_output_id_pending(response):
                 verbose_proxy_logger.info(
                     f"CheckBatchCost: batch {batch_id} is completed but its output file id "

litellm/proxy/batches_endpoints/endpoints.py

@@ -67,6 +67,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (  # noqa: F401  #
     prepare_data_with_credentials,
     update_batch_in_database,
     validate_managed_id_requirement,
+    with_created_batch_endpoint,
 )
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_request_data
 from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
@@ -644,6 +645,7 @@ async def retrieve_batch(
             proxy_logging_obj=proxy_logging_obj,
             user_api_key_dict=user_api_key_dict,
         )
+        stored_batch: Final = response
 
         if executed_batch and response is None:
             raise batch_error(404, f"No batch found with id '{batch_id}'.")
@@ -788,11 +790,12 @@ async def retrieve_batch(
                 **data,
             )
 
+        retrieved_batch: Final = with_created_batch_endpoint(response, stored_batch)
         # FIX: Update the database with the latest state from provider
         await update_batch_in_database(
             batch_id=batch_id,
             unified_batch_id=unified_batch_id,
-            response=response,
+            response=retrieved_batch,
             managed_files_obj=managed_files_obj,
             prisma_client=prisma_client,
             verbose_proxy_logger=verbose_proxy_logger,
@@ -803,7 +806,7 @@ async def retrieve_batch(
 
         ### CALL HOOKS ### - modify outgoing data
         response = await proxy_logging_obj.post_call_success_hook(
-            data=data, user_api_key_dict=user_api_key_dict, response=response
+            data=data, user_api_key_dict=user_api_key_dict, response=retrieved_batch
         )
 
         # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id
@@ -1230,20 +1233,29 @@ async def cancel_batch(
         managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
         from litellm.proxy.proxy_server import prisma_client
 
-        await update_batch_in_database(
+        db_batch_object, stored_batch = await get_batch_from_database(
             batch_id=batch_id,
             unified_batch_id=unified_batch_id,
-            response=response,
             managed_files_obj=managed_files_obj,
             prisma_client=prisma_client,
             verbose_proxy_logger=verbose_proxy_logger,
+        )
+        cancelled_batch: Final = with_created_batch_endpoint(response, stored_batch)
+        await update_batch_in_database(
+            batch_id=batch_id,
+            unified_batch_id=unified_batch_id,
+            response=cancelled_batch,
+            managed_files_obj=managed_files_obj,
+            prisma_client=prisma_client,
+            verbose_proxy_logger=verbose_proxy_logger,
+            db_batch_object=db_batch_object,
             operation="cancel",
             user_api_key_dict=user_api_key_dict,
         )
 
         ### CALL HOOKS ### - modify outgoing data
         response = await proxy_logging_obj.post_call_success_hook(
-            data=data, user_api_key_dict=user_api_key_dict, response=response
+            data=data, user_api_key_dict=user_api_key_dict, response=cancelled_batch
         )
 
         ### ALERTING ###

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1593,6 +1593,12 @@ def batch_cost_poller_is_active() -> bool:
         return False
 
 
+def with_created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> "LiteLLMBatch":
+    if stored_batch is None or not stored_batch.endpoint or stored_batch.endpoint == provider_batch.endpoint:
+        return provider_batch
+    return provider_batch.model_copy(update={"endpoint": stored_batch.endpoint})
+
+
 def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
     """Whether a "completed" batch may be retired from cost recovery.
 

tests/unit/proxy/batches_endpoints/test_endpoints.py

@@ -1878,6 +1878,20 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn
     assert retrieve_harness.update_batch_in_db.call_count == 1
 
 
+@pytest.mark.asyncio
+async def test_retrieve__provider_sync_keeps_the_endpoint_the_batch_was_created_for(retrieve_harness):
+    stored = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": "/v1/embeddings"})
+    retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored)
+    retrieve_harness.router.aretrieve_batch.return_value = make_batch(id="batch-raw-xyz", status="in_progress")
+
+    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
+        resp = await call_retrieve(retrieve_harness, "batch-unified-blob")
+
+    assert (resp.endpoint, resp.status) == ("/v1/embeddings", "in_progress")
+    assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings"
+    assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
 async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status):
@@ -2632,6 +2646,23 @@ async def test_cancel__db_write_receives_caller_auth(cancel_harness):
     assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller
 
 
+@pytest.mark.asyncio
+async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harness):
+    stored_row = MagicMock(status="in_progress")
+    stored = make_batch(id="batch-unified-blob", status="in_progress").model_copy(update={"endpoint": "/v1/embeddings"})
+    cancel_harness.router.acancel_batch.return_value = make_batch(id="batch-raw-xyz", status="cancelling")
+
+    with (
+        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID),
+        patch.object(endpoints, "get_batch_from_database", AsyncMock(return_value=(stored_row, stored))),
+    ):
+        resp = await call_cancel(cancel_harness, "batch-unified-blob")
+
+    assert (resp.endpoint, resp.status) == ("/v1/embeddings", "cancelling")
+    written = cancel_harness.update_batch_in_db.call_args.kwargs
+    assert (written["response"].endpoint, written["db_batch_object"]) == ("/v1/embeddings", stored_row)
+
+
 @pytest.mark.asyncio
 async def test_cancel__unified_missing_model_id_400(cancel_harness):
     # unified id with no model_id segment -> get_model_id returns None -> 400.

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -3784,6 +3784,20 @@ class TestMultiPodBatchCostClaim:
         assert row.batch_processed is True
         await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID)
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("has_batch_processed_column", (True, False))
+    async def test_billing_a_batch_keeps_the_endpoint_it_was_created_for(self, has_batch_processed_column: bool) -> None:
+        row: Final = _FakeManagedObjectRow()
+        row.file_object = json.dumps({**json.loads(row.file_object), "endpoint": "/v1/embeddings"})
+        instance: Final = self._instance(self._prisma(row, []), self._router())
+        instance._has_batch_processed_column = has_batch_processed_column
+
+        with self._billing_patches([]):
+            await instance.check_batch_cost()
+
+        assert row.status == "complete"
+        assert json.loads(row.file_object)["endpoint"] == "/v1/embeddings"
+
     @pytest.mark.asyncio
     async def test_schema_without_batch_processed_still_bills(self):
         """Older schemas have no column to claim, so they keep the pre-fix behavior instead