From 02b58534a0e22f0c80e3a2ab924254745fe54647 Mon Sep 17 00:00:00 2001
From: Jeremy Ruppel <jeremyruppel@users.noreply.github.com>
Date: Thu, 26 Mar 2026 15:50:32 -0400
Subject: [PATCH] fix: use TokenBadges for session list row (#23619)

The sessions list row uses a bespoke token badge instead of the
`<TokenBadge />` component, so this fixes that.
---
 .../ListSessionsPage/ListSessionsRow.tsx      | 47 +++----------------
 site/src/pages/AIBridgePage/TokenBadges.tsx   |  4 +-
 2 files changed, 9 insertions(+), 42 deletions(-)

diff --git a/site/src/pages/AIBridgePage/ListSessionsPage/ListSessionsRow.tsx b/site/src/pages/AIBridgePage/ListSessionsPage/ListSessionsRow.tsx
index 68bcf73279..29625d20d4 100644
--- a/site/src/pages/AIBridgePage/ListSessionsPage/ListSessionsRow.tsx
+++ b/site/src/pages/AIBridgePage/ListSessionsPage/ListSessionsRow.tsx
@@ -7,17 +7,14 @@ import {
 	TooltipProvider,
 	TooltipTrigger,
 } from "components/Tooltip/Tooltip";
-import { ArrowDownIcon, ArrowUpIcon, ChevronRightIcon } from "lucide-react";
+import { ChevronRightIcon } from "lucide-react";
 import { AIBridgeClientIcon } from "pages/AIBridgePage/RequestLogsPage/icons/AIBridgeClientIcon";
 import { AIBridgeProviderIcon } from "pages/AIBridgePage/RequestLogsPage/icons/AIBridgeProviderIcon";
 import type { FC } from "react";
 import { DATE_FORMAT, formatDateTime } from "utils/time";
 import type { AIBridgeSession } from "#/api/typesGenerated";
-import {
-	getProviderDisplayName,
-	getProviderIconName,
-	roundTokenDisplay,
-} from "../utils";
+import { TokenBadges } from "../TokenBadges";
+import { getProviderDisplayName, getProviderIconName } from "../utils";
 
 type ListSessionsRowProps = {
 	session: AIBridgeSession;
@@ -111,40 +108,10 @@ export const ListSessionsRow: FC<ListSessionsRowProps> = ({
 			</TableCell>
 			<TableCell className="w-32">
 				<div className="flex items-center">
-					<TooltipProvider>
-						<Tooltip>
-							<TooltipTrigger asChild>
-								<Badge className="gap-0 rounded-e-none">
-									<ArrowDownIcon className="size-icon-lg flex-shrink-0" />
-									<span className="truncate min-w-0 w-full">
-										{roundTokenDisplay(
-											session.token_usage_summary.input_tokens,
-										)}
-									</span>
-								</Badge>
-							</TooltipTrigger>
-							<TooltipContent>
-								{session.token_usage_summary.input_tokens} Input Tokens
-							</TooltipContent>
-						</Tooltip>
-					</TooltipProvider>
-					<TooltipProvider>
-						<Tooltip>
-							<TooltipTrigger asChild>
-								<Badge className="gap-0 bg-surface-tertiary rounded-s-none">
-									<ArrowUpIcon className="size-icon-lg flex-shrink-0" />
-									<span className="truncate min-w-0 w-full">
-										{roundTokenDisplay(
-											session.token_usage_summary.output_tokens,
-										)}
-									</span>
-								</Badge>
-							</TooltipTrigger>
-							<TooltipContent>
-								{session.token_usage_summary.output_tokens} Output Tokens
-							</TooltipContent>
-						</Tooltip>
-					</TooltipProvider>
+					<TokenBadges
+						inputTokens={session.token_usage_summary.input_tokens}
+						outputTokens={session.token_usage_summary.output_tokens}
+					/>
 				</div>
 			</TableCell>
 			<TableCell className="w-32">
diff --git a/site/src/pages/AIBridgePage/TokenBadges.tsx b/site/src/pages/AIBridgePage/TokenBadges.tsx
index fb436b09e6..11593668d4 100644
--- a/site/src/pages/AIBridgePage/TokenBadges.tsx
+++ b/site/src/pages/AIBridgePage/TokenBadges.tsx
@@ -56,7 +56,7 @@ export const TokenBadges: FC<TokenBadgesProps> = ({
 							<div className="flex items-center justify-between gap-4">
 								<div className="text-xs text-content-secondary">Input</div>
 								<div className="text-xs text-content-secondary">
-									{roundTokenDisplay(inputTokens)}
+									{inputTokens}
 								</div>
 							</div>
 						</div>
@@ -71,7 +71,7 @@ export const TokenBadges: FC<TokenBadgesProps> = ({
 							<div className="flex items-center justify-between gap-4">
 								<div className="text-xs text-content-secondary">Output</div>
 								<div className="text-xs text-content-secondary">
-									{roundTokenDisplay(outputTokens)}
+									{outputTokens}
 								</div>
 							</div>
 						</div>