Merge pull request #34192 from BerriAI/litellm_fix_batch_list_pagination_lit4678

fix(batches): paginate managed batch list by unified_object_id cursor

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -316,26 +316,34 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter}
 
         if after:
-            where_clause["id"] = {"gt": after}
+            cursor_row = (
+                await self.prisma_client.db.litellm_managedobjecttable.find_first(
+                    where={**where_clause, "unified_object_id": after}
+                )
+            )
+            if cursor_row is None:
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Invalid 'after' cursor: no batch found with id '{after}'.",
+                )
 
-        fetch_limit = limit or 20
-        if target_model_names:
-            # Oversample so post-fetch model-name filtering still has enough rows.
-            fetch_limit = max(fetch_limit * 3, 100)
+        page_size = limit or 20
+        cursor_args: Dict[str, Any] = (
+            {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
+        )
 
         batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
             where=where_clause,
-            take=fetch_limit,
-            order={"created_at": "desc"},
+            take=page_size + 1,
+            order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
+            **cursor_args,
         )
 
-        batch_objects: List[LiteLLMBatch] = []
-        for batch in batches:
-            try:
-                # Stop once we have enough after filtering
-                if len(batch_objects) >= (limit or 20):
-                    break
+        has_more = len(batches) > page_size
 
+        batch_objects: List[LiteLLMBatch] = []
+        for batch in batches[:page_size]:
+            try:
                 batch_data = (
                     json.loads(batch.file_object)
                     if isinstance(batch.file_object, str)
@@ -351,9 +359,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 )
                 continue
 
-        return build_list_page(
-            batch_objects, has_more=len(batch_objects) == (limit or 20)
-        )
+        return build_list_page(batch_objects, has_more=has_more)
 
     async def get_user_created_file_ids(
         self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str]

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

@@ -1771,8 +1771,8 @@ async def test_list_batches_from_managed_objects_table():
     # Should filter by user_id (created_by)
     prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
         where={"file_purpose": "batch", "created_by": "test-user"},
-        take=10,
-        order={"created_at": "desc"},
+        take=11,
+        order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
     )
 
 
@@ -1801,8 +1801,8 @@ async def test_list_batches_from_managed_objects_table_empty_list():
     # Default take is 20 when no limit is provided
     prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
         where={"file_purpose": "batch", "created_by": "test-user"},
-        take=20,
-        order={"created_at": "desc"},
+        take=21,
+        order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
     )
 
 
@@ -1918,8 +1918,8 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by():
     assert result_user1["data"][0].id == "unified-batch-user1"
     prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with(
         where={"file_purpose": "batch", "created_by": "user1"},
-        take=10,
-        order={"created_at": "desc"},
+        take=11,
+        order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
     )
 
     # Query with user2's API key - should only return user2's batch
@@ -1933,11 +1933,505 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by():
     assert result_user2["data"][0].id == "unified-batch-user2"
     prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with(
         where={"file_purpose": "batch", "created_by": "user2"},
-        take=10,
-        order={"created_at": "desc"},
+        take=11,
+        order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
     )
 
 
