fix: parse fenced ROI estimates and retain completed reports

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

litellm/proxy/roi_calculator/estimator.py

@@ -8,6 +8,7 @@ from pydantic import ValidationError
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 from litellm.proxy.roi_calculator.github import SourceError
+from litellm.router_strategy.complexity_router.capability_classifier import extract_classifier_json
 from litellm.types.roi_calculator import (
     ROICompletionMessage,
     ROICompletionMetadata,
@@ -159,7 +160,7 @@ class Estimator:
             choice: Final = parsed_response.choices[0]
             if choice.finish_reason not in (None, "stop") or choice.message.content is None:
                 raise ValueError("incomplete estimator response")
-            result: Final = ROIEstimatorResult.model_validate_json(choice.message.content)
+            result: Final = ROIEstimatorResult.model_validate_json(extract_classifier_json(choice.message.content))
         except Exception:
             raise SourceError(
                 "The estimator did not return valid hours and reasoning. Check the selected model and prompt."

tests/unit/proxy/roi_calculator/test_estimator.py

@@ -32,9 +32,7 @@ def _pull() -> ROIPullEvidence:
         "additions": 1,
         "deletions": 1,
         "changed_files": 1,
-        "files": (
-            {"filename": "time.py", "status": "modified", "additions": 1, "deletions": 1},
-        ),
+        "files": ({"filename": "time.py", "status": "modified", "additions": 1, "deletions": 1},),
         "commits": ({"sha": "abcdef", "message": "Fix timezone conversion"},),
         "commit_count": 1,
         "incomplete_metadata": False,
@@ -53,8 +51,16 @@ def _completion(content: str) -> Mapping[str, object]:
     return response
 
 
+@pytest.mark.parametrize(
+    "content",
+    (
+        '{"hours": 4.25, "reasoning": "Timezone conversion and regression verification."}',
+        '```json\n{"hours": 4.25, "reasoning": "Timezone conversion and regression verification."}\n```',
+        'The estimate is:\n{"hours": 4.25, "reasoning": "Timezone conversion and regression verification."}\nDone.',
+    ),
+)
 @pytest.mark.asyncio
-async def test_estimator_sends_metadata_only_json_request_and_parses_valid_result() -> None:
+async def test_estimator_sends_metadata_only_json_request_and_parses_valid_result(content: str) -> None:
     async def complete(request: ROICompletionRequest) -> object:
         evidence: Final = TypeAdapter(ROIEstimatorEvidence).validate_json(request.messages[1]["content"])
         assert request.temperature == 0
@@ -66,7 +72,7 @@ async def test_estimator_sends_metadata_only_json_request_and_parses_valid_resul
         assert evidence.changes == expected_changes
         assert evidence.commits[0].message == "Fix timezone conversion"
         assert "without AI assistance" in request.messages[0]["content"]
-        return _completion('{"hours": 4.25, "reasoning": "Timezone conversion and regression verification."}')
+        return _completion(content)
 
     result: Final = await Estimator(_settings(), complete).estimate(_pull())
 
@@ -83,6 +89,8 @@ async def test_estimator_sends_metadata_only_json_request_and_parses_valid_resul
         '{"hours": true, "reasoning": "invalid"}',
         '{"hours": 4}',
         '{"hours": 4, "reasoning": "  "}',
+        '```json\n{"hours": -1, "reasoning": "invalid"}\n```',
+        '```json\n{"hours": "4", "reasoning": "invalid"}\n```',
         "not json",
     ),
 )

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

@@ -187,4 +187,59 @@ describe("ROICalculatorView", () => {
     expect(screen.getByLabelText("GitHub token")).toHaveAttribute("type", "password");
     expect(screen.getAllByText("Connect GitHub to get started")).toHaveLength(1);
   });
+
+  it("shows the completed report after polling a running sync", 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)
+      .mockImplementationOnce(
+        () =>
+          new Promise((resolve) => {
+            window.setTimeout(() => resolve({ report: summary }), 25);
+          }),
+      );
+
+    render(<ROICalculatorView accessToken="token" />);
+
+    expect(await screen.findByRole("progressbar", { name: "Sync progress" })).toBeInTheDocument();
+    expect(await screen.findByText("Spend per estimated engineering hour", {}, { timeout: 5000 })).toBeInTheDocument();
+    expect(screen.queryByRole("heading", { name: "Connect GitHub to get started" })).not.toBeInTheDocument();
+  });
+
+  it("shows the sync error returned by the status endpoint", async () => {
+    const runningStatus = {
+      ...idleStatus,
+      running: true,
+      phase: "estimating",
+      stage: "Estimating pull requests",
+      total: 1,
+    };
+    const errorStatus = {
+      ...idleStatus,
+      phase: "error",
+      error: "The estimator could not score a pull request.",
+    };
+    vi.mocked(apiClient.get)
+      .mockResolvedValueOnce(settings)
+      .mockResolvedValueOnce({ report: null })
+      .mockResolvedValueOnce(runningStatus)
+      .mockResolvedValueOnce(errorStatus);
+
+    render(<ROICalculatorView accessToken="token" />);
+
+    expect(await screen.findByRole("alert", {}, { timeout: 5000 })).toHaveTextContent(
+      "The estimator could not score a pull request.",
+    );
+    expect(screen.getByText("Sync failed")).toBeInTheDocument();
+  });
 });

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

@@ -12,11 +12,7 @@ import { Skeleton } from "@/components/ui/skeleton";
 import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { extractErrorMessage } from "@/utils/errorUtils";
 import ROISettingsPanel from "./ROISettingsPanel";
-import {
-  IdentityMatchDialog,
-  type PersonMatchSelection,
-  PullReasoningDialog,
-} from "./ROICalculatorDialogs";
+import { IdentityMatchDialog, type PersonMatchSelection, PullReasoningDialog } from "./ROICalculatorDialogs";
 import { ROIOverview, ROIPeopleView } from "./ROICalculatorViews";
 import { filterPulls } from "./roiCalculatorData";
 import type {
@@ -93,11 +89,12 @@ export default function ROICalculatorView({ accessToken }: { accessToken: string
         .get<ROISyncStatus>("/roi-calculator/sync", { accessToken })
         .then(async (nextStatus) => {
           if (cancelled) return;
-          setStatus(nextStatus);
           if (!nextStatus.running && nextStatus.phase === "complete") {
             const report = await loadReport();
-            if (!cancelled) setSummary(report);
+            if (cancelled) return;
+            setSummary(report);
           }
+          if (!cancelled) setStatus(nextStatus);
         })
         .catch((reason: unknown) => {
           if (!cancelled) setError(extractErrorMessage(reason));
@@ -144,10 +141,7 @@ export default function ROICalculatorView({ accessToken }: { accessToken: string
     [accessToken],
   );
 
-  const filteredPulls = React.useMemo(
-    () => (summary ? filterPulls(summary.pulls, query) : []),
-    [query, summary],
-  );
+  const filteredPulls = React.useMemo(() => (summary ? filterPulls(summary.pulls, query) : []), [query, summary]);
 
   if (error && !settings) {
     return (