fix(batches): set the created endpoint on the batch object the hook saves and the logger reads

Retrieve, cancel and the cost poller now write the endpoint saved at create time onto the
provider's batch object itself instead of a copy. The managed files hook persists that object
and swaps its raw provider id for the unified id in place, and the spend logger reads the same
object, so a copy either let the hook save the provider's endpoint back or kept the raw id on
the inline cost row.

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -1163,7 +1163,7 @@ class CheckBatchCost:
         - if not, return False
         - if so, return True
         """
-        from litellm.proxy.openai_files_endpoints.common_utils import with_created_batch_endpoint
+        from litellm.proxy.openai_files_endpoints.common_utils import created_batch_endpoint
 
         try:
             from litellm.integrations.prometheus import PrometheusLogger
@@ -1226,7 +1226,7 @@ class CheckBatchCost:
             )
 
             try:
-                provider_response = await self.llm_router.aretrieve_batch(
+                response = await self.llm_router.aretrieve_batch(
                     model=model_id,
                     batch_id=batch_id,
                     litellm_metadata={
@@ -1245,7 +1245,7 @@ class CheckBatchCost:
                 else:
                     await self._expire_if_old_and_unfinished(job)
                 continue
-            response = with_created_batch_endpoint(provider_response, self._stored_batch(job))
+            response.endpoint = created_batch_endpoint(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

@@ -52,6 +52,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (  # noqa: F401  #
     apply_team_provider_credentials,
     authorize_model_for_key,
     batch_cost_poller_is_active,
+    created_batch_endpoint,
     decode_model_from_file_id,
     encode_batch_response_ids,
     encode_file_id_with_model,
@@ -67,7 +68,6 @@ 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
@@ -790,12 +790,13 @@ async def retrieve_batch(
                 **data,
             )
 
-        retrieved_batch: Final = with_created_batch_endpoint(response, stored_batch)
+        response.endpoint = 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=retrieved_batch,
+            response=response,
             managed_files_obj=managed_files_obj,
             prisma_client=prisma_client,
             verbose_proxy_logger=verbose_proxy_logger,
@@ -806,7 +807,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=retrieved_batch
+            data=data, user_api_key_dict=user_api_key_dict, response=response
         )
 
         # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id
@@ -1240,11 +1241,11 @@ async def cancel_batch(
             prisma_client=prisma_client,
             verbose_proxy_logger=verbose_proxy_logger,
         )
-        cancelled_batch: Final = with_created_batch_endpoint(response, stored_batch)
+        response.endpoint = created_batch_endpoint(response, stored_batch)
         await update_batch_in_database(
             batch_id=batch_id,
             unified_batch_id=unified_batch_id,
-            response=cancelled_batch,
+            response=response,
             managed_files_obj=managed_files_obj,
             prisma_client=prisma_client,
             verbose_proxy_logger=verbose_proxy_logger,
@@ -1255,7 +1256,7 @@ async def cancel_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=cancelled_batch
+            data=data, user_api_key_dict=user_api_key_dict, response=response
         )
 
         ### ALERTING ###

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -1607,10 +1607,8 @@ 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 created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> str:
+    return stored_batch.endpoint if stored_batch is not None and stored_batch.endpoint else provider_batch.endpoint
 
 
 def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:

tests/integration/spend/test_batch_observability.py

@@ -2,12 +2,16 @@ from __future__ import annotations
 
 import json
 import uuid
+from dataclasses import dataclass
 from hashlib import sha256
+from pathlib import Path
 from typing import Final
 
+import httpx
 import pytest
-from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
+from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
 from integration._support.database import read_rows
+from integration._support.process import owned_proxy
 from integration._support.upstream import delete_scenario, register_scenario
 from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
 from pydantic import JsonValue
@@ -59,11 +63,11 @@ def _failed_line(index: int) -> str:
     )
 
 
-def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]:
+def _batch(status: str, *, files_ready: bool, endpoint: str = "/v1/chat/completions") -> dict[str, JsonValue]:
     return {
         "id": "batch-$REQUEST_ID",
         "object": "batch",
-        "endpoint": "/v1/chat/completions",
+        "endpoint": endpoint,
         "errors": None,
         "input_file_id": "file-in-$REQUEST_ID",
         "completion_window": "24h",
@@ -79,7 +83,7 @@ def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]:
     }
 
 
-def _provider_routes() -> RoutedResponse:
+def _provider_routes(created_endpoint: str = "/v1/chat/completions") -> RoutedResponse:
     output_lines: Final = (
         _successful_line(1, REASONING_TOKENS[0]),
         _failed_line(2),
@@ -102,7 +106,8 @@ def _provider_routes() -> RoutedResponse:
                 },
             ),
             "POST /batches": JsonResponse(
-                content_type="application/json", body=_batch("validating", files_ready=False)
+                content_type="application/json",
+                body=_batch("validating", files_ready=False, endpoint=created_endpoint),
             ),
             "GET /batches/batch-$REQUEST_ID": JsonResponse(
                 content_type="application/json", body=_batch("completed", files_ready=True)
@@ -150,44 +155,53 @@ def _metadata(row: dict[str, JsonValue]) -> dict[str, JsonValue]:
     return object_value(JSON_OBJECT.validate_json(value) if isinstance(value, str) else value)
 
 
+@dataclass(frozen=True, slots=True)
+class CompletedBatch:
+    key: str
+    batch_id: str
+    retrieval: httpx.Response
+
+
+def _complete_a_batch(gateway: Gateway, scenario: Scenario, created_endpoint: str, *, managed: bool) -> CompletedBatch:
+    scenario_id: Final = f"batch-observability-{uuid.uuid4().hex[:12]}"
+    handle: Final = register_scenario(scenario_id, _provider_routes(created_endpoint))
+    scenario.cleanups.callback(delete_scenario, handle)
+    model_name: Final = scenario.model(api_base=handle.api_base())
+    key: Final = scenario.key(models=[model_name])
+    file_response: Final = gateway.request_multipart(
+        "/v1/files",
+        {"purpose": "batch", "target_model_names" if managed else "model": model_name},
+        {"file": ("in.jsonl", _input_file(model_name), "application/jsonl")},
+        key=key,
+    )
+    assert file_response.status_code == 200, file_response.text
+    input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"])
+    batch_response: Final = gateway.request(
+        "POST",
+        "/v1/batches",
+        {"input_file_id": input_file_id, "endpoint": created_endpoint, "completion_window": "24h"}
+        | ({} if managed else {"model": model_name}),
+        key=key,
+    )
+    assert batch_response.status_code == 200, batch_response.text
+    batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"])
+    retrieval: Final = eventually(
+        lambda: gateway.request("GET", f"/v1/batches/{batch_id}", key=key),
+        lambda response: response.status_code == 200 and response.json()["status"] == "completed",
+        seconds=30,
+    )
+    assert retrieval.status_code == 200, retrieval.text
+    return CompletedBatch(key=key, batch_id=batch_id, retrieval=retrieval)
+
+
 @pytest.mark.covers("spend.batches.retrieval_row_aggregates_reasoning_tokens_and_per_request_counts")
 def test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_file_failures(
     gateway: Gateway,
 ) -> None:
     with gateway.scenario() as scenario:
-        scenario_id: Final = f"batch-observability-{uuid.uuid4().hex[:12]}"
-        handle: Final = register_scenario(scenario_id, _provider_routes())
-        scenario.cleanups.callback(delete_scenario, handle)
-        model_name: Final = scenario.model(api_base=handle.api_base())
-        key: Final = scenario.key(models=[model_name])
-        file_response: Final = gateway.request_multipart(
-            "/v1/files",
-            {"purpose": "batch", "model": model_name},
-            {"file": ("in.jsonl", _input_file(model_name), "application/jsonl")},
-            key=key,
-        )
-        assert file_response.status_code == 200, file_response.text
-        input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"])
-        batch_response: Final = gateway.request(
-            "POST",
-            "/v1/batches",
-            {
-                "input_file_id": input_file_id,
-                "endpoint": "/v1/chat/completions",
-                "completion_window": "24h",
-                "model": model_name,
-            },
-            key=key,
-        )
-        assert batch_response.status_code == 200, batch_response.text
-        batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"])
-        retrieval: Final = eventually(
-            lambda: gateway.request("GET", f"/v1/batches/{batch_id}", key=key),
-            lambda response: response.status_code == 200 and response.json()["status"] == "completed",
-            seconds=30,
-        )
-        assert retrieval.status_code == 200, retrieval.text
-        rows: Final = eventually(lambda: _retrieval_rows(key), lambda values: len(values) == 1, seconds=70)
+        batch: Final = _complete_a_batch(gateway, scenario, "/v1/chat/completions", managed=False)
+        retrieval: Final = batch.retrieval
+        rows: Final = eventually(lambda: _retrieval_rows(batch.key), lambda values: len(values) == 1, seconds=70)
         row: Final = rows[0]
         metadata: Final = _metadata(row)
         usage: Final = object_value(metadata["usage_object"])
@@ -199,3 +213,19 @@ def test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_f
         )
         assert metadata["batch_successful_requests"] == 2, (retrieval.text, metadata)
         assert metadata["batch_failed_requests"] == 1 + ERROR_FILE_FAILURES, (retrieval.text, metadata)
+
+
+@pytest.mark.timeout(240)
+def test_managed_batch_keeps_its_created_endpoint_when_the_provider_reports_another(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with (
+        owned_proxy(gateway, tmp_path, {"PROXY_BATCH_POLLING_ENABLED": "false"}) as candidate,
+        candidate.scenario() as scenario,
+    ):
+        batch: Final = _complete_a_batch(candidate, scenario, "/v1/embeddings", managed=True)
+        stored: Final = candidate.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.key)
+        assert (batch.retrieval.json()["endpoint"], stored.json()["endpoint"]) == (
+            "/v1/embeddings",
+            "/v1/embeddings",
+        ), stored.text

tests/unit/proxy/batches_endpoints/test_endpoints.py

@@ -1909,20 +1909,48 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn
     assert retrieve_harness.update_batch_in_db.call_count == 1
 
 
+def _hook_rewriting_the_batch_id_to(batch_id):
+    def hook(**kwargs):
+        kwargs["response"].id = batch_id
+        return kwargs["response"]
+
+    return hook
+
+
 @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"})
+    provider_batch = make_batch(id="batch-raw-xyz", status="in_progress")
     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")
+    retrieve_harness.router.aretrieve_batch.return_value = provider_batch
+    retrieve_harness.logging.post_call_success_hook.side_effect = _hook_rewriting_the_batch_id_to("batch-unified-blob")
 
     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 retrieve_harness.logging.post_call_success_hook.call_args.kwargs["response"].endpoint == "/v1/embeddings"
+    assert (provider_batch.id, provider_batch.endpoint) == ("batch-unified-blob", "/v1/embeddings")
     assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID
 
 
+@pytest.mark.asyncio
+async def test_retrieve__a_saved_batch_without_an_endpoint_keeps_the_providers(retrieve_harness):
+    stored = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": ""})
+    provider_batch = make_batch(id="batch-raw-xyz", status="in_progress").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 = provider_batch
+
+    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 == "/v1/embeddings"
+    assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings"
+
+
 @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):
@@ -2695,7 +2723,9 @@ async def test_cancel__db_write_receives_caller_auth(cancel_harness):
 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")
+    provider_batch = make_batch(id="batch-raw-xyz", status="cancelling")
+    cancel_harness.router.acancel_batch.return_value = provider_batch
+    cancel_harness.logging.post_call_success_hook.side_effect = _hook_rewriting_the_batch_id_to("batch-unified-blob")
 
     with (
         patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID),
@@ -2706,6 +2736,8 @@ async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harne
     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)
+    assert cancel_harness.logging.post_call_success_hook.call_args.kwargs["response"].endpoint == "/v1/embeddings"
+    assert (provider_batch.id, provider_batch.endpoint) == ("batch-unified-blob", "/v1/embeddings")
 
 
 @pytest.mark.asyncio