+@pytest.mark.asyncio
+async def test_list_batches_pagination_uses_unified_object_id_cursor():
+    """Regression for LIT-4678.
+
+    The ``after`` cursor a client sends back is a batch's ``unified_object_id``
+    (that is what is returned as ``.id`` / ``last_id``). Paginating must use a
+    Prisma cursor on the unique ``unified_object_id`` column, not a
+    ``where id > after`` filter against the random-uuid primary key.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.find_first.return_value = MagicMock()
+    prisma_client.db.litellm_managedobjecttable.find_many.return_value = []
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    await proxy_managed_files.list_user_batches(
+        user_api_key_dict=UserAPIKeyAuth(user_id="test-user"),
+        limit=5,
+        after="unified-batch-id-7",
+    )
+
+    prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
+        where={"file_purpose": "batch", "created_by": "test-user"},
+        take=6,
+        order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
+        cursor={"unified_object_id": "unified-batch-id-7"},
+        skip=1,
+    )
+
+    _, call_kwargs = prisma_client.db.litellm_managedobjecttable.find_many.call_args
+    assert "id" not in call_kwargs["where"]
+
+
+@pytest.mark.asyncio
+async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps():
+    """Regression for LIT-4678.
+
+    Simulates the managed-objects table (random-uuid ``id`` primary key,
+    base64 ``unified_object_id``, reverse-chronological ``created_at``) and
+    walks every page the way a client would, feeding ``last_id`` back as
+    ``after``. With the old ``where id > after`` cursor this loops and drops
+    batches; the fixed cursor returns each batch exactly once, newest first.
+    """
+    import uuid as _uuid
+
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    def _unified_id(i: int) -> str:
+        raw = f"litellm_proxy;model_id:gpt-4o-batch;llm_batch_id:batch_{i:03d}"
+        return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
+
+    total = 10
+    rows = []
+    for i in range(total):
+        row = MagicMock()
+        row.id = str(_uuid.uuid4())
+        row.unified_object_id = _unified_id(i)
+        row.created_at = 1_000_000 + i
+        row.file_object = json.dumps(
+            {
+                "id": f"batch_provider_{i:03d}",
+                "object": "batch",
+                "endpoint": "/v1/chat/completions",
+                "completion_window": "24h",
+                "status": "completed",
+                "created_at": 1_000_000 + i,
+                "input_file_id": f"file-input-{i:03d}",
+                "request_counts": {"total": 1, "completed": 1, "failed": 0},
+            }
+        )
+        rows.append(row)
+
+    async def fake_find_many(where, take, order, cursor=None, skip=0):
+        result = list(rows)
+        id_filter = where.get("id")
+        if isinstance(id_filter, dict) and "gt" in id_filter:
+            result = [r for r in result if r.id > id_filter["gt"]]
+        order_keys = order if isinstance(order, list) else [order]
+        for clause in reversed(order_keys):
+            (order_field, direction), = clause.items()
+            result.sort(
+                key=lambda r: getattr(r, order_field), reverse=(direction == "desc")
+            )
+        if cursor is not None:
+            (cur_field, cur_val), = cursor.items()
+            idx = next(
+                (i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val),
+                None,
+            )
+            if idx is None:
+                return []
+            result = result[idx + skip:]
+        return result[:take]
+
+    async def fake_find_first(where):
+        return next(
+            (r for r in rows if r.unified_object_id == where.get("unified_object_id")),
+            None,
+        )
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+        side_effect=fake_find_many
+    )
+    prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+        side_effect=fake_find_first
+    )
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+    user = UserAPIKeyAuth(user_id="test-user")
+
+    seen: list = []
+    after = None
+    for _ in range(total + 5):
+        resp = await proxy_managed_files.list_user_batches(
+            user_api_key_dict=user, limit=3, after=after
+        )
+        page_ids = [b.id for b in resp["data"]]
+        seen.extend(page_ids)
+        if not resp["has_more"]:
+            break
+        assert page_ids, "has_more was true but the page was empty"
+        assert resp["last_id"] != after, "cursor did not advance (pagination loop)"
+        after = resp["last_id"]
+
+    expected = [_unified_id(i) for i in reversed(range(total))]
+    assert seen == expected
+    assert len(seen) == len(set(seen))
+
+
+@pytest.mark.asyncio
+async def test_list_batches_pagination_stable_when_created_at_ties():
+    """Regression for LIT-4678.
+
+    Cursor pagination is only well-defined when the ``order`` fully determines
+    row order. If listing ordered by non-unique ``created_at`` alone, batches
+    sharing a timestamp come back in an arbitrary order that can shift between
+    page requests, so a cursor row's neighbours change and batches get skipped
+    or duplicated. Listing must add the unique ``unified_object_id`` as a
+    tie-breaker so the order is total and pagination is stable.
+    """
+    import itertools
+    import uuid as _uuid
+
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    def _unified_id(i: int) -> str:
+        raw = f"litellm_proxy;model_id:gpt-4o-batch;llm_batch_id:batch_{i:03d}"
+        return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
+
+    total = 6
+    shared_created_at = 1_000_000
+    rows = []
+    for i in range(total):
+        row = MagicMock()
+        row.id = str(_uuid.uuid4())
+        row.unified_object_id = _unified_id(i)
+        row.created_at = shared_created_at
+        row.file_object = json.dumps(
+            {
+                "id": f"batch_provider_{i:03d}",
+                "object": "batch",
+                "endpoint": "/v1/chat/completions",
+                "completion_window": "24h",
+                "status": "completed",
+                "created_at": shared_created_at,
+                "input_file_id": f"file-input-{i:03d}",
+                "request_counts": {"total": 1, "completed": 1, "failed": 0},
+            }
+        )
+        rows.append(row)
+
+    call_counter = itertools.count()
+
+    async def fake_find_many(where, take, order, cursor=None, skip=0):
+        call = next(call_counter)
+        order_keys = order if isinstance(order, list) else [order]
+        fields = [next(iter(clause)) for clause in order_keys]
+        result = list(rows)
+        for clause in reversed(order_keys):
+            field, direction = next(iter(clause.items()))
+            result.sort(
+                key=lambda r: getattr(r, field), reverse=(direction == "desc")
+            )
+
+        def order_key(r):
+            return tuple(getattr(r, f) for f in fields)
+
+        stabilized = []
+        i = 0
+        while i < len(result):
+            j = i
+            while j < len(result) and order_key(result[j]) == order_key(result[i]):
+                j += 1
+            group = result[i:j]
+            if len(group) > 1:
+                rot = call % len(group)
+                group = group[rot:] + group[:rot]
+            stabilized.extend(group)
+            i = j
+        result = stabilized
+
+        if cursor is not None:
+            cur_field, cur_val = next(iter(cursor.items()))
+            idx = next(
+                (k for k, r in enumerate(result) if getattr(r, cur_field) == cur_val),
+                None,
+            )
+            if idx is None:
+                return []
+            result = result[idx + skip:]
+        return result[:take]
+
+    async def fake_find_first(where):
+        return next(
+            (r for r in rows if r.unified_object_id == where.get("unified_object_id")),
+            None,
+        )
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+        side_effect=fake_find_many
+    )
+    prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+        side_effect=fake_find_first
+    )
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+    user = UserAPIKeyAuth(user_id="test-user")
+
+    seen: list = []
+    after = None
+    for _ in range(total + 5):
+        resp = await proxy_managed_files.list_user_batches(
+            user_api_key_dict=user, limit=2, after=after
+        )
+        page_ids = [b.id for b in resp["data"]]
+        seen.extend(page_ids)
+        if not resp["has_more"]:
+            break
+        assert page_ids, "has_more was true but the page was empty"
+        assert resp["last_id"] != after, "cursor did not advance (pagination loop)"
+        after = resp["last_id"]
+
+    assert sorted(seen) == sorted(_unified_id(i) for i in range(total))
+    assert len(seen) == len(set(seen)), "a tied batch was returned more than once"
+
+
+def _managed_batch_row(index, file_object=None):
+    row = MagicMock()
+    row.id = f"pk-{index:03d}"
+    raw = f"litellm_proxy;model_id:gpt-4o-batch;llm_batch_id:batch_{index:03d}"
+    row.unified_object_id = base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
+    row.created_at = 1_000_000 + index
+    row.file_object = (
+        file_object
+        if file_object is not None
+        else json.dumps(
+            {
+                "id": f"batch_provider_{index:03d}",
+                "object": "batch",
+                "endpoint": "/v1/chat/completions",
+                "completion_window": "24h",
+                "status": "completed",
+                "created_at": 1_000_000 + index,
+                "input_file_id": f"file-input-{index:03d}",
+                "request_counts": {"total": 1, "completed": 1, "failed": 0},
+            }
+        )
+    )
+    return row
+
+
+def _fake_managed_object_table(rows):
+    async def find_many(where, take, order, cursor=None, skip=0):
+        result = sorted(
+            rows, key=lambda r: (r.created_at, r.unified_object_id), reverse=True
+        )
+        if cursor is not None:
+            (cur_field, cur_val), = cursor.items()
+            idx = next(
+                (i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val),
+                None,
+            )
+            if idx is None:
+                return []
+            result = result[idx + skip:]
+        return result[:take]
+
+    async def find_first(where):
+        return next(
+            (r for r in rows if r.unified_object_id == where.get("unified_object_id")),
+            None,
+        )
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+        side_effect=find_many
+    )
+    prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+        side_effect=find_first
+    )
+    return prisma_client
+
+
+async def _walk_batch_pages(proxy_managed_files, user, limit, max_pages=20):
+    pages = []
+    after = None
+    for _ in range(max_pages):
+        resp = await proxy_managed_files.list_user_batches(
+            user_api_key_dict=user, limit=limit, after=after
+        )
+        pages.append(resp)
+        if not resp["has_more"]:
+            break
+        assert resp["last_id"] is not None, "has_more was true but there is no cursor"
+        after = resp["last_id"]
+    return pages
+
+
+@pytest.mark.asyncio
+async def test_list_batches_rejects_unknown_after_cursor():
+    """An ``after`` that does not resolve to a batch the caller can see is a
+    client error, not an empty page.
+
+    Returning ``[]`` for an unresolvable cursor is indistinguishable from
+    "you have reached the end of the list", so a client walking pages with a
+    stale or malformed cursor silently sees a truncated batch list instead of
+    an error it can act on.
+    """
+    from fastapi import HTTPException
+
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+        return_value=None
+    )
+    prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[])
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await proxy_managed_files.list_user_batches(
+            user_api_key_dict=UserAPIKeyAuth(user_id="test-user"),
+            limit=3,
+            after="does-not-exist-xyz",
+        )
+
+    assert exc_info.value.status_code == 400
+    assert "does-not-exist-xyz" in str(exc_info.value.detail)
+    prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called()
+    prisma_client.db.litellm_managedobjecttable.find_first.assert_called_once_with(
+        where={
+            "file_purpose": "batch",
+            "created_by": "test-user",
+            "unified_object_id": "does-not-exist-xyz",
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_list_batches_treats_empty_after_as_no_cursor():
+    """``?after=`` means "start from the beginning", as it always has.
+
+    Only a cursor the client actually sent is validated, so an SDK that always
+    emits the query parameter does not get a 400 on its first page.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    rows = [_managed_batch_row(i) for i in range(2)]
+    prisma_client = _fake_managed_object_table(rows)
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    page = await proxy_managed_files.list_user_batches(
+        user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2, after=""
+    )
+
+    assert [batch.id for batch in page["data"]] == [
+        rows[1].unified_object_id,
+        rows[0].unified_object_id,
+    ]
+    prisma_client.db.litellm_managedobjecttable.find_first.assert_not_called()
+    _, call_kwargs = prisma_client.db.litellm_managedobjecttable.find_many.call_args
+    assert "cursor" not in call_kwargs
+
+
+@pytest.mark.asyncio
+async def test_list_batches_rejects_after_cursor_owned_by_another_user():
+    """The cursor lookup must be scoped to the rows the caller can list.
+
+    A Prisma cursor resolves by unique column regardless of the ``where``
+    filter, so an unscoped cursor would let one user anchor their page window
+    to another user's batch and learn when it was created.
+    """
+    from fastapi import HTTPException
+
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    other_users_batch = _managed_batch_row(0)
+
+    async def find_first(where):
+        if where.get("created_by") != "user-b":
+            return None
+        return other_users_batch
+
+    prisma_client = AsyncMock()
+    prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
+        side_effect=find_first
+    )
+    prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[])
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await proxy_managed_files.list_user_batches(
+            user_api_key_dict=UserAPIKeyAuth(user_id="user-a"),
+            limit=3,
+            after=other_users_batch.unified_object_id,
+        )
+
+    assert exc_info.value.status_code == 400
+    prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_list_batches_has_more_false_on_exactly_full_final_page():
+    """``has_more`` must mean "another row exists", not "this page is full".
+
+    With a batch count that is an exact multiple of ``limit``, reporting
+    ``has_more`` off page fullness makes every client fetch one extra empty
+    page before it can stop.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    rows = [_managed_batch_row(i) for i in range(4)]
+    prisma_client = _fake_managed_object_table(rows)
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    pages = await _walk_batch_pages(
+        proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2
+    )
+
+    seen = [batch.id for page in pages for batch in page["data"]]
+    assert seen == [r.unified_object_id for r in reversed(rows)]
+    assert [page["has_more"] for page in pages] == [True, False]
+    assert prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_list_batches_unparseable_row_does_not_truncate_pagination():
+    """A row that fails to parse must not end pagination early.
+
+    Skipping a corrupt row shortens the page, so deriving ``has_more`` from
+    the number of returned batches reports "no more results" while older
+    batches are still unread, silently hiding them from the caller.
+    """
+    from litellm.proxy._types import UserAPIKeyAuth
+
+    rows = [_managed_batch_row(i) for i in range(4)]
+    rows[2].file_object = "{ not valid json"
+    prisma_client = _fake_managed_object_table(rows)
+
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    pages = await _walk_batch_pages(
+        proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2
+    )
+
+    seen = [batch.id for page in pages for batch in page["data"]]
+    assert seen == [rows[3].unified_object_id, rows[1].unified_object_id, rows[0].unified_object_id]
+    assert len(seen) == len(set(seen))
+
+
 @pytest.mark.asyncio
 async def test_return_unified_file_id_includes_expires_at():
     from litellm.types.llms.openai import OpenAIFileObject
@@ -2312,8 +2806,8 @@ async def test_list_batches_only_returns_user_own_batches():
     # Verify the database query filtered by user_id
     prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
         where={"file_purpose": "batch", "created_by": "user_a_id"},
-        take=10,
-        order={"created_at": "desc"},
+        take=11,
+        order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
     )