fix(batches): register ownership for every batch create path (#39810)

* 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.

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -46,6 +46,7 @@ from litellm.proxy._types import (
     UserAPIKeyAuth,
 )
 from litellm.proxy.openai_files_endpoints.common_utils import (
+    BATCH_CREATE_HIDDEN_PARAM,
     FILE_LIST_CONTINUATION_CHUNK_SIZE,
     MAX_FILE_LIST_LIMIT,
     _is_base64_encoded_unified_file_id,
@@ -1321,7 +1322,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             ## Check if unified_file_id is in the response
             unified_file_id = response._hidden_params.get("unified_file_id")  # managed file id
             unified_batch_id = response._hidden_params.get("unified_batch_id")  # managed batch id
-            is_batch_create: Final = unified_file_id is not None
+            is_batch_create: Final = response._hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True
             model_id = cast(Optional[str], response._hidden_params.get("model_id"))
             model_name = cast(Optional[str], response._hidden_params.get("model_name"))
 
@@ -1410,10 +1411,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
                 )
 
-            # Only record batch creation metric on actual create (not retrieve/cancel).
-            # unified_file_id in _hidden_params is only set by the create_batch endpoint.
-            original_unified_file_id = response._hidden_params.get("unified_file_id")
-            if original_unified_file_id:
+            if is_batch_create:
                 prom_logger = self._get_prometheus_logger()
                 if prom_logger:
                     batch_provider = ""

litellm/proxy/batches_endpoints/endpoints.py

@@ -25,6 +25,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
     get_custom_llm_provider_from_request_query,
 )
 from litellm.proxy.openai_files_endpoints.common_utils import (
+    BATCH_CREATE_HIDDEN_PARAM,
     _is_base64_encoded_unified_file_id,
     add_internal_model_credentials,
     apply_team_provider_credentials,
@@ -347,6 +348,8 @@ async def create_batch(
                     **_create_batch_data,
                 )
 
+        response._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True
+
         ### 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

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -37,6 +37,8 @@ MAX_FILE_LIST_LIMIT: Final = 10000
 
 FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500
 
+BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create"
+
 
 def validate_file_list_limit(limit: int | None) -> None:
     """Reject a ``limit`` outside the range OpenAI documents for GET /v1/files."""

tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py

@@ -10,6 +10,7 @@ from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFi
 from litellm.caching import DualCache
 from litellm.proxy._types import CallTypes
 from litellm.proxy.openai_files_endpoints.common_utils import (
+    BATCH_CREATE_HIDDEN_PARAM,
     _is_base64_encoded_unified_file_id,
     encode_file_id_with_model,
 )
@@ -3185,7 +3186,7 @@ def _batch_response(batch_id, output_file_id=None, is_create=False):
         output_file_id=output_file_id,
     )
     if is_create:
-        batch._hidden_params["unified_file_id"] = "unified-input-file-id"
+        batch._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True
     return batch
 
 
@@ -3411,11 +3412,8 @@ async def test_provider_format_file_without_ownership_row_stays_accessible():
 
 
 @pytest.mark.asyncio
-async def test_post_call_batch_create_stores_ownership_row():
-    """
-    Batch creation (response hidden params carry the unified input file id)
-    must write an ownership row attributed to the creating key.
-    """
+@pytest.mark.parametrize("batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID])
+async def test_post_call_batch_create_stores_ownership_row(batch_id):
     from litellm.proxy._types import UserAPIKeyAuth
 
     prisma_client = AsyncMock()
@@ -3432,13 +3430,11 @@ async def test_post_call_batch_create_stores_ownership_row():
         user_api_key_dict=UserAPIKeyAuth(
             user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()
         ),
-        response=_batch_response(MODEL_ENCODED_BATCH_ID, is_create=True),
+        response=_batch_response(batch_id, is_create=True),
     )
 
     upsert_call = prisma_client.db.litellm_managedobjecttable.upsert.await_args
-    assert upsert_call.kwargs["where"] == {
-        "unified_object_id": MODEL_ENCODED_BATCH_ID
-    }
+    assert upsert_call.kwargs["where"] == {"unified_object_id": batch_id}
     create_data = upsert_call.kwargs["data"]["create"]
     assert create_data["created_by"] == "user_a"
     assert create_data["team_id"] == "team_a"

tests/test_litellm/enterprise/proxy/test_managed_files_hook.py

@@ -15,6 +15,7 @@ from typing import Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
+from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM
 from litellm.types.llms.openai import FileListPage, OpenAIFileObject
 from litellm.types.utils import LiteLLMBatch
 
@@ -1540,6 +1541,11 @@ async def test_batch_create_hook_persists_creating_key_and_tags():
     managed_files = _make_managed_files_instance()
     creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None)
     create_response = _make_batch_response(status="validating", output_file_id=None)
