diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 3f42867d90e..8821736d0ff 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/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] 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..50af6465d06 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/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"}], )