From b221aa0cc20c69e5baa228ef950b5282b53e9f02 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 11 Feb 2026 19:25:01 -0800
Subject: [PATCH] backend fix

---
 .../spend_management_endpoints.py             | 87 +++++++++++++++++--
 1 file changed, 80 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index c05eda85c81..b46deacc1cc 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -1850,13 +1850,15 @@ async def ui_view_spend_logs(  # noqa: PLR0915
 
         verbose_proxy_logger.debug("data= %s", json.dumps(data, indent=4, default=str))
 
-        return {
-            "data": data,
-            "total": total_records,
-            "page": page,
-            "page_size": page_size,
-            "total_pages": total_pages,
-        }
+        return await _build_ui_spend_logs_response(
+            prisma_client,
+            data,
+            total_records,
+            page,
+            page_size,
+            total_pages,
+            enrich_session_counts=not is_v2,
+        )
     except Exception as e:
         verbose_proxy_logger.exception(f"Error in ui_view_spend_logs: {e}")
         raise handle_exception_on_proxy(e)
@@ -3101,6 +3103,77 @@ async def ui_view_session_spend_logs(
             )
 
 
+async def _build_ui_spend_logs_response(
+    prisma_client: "PrismaClient",
+    data: list,
+    total_records: int,
+    page: int,
+    page_size: int,
+    total_pages: int,
+    enrich_session_counts: bool = True,
+) -> dict:
+    """
+    Build the paginated response for the UI spend-logs endpoint.
+
+    When ``enrich_session_counts`` is ``True`` (the default for the v1/UI
+    endpoint), each row is enriched with ``session_total_count`` so the
+    frontend knows which sessions are expandable (multi-call sessions).
+    For every row that carries a ``session_id``, a single ``GROUP BY`` query
+    fetches the total number of logs in each referenced session.  Rows without
+    a ``session_id`` default to ``1``.
+
+    When ``enrich_session_counts`` is ``False`` (v2 endpoint), rows are
+    serialised without the extra query.
+
+    Args:
+        prisma_client: The connected Prisma client instance.
+        data: A list of Prisma model instances (must support ``.model_dump()``
+              and have a ``session_id`` attribute).
+        total_records: Total number of matching records (for pagination).
+        page: Current page number.
+        page_size: Number of items per page.
+        total_pages: Total number of pages.
+        enrich_session_counts: Whether to add ``session_total_count`` to each
+            row.  Defaults to ``True``.
+
+    Returns:
+        A dict with ``data`` (enriched rows), ``total``, ``page``,
+        ``page_size``, and ``total_pages``.
+    """
+    count_map: dict[str, int] = {}
+    if enrich_session_counts:
+        session_ids = list(
+            {row.session_id for row in data if getattr(row, "session_id", None)}
+        )
+        if session_ids:
+            counts = await prisma_client.db.litellm_spendlogs.group_by(
+                by=["session_id"],
+                where={"session_id": {"in": session_ids}},
+                count={"session_id": True},
+            )
+            count_map = {
+                r["session_id"]: r["_count"]["session_id"]
+                for r in counts
+                if r.get("session_id")
+            }
+
+    enriched: List[dict] = []
+    for row in data:
+        row_dict = row.model_dump()
+        if enrich_session_counts:
+            sid = row_dict.get("session_id")
+            row_dict["session_total_count"] = count_map.get(sid, 1) if sid else 1
+        enriched.append(row_dict)
+
+    return {
+        "data": enriched,
+        "total": total_records,
+        "page": page,
+        "page_size": page_size,
+        "total_pages": total_pages,
+    }
+
+
 def _build_status_filter_condition(status_filter: Optional[str]) -> Dict[str, Any]:
     """
     Helper function to build the status filter condition for database queries.