fix(roi): page repository search and recover polling errors

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm/proxy/roi_calculator/github.py

@@ -50,6 +50,17 @@ class _RepositoryItem(_GitHubModel):
     archived: bool = False
 
 
+def _repository_values(repositories: tuple[_RepositoryItem, ...]) -> tuple[tuple[str, str, bool], ...]:
+    return tuple(
+        (
+            repository.full_name,
+            repository.visibility or ("private" if repository.private else "public"),
+            repository.archived,
+        )
+        for repository in repositories
+    )
+
+
 class _PullDetail(_GitHubModel):
     number: int
     title: str
@@ -178,6 +189,8 @@ class _GraphQLPayload(TypedDict):
 
 
 _REPOSITORIES: Final[TypeAdapter[tuple[_RepositoryItem, ...]]] = TypeAdapter(tuple[_RepositoryItem, ...])
+_REPOSITORY_SEARCH_PAGES: Final[int] = 10
+_REPOSITORY_PAGE_ERROR: Final[str] = "GitHub returned an unexpected repository list."
 _PULLS: Final[TypeAdapter[tuple[GitHubPullListItem, ...]]] = TypeAdapter(tuple[GitHubPullListItem, ...])
 _PULL_FILES: Final[TypeAdapter[tuple[_PullFile, ...]]] = TypeAdapter(tuple[_PullFile, ...])
 _REST_COMMITS: Final[TypeAdapter[tuple[_RestCommit, ...]]] = TypeAdapter(tuple[_RestCommit, ...])