+    create_response._hidden_params = {
+        BATCH_CREATE_HIDDEN_PARAM: True,
+        "model_id": "model-deploy-xyz",
+        "model_name": "azure/gpt-4",
+    }
 
     await managed_files.async_post_call_success_hook(
         data={"litellm_metadata": {"tags": ["env:prod", "team:ml"], "user_api_key": creator.api_key}},
@@ -1554,6 +1560,52 @@ async def test_batch_create_hook_persists_creating_key_and_tags():
     assert stored["user_api_key_dict"] is creator
 
 
+@pytest.mark.asyncio
+async def test_batch_create_hook_records_created_metric_once():
+    managed_files = _make_managed_files_instance()
+    prometheus_logger = MagicMock()
+    managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger)
+    create_response = _make_batch_response(status="validating", output_file_id=None)
+    create_response._hidden_params = {
+        BATCH_CREATE_HIDDEN_PARAM: True,
+        "model_id": "model-deploy-xyz",
+        "model_name": "azure/gpt-4",
+    }
+
+    await managed_files.async_post_call_success_hook(
+        data={},
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None),
+        response=create_response,
+    )
+
+    prometheus_logger.record_managed_batch_created.assert_called_once()
+    recorded = prometheus_logger.record_managed_batch_created.call_args.kwargs
+    assert recorded["model"] == "azure/gpt-4"
+    assert recorded["api_provider"] == "azure"
+    assert recorded["user"] == "alice"
+
+
+@pytest.mark.asyncio
+async def test_batch_retrieve_hook_does_not_record_created_metric():
+    managed_files = _make_managed_files_instance()
+    prometheus_logger = MagicMock()
+    managed_files._get_prometheus_logger = MagicMock(return_value=prometheus_logger)
+    retrieve_response = _make_batch_response(status="in_progress", output_file_id=None)
+    retrieve_response._hidden_params = {
+        "unified_batch_id": "some-unified-batch-id",
+        "model_id": "model-deploy-xyz",
+        "model_name": "azure/gpt-4",
+    }
+
+    await managed_files.async_post_call_success_hook(
+        data={},
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None),
+        response=retrieve_response,
+    )
+
+    prometheus_logger.record_managed_batch_created.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_batch_retrieve_hook_does_not_claim_attribution():
     """A retrieve carries unified_batch_id but no unified_file_id, so it must not rewrite

tests/test_litellm/proxy/batches_endpoints/test_endpoints.py

@@ -29,6 +29,7 @@ added to this layer raises instead of silently passing - the inventory of seams
 cannot drift without a test failure.
 """
 
+import base64
 import json
 from contextlib import ExitStack
 from dataclasses import dataclass
@@ -36,7 +37,7 @@ from typing import Any, Dict, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-
+from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
 
 import litellm
 import litellm.proxy.batches_endpoints.endpoints as endpoints
@@ -989,6 +990,67 @@ async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds):
     assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch"
 
 
+def install_managed_files_hook(harness: Harness) -> AsyncMock:
+    prisma_client = AsyncMock()
+    managed_files = _PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client)
+    harness.logging.post_call_success_hook = AsyncMock(side_effect=managed_files.async_post_call_success_hook)
+    harness.router.model_list = []
+    return prisma_client
+
+
+TEAM_A_KEY = UserAPIKeyAuth(api_key="sk-team-a", user_id="user_a", team_id="team_a")
+
+
+def assert_ownership_registered_for_team_a(prisma_client: AsyncMock, batch_id: str) -> None:
+    upsert = prisma_client.db.litellm_managedobjecttable.upsert
+    upsert.assert_awaited_once()
+    assert upsert.await_args.kwargs["where"] == {"unified_object_id": batch_id}
+    created = upsert.await_args.kwargs["data"]["create"]
+    assert created["created_by"] == "user_a"
+    assert created["team_id"] == "team_a"
+    prisma_client.db.litellm_managedobjecttable.update_many.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "body",
+    [
+        {"input_file_id": AZURE_FILE_ID},
+        {"input_file_id": "file-plain", "model": "vertex-model"},
+        {"input_file_id": "file-plain"},
+    ],
+    ids=["model_encoded_file_id", "model_param", "provider_fallback"],
+)
+async def test_create__registers_ownership_for_creator(harness, openai_env_creds, body):
+    set_body(harness, {**body, "endpoint": "/v1/chat/completions", "completion_window": "24h"})
+    prisma_client = install_managed_files_hook(harness)
+
+    resp = await call_create(harness, user=TEAM_A_KEY)
+
+    assert_ownership_registered_for_team_a(prisma_client, resp.id)
+
+
+@pytest.mark.asyncio
+async def test_create__unified_file_id_registers_ownership_for_creator(harness):
+    unified_input_file_id = base64.urlsafe_b64encode(
+        b"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,gpt-4o-mini"
+    ).decode()
+    set_body(
+        harness,
+        {
+            "input_file_id": unified_input_file_id,
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+        },
+    )
+    prisma_client = install_managed_files_hook(harness)
+
+    resp = await call_create(harness, user=TEAM_A_KEY)
+
+    assert harness.router_acreate.call_count == 1
+    assert_ownership_registered_for_team_a(prisma_client, resp.id)
+
+
 @pytest.mark.asyncio
 async def test_create__metadata_sanitized_before_forwarding(harness, openai_env_creds):
     set_body(