From d8c7e2529655a0a8144973e81cb1b055174433ec Mon Sep 17 00:00:00 2001 From: shivam Date: Tue, 21 Jul 2026 23:25:26 +0000 Subject: [PATCH 1/4] fix(batches): paginate managed batch list by unified_object_id cursor GET /batches served from the managed-objects table paged with a where id > after filter, but the after cursor clients send back is a batch's unified_object_id (the value returned as .id and last_id), and id is the table's random-uuid primary key. Comparing the two unrelated fields, while ordering by created_at desc but filtering with gt, made pages repeat the same last_id (pagination loops) and silently drop batches. Switch to Prisma cursor pagination on the unique unified_object_id column so listing walks every batch exactly once in reverse-chronological order, matching OpenAI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/managed_files.py | 8 +- .../proxy/hooks/test_managed_files.py | 123 ++++++++++++++++++ 2 files changed, 128 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 3f42867d90e..bbab80eb8a5 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -315,18 +315,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} - if after: - where_clause["id"] = {"gt": 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) + 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"}, + **cursor_args, ) batch_objects: List[LiteLLMBatch] = [] diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 68dc3269f34..e081ffaf016 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1938,6 +1938,129 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by(): ) +@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_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=5, + order={"created_at": "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_field, direction), = order.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] + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + side_effect=fake_find_many + ) + + 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"]] + if not page_ids: + break + seen.extend(page_ids) + 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_return_unified_file_id_includes_expires_at(): from litellm.types.llms.openai import OpenAIFileObject From 93f27641ae2ddb79f2ceafe2554e413b8d894e9b Mon Sep 17 00:00:00 2001 From: shivam Date: Tue, 21 Jul 2026 23:41:16 +0000 Subject: [PATCH 2/4] fix(batches): stabilize managed batch pagination with unified_object_id tie-breaker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/managed_files.py | 2 +- .../proxy/hooks/test_managed_files.py | 132 ++++++++++++++++-- 2 files changed, 123 insertions(+), 11 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bbab80eb8a5..cf5c2b0905d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -327,7 +327,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( where=where_clause, take=fetch_limit, - order={"created_at": "desc"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], **cursor_args, ) diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index e081ffaf016..347a8fcd023 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1772,7 +1772,7 @@ async def test_list_batches_from_managed_objects_table(): 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"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1802,7 +1802,7 @@ async def test_list_batches_from_managed_objects_table_empty_list(): 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"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1919,7 +1919,7 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by(): prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with( where={"file_purpose": "batch", "created_by": "user1"}, take=10, - order={"created_at": "desc"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) # Query with user2's API key - should only return user2's batch @@ -1934,7 +1934,7 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by(): prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with( where={"file_purpose": "batch", "created_by": "user2"}, take=10, - order={"created_at": "desc"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1965,7 +1965,7 @@ async def test_list_batches_pagination_uses_unified_object_id_cursor(): prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( where={"file_purpose": "batch", "created_by": "test-user"}, take=5, - order={"created_at": "desc"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], cursor={"unified_object_id": "unified-batch-id-7"}, skip=1, ) @@ -2018,10 +2018,12 @@ async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps(): 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_field, direction), = order.items() - result.sort( - key=lambda r: getattr(r, order_field), reverse=(direction == "desc") - ) + 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( @@ -2061,6 +2063,116 @@ async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps(): 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] + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + side_effect=fake_find_many + ) + + 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"]] + if not page_ids: + break + seen.extend(page_ids) + 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" + + @pytest.mark.asyncio async def test_return_unified_file_id_includes_expires_at(): from litellm.types.llms.openai import OpenAIFileObject @@ -2436,7 +2548,7 @@ async def test_list_batches_only_returns_user_own_batches(): 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"}, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) From c1ea54a1d09fd183fea05fa24a455bec756e1bad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:33:33 -0700 Subject: [PATCH 3/4] fix(batches): reject unresolvable list cursors and derive has_more from row count An `after` that does not resolve to a batch the caller can list now returns 400 instead of an empty page. An empty page is indistinguishable from the end of the list, so a stale or malformed cursor silently truncated a client's batch list. The lookup is scoped to the caller's own rows, so a Prisma cursor can no longer be anchored to another user's batch. `has_more` now comes from whether an extra row exists rather than from whether the page came back full. Reporting fullness made every client fetch one extra empty page when the batch count was an exact multiple of `limit`, and made a page shortened by an unparseable row look like the end of the list, hiding the older batches behind it. Also drops the unreachable `target_model_names` oversampling branch; that argument raises a few lines above it. --- .../proxy/hooks/managed_files.py | 36 +-- .../proxy/hooks/test_managed_files.py | 250 +++++++++++++++++- 2 files changed, 261 insertions(+), 25 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index cf5c2b0905d..56e7f7ba633 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -315,29 +315,37 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} - 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) + if after is not None: + 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}'.", + ) + page_size = limit or 20 cursor_args: Dict[str, Any] = ( - {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} + {"cursor": {"unified_object_id": after}, "skip": 1} + if after is not None + else {} ) batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( where=where_clause, - take=fetch_limit, + 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) @@ -353,9 +361,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] diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 347a8fcd023..bbcdbac2709 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1771,7 +1771,7 @@ 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, + take=11, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1801,7 +1801,7 @@ 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, + take=21, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1918,7 +1918,7 @@ 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, + take=11, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1933,7 +1933,7 @@ 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, + take=11, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) @@ -1950,6 +1950,7 @@ async def test_list_batches_pagination_uses_unified_object_id_cursor(): 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( @@ -1964,7 +1965,7 @@ async def test_list_batches_pagination_uses_unified_object_id_cursor(): prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( where={"file_purpose": "batch", "created_by": "test-user"}, - take=5, + take=6, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], cursor={"unified_object_id": "unified-batch-id-7"}, skip=1, @@ -2035,10 +2036,19 @@ async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps(): 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 @@ -2052,9 +2062,10 @@ async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps(): user_api_key_dict=user, limit=3, after=after ) page_ids = [b.id for b in resp["data"]] - if not page_ids: - break 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"] @@ -2146,10 +2157,19 @@ async def test_list_batches_pagination_stable_when_created_at_ties(): 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 @@ -2163,9 +2183,10 @@ async def test_list_batches_pagination_stable_when_created_at_ties(): user_api_key_dict=user, limit=2, after=after ) page_ids = [b.id for b in resp["data"]] - if not page_ids: - break 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"] @@ -2173,6 +2194,215 @@ async def test_list_batches_pagination_stable_when_created_at_ties(): 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_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 @@ -2547,7 +2777,7 @@ 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, + take=11, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], ) From 85ad6971e997f2507d6fe0be51423f22d6bf782d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:49:13 -0700 Subject: [PATCH 4/4] fix(batches): keep an empty `after` meaning "start from the beginning" Validating the cursor whenever `after` was non-None turned `?after=` into a 400, which the listing has always read as "no cursor". Only a cursor the client actually sent is looked up now, matching the sibling managed-resource listing. --- .../proxy/hooks/managed_files.py | 6 ++-- .../proxy/hooks/test_managed_files.py | 29 +++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 56e7f7ba633..8821736d0ff 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -315,7 +315,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} - if after is not None: + if after: cursor_row = ( await self.prisma_client.db.litellm_managedobjecttable.find_first( where={**where_clause, "unified_object_id": after} @@ -329,9 +329,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): page_size = limit or 20 cursor_args: Dict[str, Any] = ( - {"cursor": {"unified_object_id": after}, "skip": 1} - if after is not None - else {} + {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} ) batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index bbcdbac2709..50af6465d06 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -2309,6 +2309,35 @@ async def test_list_batches_rejects_unknown_after_cursor(): ) +@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.