fix(proxy): align list page continuation

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/fine_tuning_endpoints/endpoints.py

@@ -49,6 +49,9 @@ async def _provider_job_cursor(
 ) -> str | None:
     if after is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
         return after
+    decoded: Final = is_base64_encoded_unified_file_id(after)
+    if not decoded:
+        return after
     await validate_managed_id_requirement(
         resource_id=after,
         resource_kind="fine-tuning job",
@@ -56,8 +59,7 @@ async def _provider_job_cursor(
         managed_files_obj=managed_files_obj,
         require_owned_resource=True,
     )
-    decoded: Final = is_base64_encoded_unified_file_id(after)
-    return get_batch_id_from_unified_batch_id(decoded) if decoded else after
+    return get_batch_id_from_unified_batch_id(decoded)
 
 
 async def _owned_jobs(response: object, user_api_key_dict: UserAPIKeyAuth, managed_files_obj: object) -> object:
@@ -118,19 +120,28 @@ async def _collect_owned_job_pages(
     user_api_key_dict: UserAPIKeyAuth,
     managed_files_obj: object,
 ) -> object:
+    loop: Final = asyncio.get_running_loop()
+    deadline: Final = loop.time() + _JOB_LIST_SCAN_TIMEOUT
+
+    async def owned_page(cursor: str | None) -> _OwnedJobPage:
+        return await _owned_job_page(await next_page(cursor), user_api_key_dict, managed_files_obj)
+
     state: tuple[str | None, _OwnedJobPage] = (  # rebind-ok: advance through provider pages without recursion
         None,
         await _owned_job_page(response, user_api_key_dict, managed_files_obj),
     )
     for _ in range(_JOB_LIST_MAX_PAGES - 1):
-        if state[1].next_cursor is None or state[1].next_cursor == state[0]:
+        if state[1].next_cursor is None or state[1].next_cursor == state[0] or loop.time() >= deadline:
+            break
+        try:
+            state = (
+                state[1].next_cursor,
+                await asyncio.wait_for(owned_page(state[1].next_cursor), timeout=deadline - loop.time()),
+            )
+        except asyncio.TimeoutError:
             break
-        state = (
-            state[1].next_cursor,
-            await _owned_job_page(await next_page(state[1].next_cursor), user_api_key_dict, managed_files_obj),
-        )
     if state[1].next_cursor is not None and state[1].next_cursor != state[0]:
-        raise HTTPException(status_code=503, detail="Fine-tuning job listing exceeded its page budget")
+        return {**state[1].response, "has_more": True, "last_id": state[1].next_cursor}
     return {**state[1].response, "has_more": False} if state[1].next_cursor is not None else state[1].response
 
 
@@ -142,13 +153,7 @@ async def _owned_job_pages(
 ) -> object:
     if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
         return response
-    try:
-        return await asyncio.wait_for(
-            _collect_owned_job_pages(response, next_page, user_api_key_dict, managed_files_obj),
-            timeout=_JOB_LIST_SCAN_TIMEOUT,
-        )
-    except asyncio.TimeoutError as e:
-        raise HTTPException(status_code=504, detail="Fine-tuning job listing timed out") from e
+    return await _collect_owned_job_pages(response, next_page, user_api_key_dict, managed_files_obj)
 
 
 def set_fine_tuning_config(config):

litellm/proxy/vector_store_files_endpoints/endpoints.py

@@ -189,21 +189,33 @@ async def _collect_managed_file_list_pages(
     user_api_key_dict: UserAPIKeyAuth,
     before: bool,
 ) -> VectorStoreFileListResponse:
+    loop: Final = asyncio.get_running_loop()
+    deadline: Final = loop.time() + _FILE_LIST_SCAN_TIMEOUT
+
+    async def owned_page(cursor: str) -> _ManagedFilePage:
+        return await _managed_file_list_page(await next_page(cursor), managed_files_obj, user_api_key_dict, before)
+
     state: tuple[str | None, _ManagedFilePage] = (  # rebind-ok: advance a bounded page scan
         None,
         await _managed_file_list_page(response, managed_files_obj, user_api_key_dict, before),
     )
     for _ in range(_FILE_LIST_MAX_PAGES - 1):
-        if state[1].next_cursor is None or state[1].next_cursor == state[0]:
+        if state[1].next_cursor is None or state[1].next_cursor == state[0] or loop.time() >= deadline:
+            break
+        try:
+            state = (
+                state[1].next_cursor,
+                await asyncio.wait_for(owned_page(state[1].next_cursor), timeout=deadline - loop.time()),
+            )
+        except asyncio.TimeoutError:
             break
-        state = (
-            state[1].next_cursor,
-            await _managed_file_list_page(
-                await next_page(state[1].next_cursor), managed_files_obj, user_api_key_dict, before
-            ),
-        )
     if state[1].next_cursor is not None and state[1].next_cursor != state[0]:
-        raise HTTPException(status_code=503, detail="Vector store file listing exceeded its page budget")
+        return {
+            **state[1].response,
+            "has_more": True,
+            "first_id": state[1].next_cursor if before else None,
+            "last_id": state[1].next_cursor if not before else None,
+        }
     return {**state[1].response, "has_more": False} if state[1].next_cursor is not None else state[1].response
 
 
@@ -216,13 +228,7 @@ async def _with_managed_file_list_pages(
 ) -> VectorStoreFileListResponse:
     if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
         return await _with_managed_file_list_ids(response, managed_files_obj, user_api_key_dict)
-    try:
-        return await asyncio.wait_for(
-            _collect_managed_file_list_pages(response, next_page, managed_files_obj, user_api_key_dict, before),
-            timeout=_FILE_LIST_SCAN_TIMEOUT,
-        )
-    except asyncio.TimeoutError as e:
-        raise HTTPException(status_code=504, detail="Vector store file listing timed out") from e
+    return await _collect_managed_file_list_pages(response, next_page, managed_files_obj, user_api_key_dict, before)
 
 
 async def _update_request_data_with_managed_file_id(

tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py

@@ -10,6 +10,7 @@ Each test asserts BOTH that the request is rejected AND that every downstream pr
 seam stayed untouched, so a guard that raises after the provider call would still fail.
 """
 
+import asyncio
 import base64
 from contextlib import ExitStack
 from collections.abc import Mapping, Sequence
@@ -341,14 +342,14 @@ PROVIDER_JOBS_PAGE: Final = {
 }
 
 
-async def _list(custom_llm_provider: str | None):
+async def _list(custom_llm_provider: str | None, *, after: str | None = None, limit: int | None = None):
     return await endpoints.list_fine_tuning_jobs(
         request=FakeRequest(),
         fastapi_response=Response(),
         custom_llm_provider=custom_llm_provider,
         target_model_names=None,
-        after=None,
-        limit=None,
+        after=after,
+        limit=limit,
         user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
     )
 
@@ -433,10 +434,8 @@ async def test_admin_job_cursor_keeps_provider_value(cursor: str | None) -> None
 
 
 @pytest.mark.asyncio
-async def test_job_cursor_requires_a_managed_reference() -> None:
-    with pytest.raises(HTTPException) as exc:
-        await endpoints._provider_job_cursor(RAW_JOB_ID, UserAPIKeyAuth(api_key="test"), None)
-    assert exc.value.status_code == 403
+async def test_job_list_cursor_accepts_a_provider_page_boundary() -> None:
+    assert await endpoints._provider_job_cursor(RAW_JOB_ID, UserAPIKeyAuth(api_key="test"), None) == RAW_JOB_ID
 
 
 @pytest.mark.asyncio
@@ -460,12 +459,71 @@ async def test_job_listing_continues_until_an_owned_page(seams) -> None:
 
 
 @pytest.mark.asyncio
-async def test_job_listing_stops_at_its_page_budget() -> None:
+async def test_job_listing_returns_a_resumable_page_at_its_budget() -> None:
     async def next_page(cursor: str | None) -> object:
         assert cursor is not None
         return {"data": [{"id": cursor + "-next"}], "has_more": True}
 
     fetch: Final = AsyncMock(side_effect=next_page)
+    assert await endpoints._owned_job_pages(
+        {"data": [{"id": "job-first"}], "has_more": True},
+        fetch,
+        UserAPIKeyAuth(api_key="test"),
+        _OwnedJobIds({}),
+    ) == {"data": [], "has_more": True, "last_id": "job-first" + "-next" * (endpoints._JOB_LIST_MAX_PAGES - 1)}
+    assert fetch.await_count == endpoints._JOB_LIST_MAX_PAGES - 1
+
+
+@pytest.mark.asyncio
+async def test_job_listing_returns_a_resumable_page_at_its_deadline(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(endpoints, "_JOB_LIST_SCAN_TIMEOUT", 0)
+
+    async def next_page(cursor: str | None) -> object:
+        pytest.fail("An expired scan must not request another page")
+
+    assert await endpoints._owned_job_pages(
+        {"data": [{"id": "job-first"}], "has_more": True}, next_page, UserAPIKeyAuth(api_key="test"), _OwnedJobIds({})
+    ) == {"data": [], "has_more": True, "last_id": "job-first"}
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_job_listing_resumes_after_an_empty_scan(seams) -> None:
+    job: Final = _job().model_dump()
+    route: Final = respx.get(f"{FINE_TUNING_API_BASE}/fine_tuning/jobs").mock(
+        side_effect=[
+            *(
+                httpx.Response(200, json={"data": [{**job, "id": f"ftjob-page-{index}"}], "has_more": True})
+                for index in range(endpoints._JOB_LIST_MAX_PAGES)
+            ),
+            httpx.Response(200, json={"data": [job], "has_more": False}),
+        ]
+    )
+    seams.logging.get_proxy_hook.return_value = _OwnedJobIds({RAW_JOB_ID: _unified_job_id()})
+    provider: Final = {"custom_llm_provider": "openai", "api_key": "sk-provider", "api_base": FINE_TUNING_API_BASE}
+    with patch.object(endpoints, "fine_tuning_config", [provider]):
+        page: Final = await _list(custom_llm_provider="openai", limit=1)
+        assert page == {"data": [], "has_more": True, "last_id": f"ftjob-page-{endpoints._JOB_LIST_MAX_PAGES - 1}"}
+        following: Final = await _list(custom_llm_provider="openai", after=page["last_id"], limit=1)
+    assert following == {"data": [{**job, "id": _unified_job_id()}], "has_more": False}
+    assert route.calls[-1].request.url.params["after"] == page["last_id"]
+    assert route.calls[-1].request.url.params["limit"] == "1"
+
+
+@pytest.mark.asyncio
+async def test_job_scan_timeout_keeps_the_last_completed_page() -> None:
+    fetch: Final = AsyncMock(side_effect=asyncio.TimeoutError)
+    assert await endpoints._owned_job_pages(
+        {"data": [{"id": "job-first"}], "has_more": True},
+        fetch,
+        UserAPIKeyAuth(api_key="test"),
+        _OwnedJobIds({}),
+    ) == {"data": [], "has_more": True, "last_id": "job-first"}
+
+
+@pytest.mark.asyncio
+async def test_job_scan_preserves_an_upstream_error() -> None:
+    fetch: Final = AsyncMock(side_effect=HTTPException(status_code=502, detail="upstream unavailable"))
     with pytest.raises(HTTPException) as exc:
         await endpoints._owned_job_pages(
             {"data": [{"id": "job-first"}], "has_more": True},
@@ -473,21 +531,4 @@ async def test_job_listing_stops_at_its_page_budget() -> None:
             UserAPIKeyAuth(api_key="test"),
             _OwnedJobIds({}),
         )
-    assert exc.value.status_code == 503
-    assert exc.value.detail == "Fine-tuning job listing exceeded its page budget"
-    assert fetch.await_count == endpoints._JOB_LIST_MAX_PAGES - 1
-
-
-@pytest.mark.asyncio
-async def test_job_listing_reports_an_expired_scan_deadline(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(endpoints, "_JOB_LIST_SCAN_TIMEOUT", 0)
-
-    async def next_page(cursor: str | None) -> object:
-        pytest.fail("An expired scan must not request another page")
-
-    with pytest.raises(HTTPException) as exc:
-        await endpoints._owned_job_pages(
-            {"data": [], "has_more": False}, next_page, UserAPIKeyAuth(api_key="test"), _OwnedJobIds({})
-        )
-    assert exc.value.status_code == 504
-    assert exc.value.detail == "Fine-tuning job listing timed out"
+    assert exc.value.status_code == 502

tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py

@@ -9,6 +9,7 @@ A raw or forged managed-looking file id has no ownership row, so without the gua
 is attached to a vector store or read back under shared provider credentials.
 """
 
+import asyncio
 import base64
 from collections.abc import Mapping, Sequence
 from copy import deepcopy
@@ -316,28 +317,75 @@ async def test_admin_file_listing_preserves_a_provider_page() -> None:
 
 
 @pytest.mark.asyncio
-async def test_file_listing_reports_its_page_budget() -> None:
+@pytest.mark.parametrize("before", [False, True])
+async def test_file_listing_resumes_after_its_page_budget(before: bool) -> None:
     async def next_page(cursor: str) -> VectorStoreFileListResponse:
+        if cursor.count("-next") == endpoints._FILE_LIST_MAX_PAGES - 1:
+            return {"data": [_vector_store_file_row(RAW_FILE_ID)], "has_more": False}
         return {"data": [_vector_store_file_row(cursor + "-next")], "has_more": True}
 
     fetch: Final = AsyncMock(side_effect=next_page)
     resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={}))