@@ -243,6 +256,7 @@ async def _fetch_page(
     params: Mapping[str, str | int] | None,
     page: int,
     headers: Mapping[str, str] | None = None,
+    error_message: str = "GitHub returned an unexpected pagination response.",
 ) -> tuple[tuple[_T, ...], bool]:
     response: Final = await _request(
         client,
@@ -260,7 +274,7 @@ async def _fetch_page(
     try:
         parsed: Final[tuple[_T, ...]] = adapter.validate_python(response.json())
     except Exception:
-        raise SourceError("GitHub returned an unexpected pagination response.") from None
+        raise SourceError(error_message) from None
     return parsed, 'rel="next"' in response.headers.get("link", "")
 
 
@@ -337,35 +351,51 @@ class GitHub:
         query: str = "",
         page: int = 1,
     ) -> tuple[tuple[tuple[str, str, bool], ...], bool]:
-        response: Final = await _request(
-            self.client,
-            "GET",
-            self._url("user/repos"),
-            params=MappingProxyType(
-                {
-                    "per_page": 100,
-                    "page": page,
-                    "sort": "updated",
-                    "direction": "desc",
-                    "affiliation": "owner,collaborator,organization_member",
-                }
-            ),
-            headers=self._headers,
+        params: Final = MappingProxyType(
+            {
+                "sort": "updated",
+                "direction": "desc",
+                "affiliation": "owner,collaborator,organization_member",
+            }
         )
-        try:
-            repositories: Final[tuple[_RepositoryItem, ...]] = _REPOSITORIES.validate_python(response.json())
-        except Exception:
-            raise SourceError("GitHub returned an unexpected repository list.") from None
-        filtered: Final = tuple(
-            (
-                repository.full_name,
-                repository.visibility or ("private" if repository.private else "public"),
-                repository.archived,
+        if not query:
+            repositories, has_more = await _fetch_page(
+                self.client,
+                self._url("user/repos"),
+                _REPOSITORIES,
+                params,
+                page,
+                self._headers,
+                error_message=_REPOSITORY_PAGE_ERROR,
             )
-            for repository in repositories
-            if query.casefold() in repository.full_name.casefold()
-        )
-        return filtered, 'rel="next"' in response.headers.get("link", "")
+            return _repository_values(repositories), has_more
+
+        normalized_query: Final = query.casefold()
+        first_github_page: Final = (page - 1) * _REPOSITORY_SEARCH_PAGES + 1
+
+        async def search_pages(
+            github_page: int,
+            pages_remaining: int,
+        ) -> tuple[tuple[_RepositoryItem, ...], bool]:
+            repositories, has_more = await _fetch_page(
+                self.client,
+                self._url("user/repos"),
+                _REPOSITORIES,
+                params,
+                github_page,
+                self._headers,
+                error_message=_REPOSITORY_PAGE_ERROR,
+            )
+            matches: Final = tuple(
+                repository for repository in repositories if normalized_query in repository.full_name.casefold()
+            )
+            if pages_remaining == 1 or not has_more:
+                return matches, has_more
+            later_matches, later_has_more = await search_pages(github_page + 1, pages_remaining - 1)
+            return (*matches, *later_matches), later_has_more
+
+        matches, has_more = await search_pages(first_github_page, _REPOSITORY_SEARCH_PAGES)
+        return _repository_values(matches), has_more
 
     async def pulls(self, repo: str, start: date, end: date) -> tuple[GitHubPullListItem, ...]:
         async def pull_pages() -> AsyncIterator[GitHubPullListItem]:

tests/unit/proxy/roi_calculator/test_github.py

@@ -109,16 +109,12 @@ async def test_github_maps_upstream_errors_without_returning_response_secrets()
 
 
 @pytest.mark.asyncio
-async def test_github_repository_listing_applies_search_and_reports_next_page() -> None:
+async def test_github_repository_search_starts_page_two_at_github_page_eleven() -> None:
     def respond(request: httpx.Request) -> httpx.Response:
-        assert request.url.params["page"] == "2"
+        assert request.url.params["page"] == "11"
         assert request.url.params["affiliation"] == "owner,collaborator,organization_member"
         assert request.headers["authorization"] == "Bearer test-github-token"
-        return httpx.Response(
-            200,
-            headers=_NEXT_PAGE_HEADERS,
-            content=_REPOSITORIES_JSON,
-        )
+        return httpx.Response(200, content=_REPOSITORIES_JSON)
 
     github: Final = _github(httpx.MockTransport(respond))
     try:
@@ -127,4 +123,52 @@ async def test_github_repository_listing_applies_search_and_reports_next_page()
         await github.close()
 
     assert repositories == (("org/backend", "private", False),)
-    assert has_more
+    assert not has_more
+
+
+@pytest.mark.asyncio
+async def test_github_repository_search_scans_until_a_later_page_match() -> None:
+    expected_pages: Final = iter(("1", "2", "3"))
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        page: Final = request.url.params["page"]
+        assert page == next(expected_pages)
+        if page == "3":
+            return httpx.Response(
+                200,
+                content='[{"full_name":"org/target-repo","visibility":"private","archived":false}]',
+            )
+        return httpx.Response(200, headers=_NEXT_PAGE_HEADERS, content=_REPOSITORIES_JSON)
+
+    github: Final = _github(httpx.MockTransport(respond))
+    try:
+        repositories, has_more = await github.repositories(query="TARGET", page=1)
+    finally:
+        await github.close()
+
+    assert repositories == (("org/target-repo", "private", False),)
+    assert not has_more
+    assert next(expected_pages, None) is None
+
+
+@pytest.mark.asyncio
+async def test_github_repository_search_pages_ten_github_pages_per_search_page() -> None:
+    expected_pages: Final = iter(tuple(str(page) for page in range(1, 21)))
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        page: Final = request.url.params["page"]
+        assert page == next(expected_pages)
+        return httpx.Response(200, headers=_NEXT_PAGE_HEADERS, content="[]")
+
+    github: Final = _github(httpx.MockTransport(respond))
+    try:
+        first_repositories, first_has_more = await github.repositories(query="missing", page=1)
+        second_repositories, second_has_more = await github.repositories(query="missing", page=2)
+    finally:
+        await github.close()
+
+    assert first_repositories == ()
+    assert first_has_more
+    assert second_repositories == ()
+    assert second_has_more
+    assert next(expected_pages, None) is None

ui/litellm-dashboard/src/app/(dashboard)/roi-calculator/_components/ROICalculatorView.integration.test.tsx

@@ -276,4 +276,56 @@ describe("ROICalculatorView", () => {
     );
     expect(screen.getByText("Sync failed")).toBeInTheDocument();
   });
+
+  it("shows a report error and ends progress when the completed report cannot load", async () => {
+    const runningStatus = {
+      ...idleStatus,
+      running: true,
+      phase: "estimating",
+      stage: "Estimating pull requests",
+      total: 1,
+    };
+    const completedStatus = { ...idleStatus, phase: "complete", done: 1, total: 1 };
+    vi.mocked(apiClient.get)
+      .mockResolvedValueOnce(settings)
+      .mockResolvedValueOnce({ report: null })
+      .mockResolvedValueOnce(runningStatus)
+      .mockResolvedValueOnce(completedStatus)
+      .mockRejectedValueOnce(new Error("The report could not be loaded."));
+
+    render(<ROICalculatorView accessToken="token" />);
+
+    expect(await screen.findByRole("progressbar", { name: "Sync progress" })).toBeInTheDocument();
+    expect(await screen.findByRole("alert", {}, { timeout: 5000 })).toHaveTextContent(
+      "The report could not be loaded.",
+    );
+    expect(screen.queryByRole("progressbar", { name: "Sync progress" })).not.toBeInTheDocument();
+  });
+
+  it("clears a transient poll error when the next poll completes and loads the report", async () => {
+    const runningStatus = {
+      ...idleStatus,
+      running: true,
+      phase: "estimating",
+      stage: "Estimating pull requests",
+      total: 1,
+    };
+    const completedStatus = { ...idleStatus, phase: "complete", done: 1, total: 1 };
+    vi.mocked(apiClient.get)
+      .mockResolvedValueOnce(settings)
+      .mockResolvedValueOnce({ report: null })
+      .mockResolvedValueOnce(runningStatus)
+      .mockRejectedValueOnce(new Error("The sync status could not be loaded."))
+      .mockResolvedValueOnce(completedStatus)
+      .mockResolvedValueOnce({ report: summary });
+
+    render(<ROICalculatorView accessToken="token" />);
+
+    expect(await screen.findByRole("progressbar", { name: "Sync progress" })).toBeInTheDocument();
+    expect(await screen.findByRole("alert", {}, { timeout: 5000 })).toHaveTextContent(
+      "The sync status could not be loaded.",
+    );
+    expect(await screen.findByText("Spend per estimated engineering hour", {}, { timeout: 7000 })).toBeInTheDocument();
+    expect(screen.queryByText("The sync status could not be loaded.")).not.toBeInTheDocument();
+  });
 });

ui/litellm-dashboard/src/app/(dashboard)/roi-calculator/_components/ROICalculatorView.tsx

@@ -99,11 +99,17 @@ export default function ROICalculatorView({
         .get<ROISyncStatus>("/roi-calculator/sync", { accessToken })
         .then(async (nextStatus) => {
           if (cancelled) return;
+          setError(null);
           if (!nextStatus.running && nextStatus.phase === "complete") {
-            const report = await loadReport();
-            if (cancelled) return;
-            setSummary(report);
-            if (view === "settings") setView("overview");
+            try {
+              const report = await loadReport();
+              if (cancelled) return;
+              setSummary(report);
+              setError(null);
+              if (view === "settings") setView("overview");
+            } catch (reason) {
+              if (!cancelled) setError(extractErrorMessage(reason));
+            }
           }
           if (!cancelled) setStatus(nextStatus);
         })