Merge pull request #39598 from BerriAI/litellm_lit_4929_session_token_aggregation

fix(ui): aggregate session token usage in the logs table

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -173,6 +173,9 @@ class _SessionSpendRow(TypedDict):
     session_cache_hit_count: ReadOnly[int]
     session_llm_count: ReadOnly[int]
     session_agent_count: ReadOnly[int]
+    session_total_prompt_tokens: ReadOnly[int]
+    session_total_completion_tokens: ReadOnly[int]
+    session_total_tokens: ReadOnly[int]
     session_models: ReadOnly[Sequence[str]]
 
 
@@ -188,6 +191,9 @@ class _SessionSpendStats(NamedTuple):
     session_cache_hit_count: int
     session_llm_count: int
     session_agent_count: int
+    session_total_prompt_tokens: int
+    session_total_completion_tokens: int
+    session_total_tokens: int
     session_models: Sequence[str]
     session_models_truncated: bool
 
@@ -4287,8 +4293,8 @@ async def _build_ui_spend_logs_response(
     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`` plus spend
-    and call-type aggregates so the frontend knows which sessions are
+    endpoint), each row is enriched with ``session_total_count`` plus spend,
+    token and call-type aggregates so the frontend knows which sessions are
     expandable (multi-call sessions).  One ``GROUP BY (session_id, api_key)``
     query serves every referenced session, keyed per api key so two callers
     reusing a session id never see each other's totals.  Rows without a
