Merge 29ee4ffa87 into 2dfeb80e04

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -91,20 +91,24 @@ _SESSION_KEY_EXPR: Final = "COALESCE(NULLIF(session_id, ''), request_id)"
 _SESSION_GROUP_KEY_SQL: Final = f"{_SESSION_KEY_EXPR}, api_key"
 _MCP_CALL_TYPES_SQL: Final = "('call_mcp_tool', 'list_mcp_tools')"
 _AGENT_CALL_TYPE_SQL: Final = "'asend_message'"
+_BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')"
+_AGENT_ROW_SQL: Final = (
+    f"(call_type = {_AGENT_CALL_TYPE_SQL} OR (COALESCE(agent_id, '') <> '' "
+    f"AND call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type NOT IN {_BATCH_CALL_TYPES_SQL}))"
+)
 _SESSION_REPRESENTATIVE_ORDER_SQL: Final = (
     f"(call_type = {_AGENT_CALL_TYPE_SQL}) DESC, "
     f'CASE WHEN call_type = {_AGENT_CALL_TYPE_SQL} THEN "endTime" END DESC NULLS LAST, '
     f'call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC, request_id'
 )
-_BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')"
 _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType(
     {
         "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}",
-        "agent": f"call_type = {_AGENT_CALL_TYPE_SQL}",
+        "agent": _AGENT_ROW_SQL,
         "batch": f"call_type IN {_BATCH_CALL_TYPES_SQL}",
         "llm": (
-            f"(call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} "
-            f"AND call_type NOT IN {_BATCH_CALL_TYPES_SQL})"
+            f"(call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type NOT IN {_BATCH_CALL_TYPES_SQL} "
+            f"AND NOT {_AGENT_ROW_SQL})"
         ),
     }
 )
@@ -4729,9 +4733,9 @@ async def _build_ui_spend_logs_response(
                            ), 0)::double precision AS mcp_tool_call_spend,
                            COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count,
                            COUNT(*) FILTER (
-                               WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL}
+                               WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND NOT {_AGENT_ROW_SQL}
                            )::int AS session_llm_count,
