From d8912dd9fead25ab0aad4b3546a6e13703b06558 Mon Sep 17 00:00:00 2001
From: ephraimduncan <ephraimduncan68@gmail.com>
Date: Fri, 31 Jul 2026 14:50:06 +0000
Subject: [PATCH] fix(elements): include reasoning and cache costs in context
 total

---
 packages/elements/__tests__/context.test.tsx | 54 ++++++++++++++++++++
 packages/elements/src/context.tsx            |  7 ++-
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/packages/elements/__tests__/context.test.tsx b/packages/elements/__tests__/context.test.tsx
index 2c2d6fbc..91204674 100644
--- a/packages/elements/__tests__/context.test.tsx
+++ b/packages/elements/__tests__/context.test.tsx
@@ -1,4 +1,5 @@
 import { render, screen } from "@testing-library/react";
+import { getUsage } from "tokenlens";
 
 import {
   Context,
@@ -116,6 +117,11 @@ describe("contextContentBody", () => {
   });
 });
 
+const totalCostUSD = (
+  modelId: string,
+  usage: Parameters<typeof getUsage>[0]["usage"]
+) => getUsage({ modelId, usage }).costUSD?.totalUSD ?? 0;
+
 describe("contextContentFooter", () => {
   it("renders default footer with cost", () => {
     render(
@@ -138,6 +144,54 @@ describe("contextContentFooter", () => {
     );
     expect(screen.getByText("Custom Footer")).toBeInTheDocument();
   });
+
+  it("includes reasoning and cache costs in the total", () => {
+    // Token counts chosen so omitting any single category changes the
+    // rendered (cent-rounded) total for this model's pricing.
+    const modelId = "vercel:xai/grok-3-mini-fast";
+    const usage = {
+      cachedInputTokens: 40_000_000,
+      inputTokens: 1_000_000,
+      outputTokens: 500_000,
+      reasoningTokens: 2_000_000,
+    };
+    const expectedTotal = totalCostUSD(modelId, {
+      cacheReads: usage.cachedInputTokens,
+      input: usage.inputTokens,
+      output: usage.outputTokens,
+      reasoningTokens: usage.reasoningTokens,
+    });
+    const rowTotals =
+      totalCostUSD(modelId, { input: usage.inputTokens, output: 0 }) +
+      totalCostUSD(modelId, { input: 0, output: usage.outputTokens }) +
+      totalCostUSD(modelId, { reasoningTokens: usage.reasoningTokens }) +
+      totalCostUSD(modelId, {
+        cacheReads: usage.cachedInputTokens,
+        input: 0,
+        output: 0,
+      });
+    expect(expectedTotal).toBeGreaterThan(0);
+    expect(expectedTotal).toBeCloseTo(rowTotals, 10);
+
+    render(
+      <Context
+        defaultOpen
+        maxTokens={100_000_000}
+        modelId={modelId}
+        usage={usage}
+        usedTokens={43_500_000}
+      >
+        <ContextContent>
+          <ContextContentFooter />
+        </ContextContent>
+      </Context>
+    );
+    const formattedTotal = new Intl.NumberFormat("en-US", {
+      currency: "USD",
+      style: "currency",
+    }).format(expectedTotal);
+    expect(screen.getByText(formattedTotal)).toBeInTheDocument();
+  });
 });
 
 describe("contextInputUsage", () => {
diff --git a/packages/elements/src/context.tsx b/packages/elements/src/context.tsx
index 53e28851..7c6e9659 100644
--- a/packages/elements/src/context.tsx
+++ b/packages/elements/src/context.tsx
@@ -1,5 +1,8 @@
 "use client";
 
+import type { LanguageModelUsage } from "ai";
+import type { ComponentProps } from "react";
+
 import { Button } from "@repo/shadcn-ui/components/ui/button";
 import {
   HoverCard,
@@ -8,8 +11,6 @@ import {
 } from "@repo/shadcn-ui/components/ui/hover-card";
 import { Progress } from "@repo/shadcn-ui/components/ui/progress";
 import { cn } from "@repo/shadcn-ui/lib/utils";
-import type { LanguageModelUsage } from "ai";
-import type { ComponentProps } from "react";
 import { createContext, useContext, useMemo } from "react";
 import { getUsage } from "tokenlens";
 
@@ -201,8 +202,10 @@ export const ContextContentFooter = ({
     ? getUsage({
         modelId,
         usage: {
+          cacheReads: usage?.cachedInputTokens ?? 0,
           input: usage?.inputTokens ?? 0,
           output: usage?.outputTokens ?? 0,
+          reasoningTokens: usage?.reasoningTokens ?? 0,
         },
       }).costUSD?.totalUSD
     : undefined;