From e98e2b9b2ea2679c393b4753b688bb0bd778ad7a Mon Sep 17 00:00:00 2001
From: TJ <tracy@coder.com>
Date: Tue, 25 Aug 2026 00:21:53 -0700
Subject: [PATCH] fix(site/src/pages/AIBridgePage): align token badges with
 neighboring badge sizing (#28477)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The in/out token pills on the AI Gateway sessions pages looked out of
place next to the provider/client/model badges: they used `Badge
size="sm"` while their neighbors use the default `md`, and their arrow
icons were `size-icon-lg` (24px) — taller than the badge itself,
inflating the pill height.

> [!NOTE]
> The underlying font-size problem (`Badge sm` = 10px `text-2xs`) is
fixed globally in #28517. This PR handles what that one doesn't: the
token pills' size variant, icon sizing, and API cleanup.

## Changes

- Token pills use Badge's default `md` size, matching the
provider/client/model badges next to them.
- Arrow icons drop from `size-icon-lg` (24px) to `size-icon-xs` (14px),
the same icon size the neighboring badges use, with `gap-0.5` between
icon and count.
- Removed the unused `size` prop from `TokenBadges` — no call site ever
passed one — and dropped the now-meaningless `SizeXs`/`SizeMd` stories.

Affects every `TokenBadges` consumer: the sessions list, the session
summary card, and the timeline prompt/tool-call tables.

> 🤖 Generated by Coder Agents on behalf of @tracyjohnsonux
---
 .../pages/AIBridgePage/TokenBadges.stories.tsx   | 16 ----------------
 site/src/pages/AIBridgePage/TokenBadges.tsx      | 13 ++++---------
 2 files changed, 4 insertions(+), 25 deletions(-)

diff --git a/site/src/pages/AIBridgePage/TokenBadges.stories.tsx b/site/src/pages/AIBridgePage/TokenBadges.stories.tsx
index 109c310da27..044c97bc58c 100644
--- a/site/src/pages/AIBridgePage/TokenBadges.stories.tsx
+++ b/site/src/pages/AIBridgePage/TokenBadges.stories.tsx
@@ -40,19 +40,3 @@ export const WithMetadata: Story = {
 		},
 	},
 };
-
-export const SizeXs: Story = {
-	args: {
-		size: "xs",
-		inputTokens: 1234,
-		outputTokens: 567,
-	},
-};
-
-export const SizeMd: Story = {
-	args: {
-		size: "md",
-		inputTokens: 1234,
-		outputTokens: 567,
-	},
-};
diff --git a/site/src/pages/AIBridgePage/TokenBadges.tsx b/site/src/pages/AIBridgePage/TokenBadges.tsx
index dd4b39ce562..50e30a8797b 100644
--- a/site/src/pages/AIBridgePage/TokenBadges.tsx
+++ b/site/src/pages/AIBridgePage/TokenBadges.tsx
@@ -11,14 +11,12 @@ import { JsonPrettyPrinter } from "./JsonPrettyPrinter";
 import { roundTokenDisplay } from "./utils";
 
 interface TokenBadgesProps {
-	size?: "xs" | "sm" | "md";
 	inputTokens: number;
 	outputTokens: number;
 	tokenUsageMetadata?: Record<string, unknown>;
 }
 
 export const TokenBadges: FC<TokenBadgesProps> = ({
-	size = "sm",
 	inputTokens,
 	outputTokens,
 	tokenUsageMetadata,
@@ -28,17 +26,14 @@ export const TokenBadges: FC<TokenBadgesProps> = ({
 			<Tooltip>
 				<TooltipTrigger asChild>
 					<span>
-						<Badge className="gap-0 rounded-e-none" size={size}>
-							<ArrowDownIcon className="size-icon-lg flex-shrink-0" />
+						<Badge className="gap-0.5 rounded-e-none">
+							<ArrowDownIcon className="size-icon-xs flex-shrink-0" />
 							<span className="truncate min-w-0">
 								{roundTokenDisplay(inputTokens)}
 							</span>
 						</Badge>
-						<Badge
-							className="gap-0 bg-surface-tertiary rounded-s-none"
-							size={size}
-						>
-							<ArrowUpIcon className="size-icon-lg flex-shrink-0" />
+						<Badge className="gap-0.5 bg-surface-tertiary rounded-s-none">
+							<ArrowUpIcon className="size-icon-xs flex-shrink-0" />
 							<span className="truncate min-w-0">
 								{roundTokenDisplay(outputTokens)}
 							</span>