@@ -4356,7 +4362,10 @@ async def _build_ui_spend_logs_response(
                            COUNT(*) FILTER (
                                WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
                            )::int AS session_llm_count,
-                           COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count
+                           COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count,
+                           COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens,
+                           COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens,
+                           COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens
                     FROM "LiteLLM_SpendLogs"
                     WHERE session_id = ANY($1::text[])
                       AND api_key = ANY($2::text[])
@@ -4389,6 +4398,9 @@ async def _build_ui_spend_logs_response(
                     session_cache_hit_count=int(row.get("session_cache_hit_count") or 0),
                     session_llm_count=int(row.get("session_llm_count") or 0),
                     session_agent_count=int(row.get("session_agent_count") or 0),
+                    session_total_prompt_tokens=int(row.get("session_total_prompt_tokens") or 0),
+                    session_total_completion_tokens=int(row.get("session_total_completion_tokens") or 0),
+                    session_total_tokens=int(row.get("session_total_tokens") or 0),
                     session_models=models[:_SESSION_MODELS_LIMIT],
                     session_models_truncated=len(models) > _SESSION_MODELS_LIMIT,
                 )
@@ -4418,6 +4430,9 @@ async def _build_ui_spend_logs_response(
                 row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count
                 row_dict["session_llm_count"] = session_stats.session_llm_count
                 row_dict["session_agent_count"] = session_stats.session_agent_count
+                row_dict["session_total_prompt_tokens"] = session_stats.session_total_prompt_tokens
+                row_dict["session_total_completion_tokens"] = session_stats.session_total_completion_tokens
+                row_dict["session_total_tokens"] = session_stats.session_total_tokens
                 row_dict["session_models"] = session_stats.session_models
                 row_dict["session_models_truncated"] = session_stats.session_models_truncated
             enriched.append(row_dict)

tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py

@@ -4234,6 +4234,91 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend():
     assert call_args[2] == [api_key]
 
 
+@pytest.mark.asyncio
+async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens():
+    """
+    Regression test for LIT-4929: the logs table showed the summed session cost but
+    only the last call's token usage.  Every row of a multi-round session must carry
+    the session-wide prompt, completion and total token sums from the aggregate
+    query, while rows outside a session carry none of them.
+    """
+    from litellm.proxy.spend_tracking.spend_management_endpoints import (
+        _build_ui_spend_logs_response,
+    )
+
+    session_id = "sess-multi-round-tokens"
+    api_key = "hashed-key-xyz"
+    dict_rows = [
+        {
+            "request_id": "req-1",
+            "session_id": session_id,
+            "call_type": "completion",
+            "api_key": api_key,
+            "total_tokens": 10,
+            "prompt_tokens": 7,
+            "completion_tokens": 3,
+        },
+        {
+            "request_id": "req-2",
+            "session_id": session_id,
+            "call_type": "completion",
+            "api_key": api_key,
+            "total_tokens": 50,
+            "prompt_tokens": 35,
+            "completion_tokens": 15,
+        },
+        {
+            "request_id": "req-3",
+            "session_id": None,
+            "call_type": "completion",
+            "api_key": api_key,
+            "total_tokens": 5,
+            "prompt_tokens": 4,
+            "completion_tokens": 1,
+        },
+    ]
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "session_id": session_id,
+                "api_key": api_key,
+                "session_total_count": 2,
+                "session_total_spend": 0.06,
+                "mcp_tool_call_count": 0,
+                "mcp_tool_call_spend": 0.0,
+                "session_total_prompt_tokens": 42,
+                "session_total_completion_tokens": 18,
+                "session_total_tokens": 60,
+            }
+        ]
+    )
+
+    result = await _build_ui_spend_logs_response(
+        prisma_client=mock_prisma,
+        data=dict_rows,
+        total_records=3,
+        page=1,
+        page_size=50,
+        total_pages=1,
+        enrich_session_counts=True,
+    )
+
+    rows = result["data"]
+    session_rows = rows[:2]
+    assert [row["session_total_tokens"] for row in session_rows] == [60, 60]
+    assert [row["session_total_prompt_tokens"] for row in session_rows] == [42, 42]
+    assert [row["session_total_completion_tokens"] for row in session_rows] == [18, 18]
+    assert [(row["total_tokens"], row["prompt_tokens"], row["completion_tokens"]) for row in session_rows] == [
+        (10, 7, 3),
+        (50, 35, 15),
+    ]
+
+    token_keys = ("session_total_tokens", "session_total_prompt_tokens", "session_total_completion_tokens")
+    assert all(key not in rows[2] for key in token_keys)
+
+
 @pytest.mark.asyncio
 async def test_build_ui_spend_logs_response_session_cache_hit_count():
     """

ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx

@@ -75,6 +75,41 @@ describe("Cost column", () => {
   });
 });
 
+describe("Tokens column", () => {
+  const sessionRow: Partial<LogEntry> = {
+    request_id: "req-session-tokens",
+    total_tokens: 10,
+    prompt_tokens: 7,
+    completion_tokens: 3,
+    session_id: "sess-1",
+    session_total_count: 3,
+  };
+
+  it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => {
+    const aggregatedRow: Partial<LogEntry> = {
+      ...sessionRow,
+      session_total_tokens: 60,
+      session_total_prompt_tokens: 42,
+      session_total_completion_tokens: 18,
+    };
+    renderRows([logEntry(aggregatedRow)]);
+
+    const tokensCell = screen.getByRole("cell", { name: /\(42\+18\)/ });
+    expect(tokensCell).toHaveTextContent("60");
+    expect(tokensCell).toHaveTextContent("session total");
+    expect(screen.queryByText("10")).not.toBeInTheDocument();
+    expect(screen.queryByText("(7+3)")).not.toBeInTheDocument();
+  });
+
+  it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => {
+    renderRows([logEntry(sessionRow)]);
+
+    const tokensCell = screen.getByRole("cell", { name: /\(7\+3\)/ });
+    expect(tokensCell).toHaveTextContent("10");
+    expect(tokensCell).not.toHaveTextContent("session total");
+  });
+});
+
 describe("Type column", () => {
   it("shows the conversation badge and composition even when an MCP call represents the conversation", async () => {
     const user = userEvent.setup();

ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx

@@ -263,13 +263,20 @@ export const getRequestLogsTableColumns = ({
     meta: { numeric: true },
     cell: ({ row }) => {
       const log = row.original;
+      const showSessionTotal = (log.session_total_count || 1) > 1 && log.session_total_tokens != null;
+      const total = showSessionTotal ? log.session_total_tokens : log.total_tokens;
+      const prompt = showSessionTotal ? log.session_total_prompt_tokens : log.prompt_tokens;
+      const completion = showSessionTotal ? log.session_total_completion_tokens : log.completion_tokens;
       return (
-        <span className="text-sm">
-          {String(log.total_tokens || "0")}
-          <span className="text-muted-foreground text-xs ml-1">
-            ({String(log.prompt_tokens || "0")}+{String(log.completion_tokens || "0")})
+        <div className="flex flex-col items-end">
+          <span className="text-sm">
+            {String(total || "0")}
+            <span className="text-muted-foreground text-xs ml-1">
+              ({String(prompt || "0")}+{String(completion || "0")})
+            </span>
           </span>
-        </span>
+          {showSessionTotal && <span className="text-[10px] text-muted-foreground">session total</span>}
+        </div>
       );
     },
   },

ui/litellm-dashboard/src/components/view_logs/columns.tsx

@@ -42,6 +42,9 @@ export type LogEntry = {
   request_duration_ms?: number;
   session_total_count?: number;
   session_total_spend?: number;
+  session_total_tokens?: number;
+  session_total_prompt_tokens?: number;
+  session_total_completion_tokens?: number;
   session_cache_hit_count?: number;
   mcp_tool_call_count?: number;
   mcp_tool_call_spend?: number;