-                           COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count,
+                           COUNT(*) FILTER (WHERE {_AGENT_ROW_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

tests/integration/spend/test_agent_spend.py

@@ -1,13 +1,14 @@
 import json
 import uuid
+from datetime import datetime, timedelta, timezone
 from hashlib import sha256
 from typing import Final
 
 import pytest
 from integration._support.client import Gateway, eventually
-from integration._support.database import read_rows
+from integration._support.database import read_rows, write_rows
 from integration._support.wire import Reply, Request, wire_server
-from pydantic import BaseModel, JsonValue
+from pydantic import BaseModel, JsonValue, TypeAdapter
 
 
 class _ActivityMetrics(BaseModel):
@@ -127,7 +128,12 @@ def _assert_a2a_calls_are_billed(
         key: Final = scenario.key()
         digest: Final = sha256(key.encode()).hexdigest()
         message: Final = (
-            {"kind": "message", "role": "user", "messageId": marker + "-in", "parts": [{"kind": "text", "text": "ping"}]}
+            {
+                "kind": "message",
+                "role": "user",
+                "messageId": marker + "-in",
+                "parts": [{"kind": "text", "text": "ping"}],
+            }
             if version == "0.3"
             else {"role": "ROLE_USER", "messageId": marker + "-in", "parts": [{"text": "ping"}]}
         )
@@ -146,9 +152,7 @@ def _assert_a2a_calls_are_billed(
                 **({"a2a-version": version} if version == "1.0" else {}),
             }
             if "stream" not in method.lower():
-                sent = gateway.client.post(
-                    path.format(agent=agent, name=marker), headers=headers, json=payload
-                )
+                sent = gateway.client.post(path.format(agent=agent, name=marker), headers=headers, json=payload)
                 assert sent.status_code == 200, sent.text
                 result: Final = sent.json()["result"]
                 parts: Final = result["parts"] if version == "0.3" else result["message"]["parts"]
@@ -275,3 +279,225 @@ def test_a2a_stream_bills_input_and_output_tokens_at_the_agent_per_token_prices(
         "Part RootModels without dumping them and ignores artifact-update text, while message/send bills 1.125"
     )
     _assert_a2a_calls_are_billed(gateway, _PER_TOKEN, _PER_TOKEN_COST, (("message/stream", "/a2a/{agent}"),))
+
+
+def test_agent_bound_chat_completion_is_classified_as_an_agent_in_ui_spend_logs(gateway: Gateway) -> None:
+    marker: Final = "agentlog" + uuid.uuid4().hex[:12]
+    agent_session: Final = f"{marker}-agent-session"
+    plain_session: Final = f"{marker}-plain-session"
+    agent_prompt: Final = f"{marker} agent-bound completion"
+    plain_prompt: Final = f"{marker} plain completion"
+
+    def upstream(request: Request) -> Reply:
+        assert request.method == "POST" and request.target == "/v1/chat/completions", request
+        body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.body)
+        assert body["model"] == "gpt-4o-mini", body
+        assert body["messages"] in (
+            [{"role": "user", "content": agent_prompt}],
+            [{"role": "user", "content": plain_prompt}],
+        ), body
+        return Reply(
+            body=json.dumps(
+                {
+                    "id": f"chatcmpl-{uuid.uuid4().hex}",
+                    "object": "chat.completion",
+                    "created": 1,
+                    "model": "gpt-4o-mini",
+                    "choices": [
+                        {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
+                    ],
+                    "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7},
+                }
+            ).encode()
+        )
+
+    with wire_server(upstream) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            api_base=f"{wire.url}/v1",
+            input_cost_per_token=0.001,
+            output_cost_per_token=0.002,
+        )
+        created: Final = gateway.request(
+            "POST",
+            "/v1/agents",
+            {"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker)},
+        )
+        assert created.status_code == 200, created.text
+        created_body: Final = TypeAdapter(dict[str, JsonValue]).validate_python(created.json())
+        agent_id: Final = TypeAdapter(str).validate_python(created_body["agent_id"])
+
+        def cleanup_agent() -> None:
+            deleted: Final = gateway.request("DELETE", f"/v1/agents/{agent_id}")
+            assert deleted.status_code == 200, deleted.text
+            assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent_id,)) == []
+
+        scenario.cleanups.callback(cleanup_agent)
+        agent_key: Final = scenario.key(agent_id=agent_id, models=[model])
+        plain_key: Final = scenario.key(models=[model])
+        agent_digest: Final = sha256(agent_key.encode()).hexdigest()
+        plain_digest: Final = sha256(plain_key.encode()).hexdigest()
+        body: Final = {"model": model, "messages": [{"role": "user", "content": agent_prompt}]}
+        agent_response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            body,
+            key=agent_key,
+            headers={"x-litellm-session-id": agent_session},
+        )
+        assert agent_response.status_code == 200, agent_response.text
+        plain_response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {**body, "messages": [{"role": "user", "content": plain_prompt}]},
+            key=plain_key,
+            headers={"x-litellm-session-id": plain_session},
+        )
+        assert plain_response.status_code == 200, plain_response.text
+
+        spend_rows: Final = eventually(
+            lambda: read_rows(
+                'SELECT request_id, session_id, agent_id, api_key, call_type FROM "LiteLLM_SpendLogs" '
+                "WHERE session_id IN (%s, %s) ORDER BY session_id",
+                (agent_session, plain_session),
+            ),
+            lambda rows: len(rows) == 2,
+            seconds=70,
+        )
+        rows_by_session: Final = {str(row["session_id"]): row for row in spend_rows}
+        assert set(rows_by_session) == {agent_session, plain_session}, spend_rows
+        assert (
+            rows_by_session[agent_session]["api_key"],
+            rows_by_session[agent_session]["call_type"],
+            rows_by_session[agent_session]["agent_id"],
+        ) == (agent_digest, "acompletion", agent_id), spend_rows
+        assert (
+            rows_by_session[plain_session]["api_key"],
+            rows_by_session[plain_session]["call_type"],
+            rows_by_session[plain_session]["agent_id"],
+        ) == (plain_digest, "acompletion", None), spend_rows
+        outbound_requests: Final = wire.drain()
+        assert [(request.method, request.target) for request in outbound_requests] == [
+            ("POST", "/v1/chat/completions"),
+            ("POST", "/v1/chat/completions"),
+        ]
+
+        today: Final = datetime.now(timezone.utc)
+        date_window: Final = {
+            "start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
+            "end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
+        }
+
+        def ui_rows(session_id: str, span_type: str) -> list[dict[str, JsonValue]]:
+            response: Final = gateway.request(
+                "GET",
+                "/spend/logs/ui",
+                params={**date_window, "session_id": session_id, "span_type": span_type},
+            )
+            assert response.status_code == 200, response.text
+            response_body: Final = TypeAdapter(dict[str, JsonValue]).validate_python(response.json())
+            return TypeAdapter(list[dict[str, JsonValue]]).validate_python(response_body["data"])
+
+        agent_rows: Final = ui_rows(agent_session, "agent")
+        agent_llm_rows: Final = ui_rows(agent_session, "llm")
+        plain_agent_rows: Final = ui_rows(plain_session, "agent")
+        plain_llm_rows: Final = ui_rows(plain_session, "llm")
+        assert [row["request_id"] for row in agent_rows] == [rows_by_session[agent_session]["request_id"]], agent_rows
+        assert agent_llm_rows == [], agent_llm_rows
+        assert plain_agent_rows == [], plain_agent_rows
+        assert [row["request_id"] for row in plain_llm_rows] == [rows_by_session[plain_session]["request_id"]], (
+            plain_llm_rows
+        )
+        assert (
+            agent_rows[0]["agent_id"],
+            agent_rows[0]["session_agent_count"],
+            agent_rows[0]["session_llm_count"],
+        ) == (agent_id, 1, 0), agent_rows
+        assert (
+            plain_llm_rows[0]["session_llm_count"],
+            plain_llm_rows[0]["session_agent_count"],
+        ) == (1, 0), plain_llm_rows
+
+
+def test_spend_log_ui_span_filters_classify_agent_bound_rows(gateway: Gateway) -> None:
+    marker: Final = "agentspan" + uuid.uuid4().hex[:12]
+    agent_id: Final = f"agent-{marker}"
+    api_key: Final = f"test-key-{marker}"
+    logged_at: Final = datetime.now(timezone.utc)
+    timestamp: Final = logged_at.strftime("%Y-%m-%d %H:%M:%S")
+    date_window: Final = {
+        "start_date": (logged_at - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
+        "end_date": (logged_at + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
+    }
+    rows_to_seed: Final = (
+        (f"{marker}-mcp", "call_mcp_tool", agent_id, f"{marker}-mcp-session"),
+        (f"{marker}-batch", "aretrieve_batch", agent_id, f"{marker}-batch-session"),
+        (f"{marker}-agent", "acompletion", agent_id, f"{marker}-agent-session"),
+        (f"{marker}-llm", "acompletion", "", f"{marker}-llm-session"),
+    )
+    insert_sql: Final = (
+        'INSERT INTO "LiteLLM_SpendLogs" '
+        '(request_id, call_type, api_key, "startTime", "endTime", agent_id, session_id, status) '
+        "VALUES (%s, %s, %s, %s::timestamp, %s::timestamp, %s, %s, 'success')"
+    )
+
+    with gateway.scenario() as scenario:
+        scenario.cleanups.callback(
+            write_rows,
+            'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s',
+            (f"{marker}-%",),
+        )
+        for request_id, call_type, row_agent_id, session_id in rows_to_seed:
+            write_rows(
+                insert_sql,
+                (request_id, call_type, api_key, timestamp, timestamp, row_agent_id, session_id),
+            )
+
+        def ui_rows(span_type: str) -> list[dict[str, JsonValue]]:
+            response: Final = gateway.request(
+                "GET",
+                "/spend/logs/ui",
+                params={**date_window, "session_id": marker, "span_type": span_type},
+            )
+            assert response.status_code == 200, response.text
+            response_body: Final = TypeAdapter(dict[str, JsonValue]).validate_python(response.json())
+            return TypeAdapter(list[dict[str, JsonValue]]).validate_python(response_body["data"])
+
+        rows_by_span: Final = {
+            span_type: ui_rows(span_type) for span_type in ("agent", "llm", "mcp", "batch")
+        }
+        assert {
+            span_type: frozenset(str(row["request_id"]) for row in rows)
+            for span_type, rows in rows_by_span.items()
+        } == {
+            "agent": frozenset({f"{marker}-agent"}),
+            "llm": frozenset({f"{marker}-llm"}),
+            "mcp": frozenset({f"{marker}-mcp"}),
+            "batch": frozenset({f"{marker}-batch"}),
+        }, rows_by_span
+
+        mcp_row: Final = rows_by_span["mcp"][0]
+        batch_row: Final = rows_by_span["batch"][0]
+        agent_row: Final = rows_by_span["agent"][0]
+        llm_row: Final = rows_by_span["llm"][0]
+        assert (
+            mcp_row["session_agent_count"],
+            mcp_row["session_llm_count"],
+            mcp_row.get("mcp_tool_call_count", 0),
+        ) == (0, 0, 1), mcp_row
+        assert (
+            batch_row["session_agent_count"],
+            batch_row["session_llm_count"],
+            batch_row.get("mcp_tool_call_count", 0),
+        ) == (0, 1, 0), batch_row
+        assert (
+            agent_row["agent_id"],
+            agent_row["session_agent_count"],
+            agent_row["session_llm_count"],
+            agent_row.get("mcp_tool_call_count", 0),
+        ) == (agent_id, 1, 0, 0), agent_row
+        assert (
+            llm_row["agent_id"],
+            llm_row["session_agent_count"],
+            llm_row["session_llm_count"],
+            llm_row.get("mcp_tool_call_count", 0),
+        ) == ("", 0, 1, 0), llm_row

tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py

@@ -105,6 +105,9 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
             "lte": _iso(params[end_index - 1]),
         }
     remaining = clause.group(1) if search_clause is None else clause.group(1).replace(search_clause.group(0), "")
+    agent_row_predicate_present: Final = "COALESCE(agent_id, '') <> ''" in remaining
+    if agent_row_predicate_present:
+        where["span_type"] = "llm" if "AND NOT (call_type = 'asend_message'" in remaining else "agent"
 
     eq_cols = {
         "team_id": "team_id",
@@ -143,7 +146,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
             where["cache_hit"] = "hit"
         elif cond == "(cache_hit IS NULL OR LOWER(cache_hit) != 'true')":
             where["cache_hit"] = "miss"
-        elif "call_type" in cond:
+        elif "call_type" in cond and not agent_row_predicate_present:
             if "call_type NOT IN" in cond:
                 where["span_type"] = "llm"
             elif "call_mcp_tool" in cond:

ui/litellm-dashboard/src/components/logs/constants.test.ts

@@ -0,0 +1,25 @@
+import { describe, expect, it } from "vitest";
+
+import { AGENT_CALL_TYPES, BATCH_CALL_TYPES, isAgentLog, MCP_CALL_TYPES } from "./constants";
+
+describe("isAgentLog", () => {
+  it.each(AGENT_CALL_TYPES)("recognizes %s without an agent ID", (callType) => {
+    expect(isAgentLog({ call_type: callType })).toBe(true);
+  });
+
+  it.each(["acompletion", "anthropic_messages"] as const)("recognizes %s with an agent ID", (callType) => {
+    expect(isAgentLog({ call_type: callType, agent_id: "agent-1" })).toBe(true);
+  });
+
+  it.each([undefined, null, ""] as const)("does not treat %s as an agent ID", (agentId) => {
+    expect(isAgentLog({ call_type: "acompletion", agent_id: agentId })).toBe(false);
+  });
+
+  it.each(MCP_CALL_TYPES)("excludes MCP call type %s with an agent ID", (callType) => {
+    expect(isAgentLog({ call_type: callType, agent_id: "agent-1" })).toBe(false);
+  });
+
+  it.each(BATCH_CALL_TYPES)("excludes batch call type %s with an agent ID", (callType) => {
+    expect(isAgentLog({ call_type: callType, agent_id: "agent-1" })).toBe(false);
+  });
+});

ui/litellm-dashboard/src/components/logs/constants.ts

@@ -1,3 +1,5 @@
+import type { LogEntry } from "./types";
+
 export const ERROR_CODE_OPTIONS: { label: string; value: string }[] = [
   { label: "400 - Bad Request", value: "400" },
   { label: "401 - Invalid Authentication", value: "401" },
@@ -21,6 +23,12 @@ export const AGENT_CALL_TYPES = ["asend_message"];
 /** Call types that represent Batch API operations (creation and retrieval, sync and async). */
 export const BATCH_CALL_TYPES = ["acreate_batch", "create_batch", "aretrieve_batch", "retrieve_batch"];
 
+export const isAgentLog = (log: Pick<LogEntry, "call_type" | "agent_id">): boolean => {
+  if (AGENT_CALL_TYPES.includes(log.call_type)) return true;
+  if (!log.agent_id) return false;
+  return !MCP_CALL_TYPES.includes(log.call_type) && !BATCH_CALL_TYPES.includes(log.call_type);
+};
+
 export const SPAN_TYPE_LABELS: Record<string, string> = {
   llm: "LLM",
   agent: "Agent",

ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx

@@ -80,8 +80,8 @@ const sessionLogs = [
   }),
 ];
 
-const renderSessionDrawer = () => {
-  vi.mocked(sessionSpendLogsCall).mockResolvedValue({ data: sessionLogs, total: 4, total_pages: 1 });
+const renderSessionDrawer = (logs: LogEntry[] = sessionLogs) => {
+  vi.mocked(sessionSpendLogsCall).mockResolvedValue({ data: logs, total: logs.length, total_pages: 1 });
   const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
   const drawer = (open: boolean) => (
     <QueryClientProvider client={queryClient}>
@@ -129,6 +129,67 @@ describe("LogDetailsDrawer session sidebar sorting", () => {
   });
 });
 
+describe("LogDetailsDrawer agent-bound logs", () => {
+  it("counts agent-key completions as Agent while keeping ordinary LLM and MCP rows separate", async () => {
+    const agentBoundLog: Partial<LogEntry> = {
+      request_id: "agent-bound",
+      model: "agent-bound-model",
+      call_type: "acompletion",
+      agent_id: "agent-1",
+    };
+    const ordinaryLlmLog: Partial<LogEntry> = {
+      request_id: "ordinary-llm",
+      model: "ordinary-llm-model",
+      call_type: "acompletion",
+    };
+    const mcpLogWithAgentId: Partial<LogEntry> = {
+      request_id: "mcp-agent-id",
+      model: "mcp-model",
+      call_type: "call_mcp_tool",
+      agent_id: "agent-2",
+    };
+    renderSessionDrawer([makeLog(agentBoundLog), makeLog(ordinaryLlmLog), makeLog(mcpLogWithAgentId)]);
+
+    expect(await screen.findByText("agent-bound-model")).toBeInTheDocument();
+    expect(screen.getByText("1 LLM")).toBeInTheDocument();
+    expect(screen.getByText("1 Agent")).toBeInTheDocument();
+    expect(screen.getByText("1 MCP")).toBeInTheDocument();
+
+    const agentRow = screen.getByRole("button", { name: /agent-bound-model/ });
+    const llmRow = screen.getByRole("button", { name: /ordinary-llm-model/ });
+    const mcpRow = screen.getByRole("button", { name: /mcp-model/ });
+    expect(agentRow.querySelector(".lucide-bot")).not.toBeNull();
+    expect(llmRow.querySelector(".lucide-sparkles")).not.toBeNull();
+    expect(mcpRow.querySelector(".lucide-wrench")).not.toBeNull();
+  });
+
+  it("keeps batch and empty-agent-id completions out of the Agent count", async () => {
+    const batchLogWithAgentId: Partial<LogEntry> = {
+      request_id: "batch-agent-id",
+      model: "batch-model",
+      call_type: "acreate_batch",
+      agent_id: "agent-2",
+    };
+    const completionWithEmptyAgentId: Partial<LogEntry> = {
+      request_id: "empty-agent-id",
+      model: "empty-agent-id-model",
+      call_type: "acompletion",
+      agent_id: "",
+    };
+    const agentBoundLog: Partial<LogEntry> = {
+      request_id: "agent-bound",
+      model: "agent-bound-model",
+      call_type: "acompletion",
+      agent_id: "agent-1",
+    };
+    renderSessionDrawer([makeLog(batchLogWithAgentId), makeLog(completionWithEmptyAgentId), makeLog(agentBoundLog)]);
+
+    expect(await screen.findByText("batch-model")).toBeInTheDocument();
+    expect(screen.getByText("2 LLM")).toBeInTheDocument();
+    expect(screen.getByText("1 Agent")).toBeInTheDocument();
+  });
+});
+
 describe("LogDetailsDrawer internal user email", () => {
   const renderSingleLog = (user: string | undefined) => {
     const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });

ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx

@@ -4,7 +4,7 @@ import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet";
 import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { LogEntry } from "../types";
 import { AutoRouterIcon, useIsAutoRoutedModelGroup } from "@/components/shared/table_cells";
-import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "../constants";
+import { isAgentLog, MCP_CALL_TYPES } from "../constants";
 import { getEventDisplayName } from "./eventDisplayName";
 import { ClassifyTag } from "./ClassifyTag";
 import { type DrawerBackTo, DrawerHeader } from "./DrawerHeader";
@@ -53,9 +53,15 @@ interface TraceEventRowProps {
 
 const TRACE_EVENT_ICON_CLASS = "text-muted-foreground shrink-0";
 
-function TraceEventIcon({ callType, isAutoRouted }: { callType: string; isAutoRouted: boolean }) {
-  if (MCP_CALL_TYPES.includes(callType)) return <Wrench size={12} className={TRACE_EVENT_ICON_CLASS} />;
-  if (AGENT_CALL_TYPES.includes(callType)) return <Bot size={12} className={TRACE_EVENT_ICON_CLASS} />;
+function TraceEventIcon({
+  log,
+  isAutoRouted,
+}: {
+  log: Pick<LogEntry, "call_type" | "agent_id">;
+  isAutoRouted: boolean;
+}) {
+  if (MCP_CALL_TYPES.includes(log.call_type)) return <Wrench size={12} className={TRACE_EVENT_ICON_CLASS} />;
+  if (isAgentLog(log)) return <Bot size={12} className={TRACE_EVENT_ICON_CLASS} />;
   if (isAutoRouted) return <AutoRouterIcon size={12} className={TRACE_EVENT_ICON_CLASS} />;
   return <Sparkles size={12} className={TRACE_EVENT_ICON_CLASS} />;
 }
@@ -76,7 +82,7 @@ function TraceEventRow({ row, isSelected, onClick }: TraceEventRowProps) {
       onClick={onClick}
     >
       <div className="flex items-center gap-1">
-        <TraceEventIcon callType={row.call_type} isAutoRouted={isAutoRouted} />
+        <TraceEventIcon log={row} isAutoRouted={isAutoRouted} />
         <span className="text-xs font-medium text-foreground truncate">
           {getEventDisplayName(row.call_type, row.model)}
         </span>
@@ -277,10 +283,8 @@ export function LogDetailsDrawer({
     sessionLogs.length > 0 ? new Date(Math.max(...sessionLogs.map((r) => new Date(r.endTime).getTime()))) : null;
   const sessionDurationSeconds =
     sessionStart && sessionEnd ? ((sessionEnd.getTime() - sessionStart.getTime()) / 1000).toFixed(2) : "0.00";
-  const llmCount = sessionLogs.filter(
-    (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type),
-  ).length;
-  const agentCount = sessionLogs.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length;
+  const llmCount = sessionLogs.filter((row) => !MCP_CALL_TYPES.includes(row.call_type) && !isAgentLog(row)).length;
+  const agentCount = sessionLogs.filter((row) => isAgentLog(row)).length;
   const mcpCount = sessionLogs.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length;
   const cacheHitCount = sessionLogs.filter((row) => String(row.cache_hit ?? "").toLowerCase() === "true").length;
   const logsForList = isSessionMode ? sessionLogs : currentLog ? [currentLog] : [];
@@ -350,12 +354,8 @@ export function LogDetailsDrawer({
                   {[
                     isSessionMode
                       ? llmCount
-                      : logsForList.filter(
-                          (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type),
-                        ).length,
-                    isSessionMode
-                      ? agentCount
-                      : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length,
+                      : logsForList.filter((row) => !MCP_CALL_TYPES.includes(row.call_type) && !isAgentLog(row)).length,
+                    isSessionMode ? agentCount : logsForList.filter((row) => isAgentLog(row)).length,
                     isSessionMode
                       ? mcpCount
                       : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length,

ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.integration.test.tsx

@@ -249,6 +249,39 @@ describe("Type column", () => {
     expect(screen.getByText("MCP")).toBeInTheDocument();
   });
 
+  it("uses agent identity for the badge and fallback session counts", async () => {
+    const user = userEvent.setup();
+    const agentKeyLog: Partial<LogEntry> = {
+      request_id: "req-agent-key",
+      call_type: "acompletion",
+      agent_id: "agent-1",
+      session_total_count: 3,
+    };
+    renderRows([logEntry(agentKeyLog)]);
+
+    const sessionCount = screen.getByText("3");
+    await user.hover(sessionCount);
+    expect(await screen.findByText("3 Agent")).toBeInTheDocument();
+  });
+
+  it("keeps ordinary completions on the LLM badge without agent identity", () => {
+    renderRows([logEntry({ request_id: "req-llm", call_type: "acompletion", session_total_count: 1 })]);
+
+    expect(screen.getByText("LLM")).toBeInTheDocument();
+  });
+
+  it("keeps MCP rows on the MCP badge when they carry an agent id", () => {
+    const mcpLogWithAgentId: Partial<LogEntry> = {
+      request_id: "req-mcp-agent-id",
+      call_type: "call_mcp_tool",
+      agent_id: "agent-1",
+      session_total_count: 1,
+    };
+    renderRows([logEntry(mcpLogWithAgentId)]);
+
+    expect(screen.getByText("MCP")).toBeInTheDocument();
+  });
+
   it("marks a batch cost row with the Batch badge instead of LLM", () => {
     renderRows([logEntry({ request_id: "batch_1_batch_cost", call_type: "aretrieve_batch" })]);
 

ui/litellm-dashboard/src/components/logs/request/RequestLogsTableColumns.tsx

@@ -10,7 +10,7 @@ import { copyToClipboard, getSpendString } from "@/utils/dataUtils";
 import { getProviderLogoAndName } from "../../provider_info_helpers";
 import { getBatchIdFromRequestId, getBatchRequestCounts, isBatchCallType } from "../batchLogUtils";
 import type { LogEntry } from "../types";
-import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "../constants";
+import { isAgentLog, MCP_CALL_TYPES } from "../constants";
 import { AgentBadge, AgentIcon, BatchBadge, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges";
 
 export interface RequestLogsTableColumnsDeps {
@@ -94,7 +94,7 @@ export const getRequestLogsTableColumns = ({
       const log = row.original;
       const sessionCount = log.session_total_count || 1;
       const isMcp = MCP_CALL_TYPES.includes(log.call_type);
-      const isAgent = AGENT_CALL_TYPES.includes(log.call_type);
+      const isAgent = isAgentLog(log);
       const sessionLlmCount = log.session_llm_count ?? (isMcp || isAgent ? 0 : sessionCount);
       const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0);
       const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0);

ui/litellm-dashboard/src/components/logs/types.ts

@@ -20,6 +20,7 @@ export type LogEntry = {
   model_group?: string;
   api_base?: string;
   call_type: string;
+  agent_id?: string | null;
   spend: number;
   total_tokens: number;
   prompt_tokens: number;