+    page: Final = await endpoints._with_managed_file_list_pages(
+        {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER, before
+    )
+    cursor: Final = "file-first" + "-next" * (endpoints._FILE_LIST_MAX_PAGES - 1)
+    assert page == {
+        "data": [],
+        "has_more": True,
+        "first_id": cursor if before else None,
+        "last_id": None if before else cursor,
+    }
+    assert fetch.await_count == endpoints._FILE_LIST_MAX_PAGES - 1
+    query: Final = _with_provider_file_id_cursors({"before" if before else "after": cursor})
+    assert query == {"before" if before else "after": cursor}
+    following: Final = await fetch(query["before" if before else "after"])
+    owned_resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={RAW_FILE_ID: _unified_file_id()}))
+    assert await endpoints._with_managed_file_list_pages(following, fetch, owned_resolver, CALLER, before) == {
+        "data": [_vector_store_file_row(_unified_file_id())],
+        "has_more": False,
+        "first_id": _unified_file_id(),
+        "last_id": _unified_file_id(),
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("before", [False, True])
+async def test_file_listing_returns_a_resumable_page_at_its_deadline(
+    monkeypatch: pytest.MonkeyPatch, before: bool
+) -> None:
+    monkeypatch.setattr(endpoints, "_FILE_LIST_SCAN_TIMEOUT", 0)
+    fetch: Final = AsyncMock()
+    resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={}))
+    assert await endpoints._with_managed_file_list_pages(
+        {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER, before
+    ) == {
+        "data": [],
+        "has_more": True,
+        "first_id": "file-first" if before else None,
+        "last_id": None if before else "file-first",
+    }
+    fetch.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_file_scan_timeout_keeps_the_last_completed_page() -> None:
+    resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={}))
+    fetch: Final = AsyncMock(side_effect=asyncio.TimeoutError)
+    assert await endpoints._with_managed_file_list_pages(
+        {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER
+    ) == {"data": [], "has_more": True, "first_id": None, "last_id": "file-first"}
+
+
+@pytest.mark.asyncio
+async def test_file_scan_preserves_an_upstream_error() -> None:
+    resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={}))
+    fetch: Final = AsyncMock(side_effect=HTTPException(status_code=502, detail="upstream unavailable"))
     with pytest.raises(HTTPException) as exc:
         await endpoints._with_managed_file_list_pages(
             {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER
         )
-    assert exc.value.status_code == 503
-    assert fetch.await_count == endpoints._FILE_LIST_MAX_PAGES - 1
-
-
-@pytest.mark.asyncio
-async def test_file_listing_reports_an_expired_deadline(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(endpoints, "_FILE_LIST_SCAN_TIMEOUT", 0)
-    fetch: Final = AsyncMock()
-    with pytest.raises(HTTPException) as exc:
-        await endpoints._with_managed_file_list_pages({"data": [], "has_more": False}, fetch, None, CALLER)
-    assert exc.value.status_code == 504
-    fetch.assert_not_awaited()
+    assert exc.value.status_code == 502
 
 
 @pytest.mark.asyncio