Update tokenlens to canary with server side fetching (#1196)

Co-authored-by: josh <josh@afterima.ge>
2025-09-15 11:32:28 +02:00 · 2025-09-15 11:32:28 +02:00 · 5ab695262f
commit 5ab695262f
parent 445d63e620
10 changed files with 196 additions and 286 deletions
--- a/app/(chat)/api/chat/route.ts
+++ b/app/(chat)/api/chat/route.ts
@ -2,7 +2,6 @@ import {
  convertToModelMessages,
  createUIMessageStream,
  JsonToSseTransformStream,
-  type LanguageModelUsage,
  smoothStream,
  stepCountIs,
  streamText,
@ -39,11 +38,32 @@ import { ChatSDKError } from '@/lib/errors';
 import type { ChatMessage } from '@/lib/types';
 import type { ChatModel } from '@/lib/ai/models';
 import type { VisibilityType } from '@/components/visibility-selector';
+import { unstable_cache as cache } from 'next/cache';
+import { fetchModels } from 'tokenlens/fetch';
+import { getUsage } from 'tokenlens/helpers';
+import type { ModelCatalog } from 'tokenlens/core';
+import type { AppUsage } from '@/lib/usage';

 export const maxDuration = 60;

 let globalStreamContext: ResumableStreamContext | null = null;

+const getTokenlensCatalog = cache(
+  async (): Promise<ModelCatalog | undefined> => {
+    try {
+      return await fetchModels();
+    } catch (err) {
+      console.warn(
+        'TokenLens: catalog fetch failed, using default catalog',
+        err,
+      );
+      return undefined; // tokenlens helpers will fall back to defaultCatalog
+    }
+  },
+  ['tokenlens-catalog'],
+  { revalidate: 24 * 60 * 60 }, // 24 hours
+);
+
 export function getStreamContext() {
  if (!globalStreamContext) {
    try {
@ -151,7 +171,7 @@ export async function POST(request: Request) {
    const streamId = generateUUID();
    await createStreamId({ streamId, chatId: id });

-    let finalUsage: LanguageModelUsage | undefined;
+    let finalMergedUsage: AppUsage | undefined;

    const stream = createUIMessageStream({
      execute: ({ writer: dataStream }) => {
@ -183,9 +203,31 @@ export async function POST(request: Request) {
            isEnabled: isProductionEnvironment,
            functionId: 'stream-text',
          },
-          onFinish: ({ usage }) => {
-            finalUsage = usage;
-            dataStream.write({ type: 'data-usage', data: usage });
+          onFinish: async ({ usage }) => {
+            try {
+              const providers = await getTokenlensCatalog();
+              const modelId =
+                myProvider.languageModel(selectedChatModel).modelId;
+              if (!modelId) {
+                finalMergedUsage = usage;
+                dataStream.write({ type: 'data-usage', data: finalMergedUsage });
+                return;
+              }
+
+              if (!providers) {
+                finalMergedUsage = usage;
+                dataStream.write({ type: 'data-usage', data: finalMergedUsage });
+                return;
+              }
+
+              const summary = getUsage({ modelId, usage, providers });
+              finalMergedUsage = { ...usage, ...summary, modelId } as AppUsage;
+              dataStream.write({ type: 'data-usage', data: finalMergedUsage });
+            } catch (err) {
+              console.warn('TokenLens enrichment failed', err);
+              finalMergedUsage = usage;
+              dataStream.write({ type: 'data-usage', data: finalMergedUsage });
+            }
          },
        });

@ -210,11 +252,11 @@ export async function POST(request: Request) {
          })),
        });

-        if (finalUsage) {
+        if (finalMergedUsage) {
          try {
            await updateChatLastContextById({
              chatId: id,
-              context: finalUsage,
+              context: finalMergedUsage,
            });
          } catch (err) {
            console.warn('Unable to persist last usage for chat', id, err);
--- a/components/chat.tsx
+++ b/components/chat.tsx
@ -1,6 +1,6 @@
 'use client';

-import { DefaultChatTransport, type LanguageModelUsage } from 'ai';
+import { DefaultChatTransport } from 'ai';
 import { useChat } from '@ai-sdk/react';
 import { useEffect, useState } from 'react';
 import useSWR, { useSWRConfig } from 'swr';
@ -21,6 +21,7 @@ import { useChatVisibility } from '@/hooks/use-chat-visibility';
 import { useAutoResume } from '@/hooks/use-auto-resume';
 import { ChatSDKError } from '@/lib/errors';
 import type { Attachment, ChatMessage } from '@/lib/types';
+import type { AppUsage } from '@/lib/usage';
 import { useDataStream } from './data-stream-provider';
 import {
  AlertDialog,
@ -50,7 +51,7 @@ export function Chat({
  isReadonly: boolean;
  session: Session;
  autoResume: boolean;
-  initialLastContext?: LanguageModelUsage;
+  initialLastContext?: AppUsage;
 }) {
  const { visibilityType } = useChatVisibility({
    chatId: id,
@ -61,9 +62,7 @@ export function Chat({
  const { setDataStream } = useDataStream();

  const [input, setInput] = useState<string>('');
-  const [usage, setUsage] = useState<LanguageModelUsage | undefined>(
-    initialLastContext,
-  );
+  const [usage, setUsage] = useState<AppUsage | undefined>(initialLastContext);
  const [showCreditCardAlert, setShowCreditCardAlert] = useState(false);

  const {
@ -96,9 +95,7 @@ export function Chat({
    }),
    onData: (dataPart) => {
      setDataStream((ds) => (ds ? [...ds, dataPart] : []));
-      if (dataPart.type === 'data-usage') {
-        setUsage(dataPart.data);
-      }
+      if (dataPart.type === 'data-usage') setUsage(dataPart.data);
    },
    onFinish: () => {
      mutate(unstable_serialize(getChatHistoryPaginationKey));
--- a/components/elements/context.tsx
+++ b/components/elements/context.tsx
@ -7,20 +7,13 @@ import {
 } from '@/components/ui/dropdown-menu';
 import { cn } from '@/lib/utils';
 import type { ComponentProps } from 'react';
-import type { LanguageModelUsage } from 'ai';
-import { breakdownTokens, estimateCost, normalizeUsage } from 'tokenlens';
 import { Separator } from '@/components/ui/separator';
 import { Progress } from '@/components/ui/progress';
+import type { AppUsage } from '@/lib/usage';

 export type ContextProps = ComponentProps<'button'> & {
-  /** Total context window size in tokens */
-  maxTokens: number;
-  /** Tokens used so far */
-  usedTokens: number;
  /** Optional full usage payload to enable breakdown view */
-  usage?: LanguageModelUsage | undefined;
-  /** Optional model id (canonical or alias) to compute cost */
-  modelId?: string;
+  usage?: AppUsage;
 };

 const THOUSAND = 1000;
@ -34,57 +27,6 @@ const ICON_CENTER = 12;
 const ICON_RADIUS = 10;
 const ICON_STROKE_WIDTH = 2;

-const formatTokens = (tokens?: number) => {
-  if (tokens === undefined) {
-    return;
-  }
-  if (!Number.isFinite(tokens)) {
-    return;
-  }
-  const abs = Math.abs(tokens);
-  if (abs < THOUSAND) {
-    return `${tokens}`;
-  }
-  if (abs < MILLION) {
-    return `${(tokens / THOUSAND).toFixed(1)}K`;
-  }
-  if (abs < BILLION) {
-    return `${(tokens / MILLION).toFixed(1)}M`;
-  }
-  return `${(tokens / BILLION).toFixed(1)}B`;
-};
-
-const formatPercent = (value: number) => {
-  if (!Number.isFinite(value)) {
-    return '0%';
-  }
-  const rounded = Math.round(value * 10) / 10;
-  return Number.isInteger(rounded)
-    ? `${Math.trunc(rounded)}%`
-    : `${rounded.toFixed(1)}%`;
-};
-
-const formatUSD = (value?: number) => {
-  if (value === undefined || !Number.isFinite(value)) return undefined;
-  const abs = Math.abs(value);
-  // Finer precision for very small amounts common in LLM pricing
-  let decimals = 2;
-  if (abs < 0.001) decimals = 5;
-  else if (abs < 0.01) decimals = 4;
-  else if (abs < 0.1) decimals = 3;
-  else if (abs < 10) decimals = 2;
-  else decimals = 1;
-  const text = value.toFixed(decimals);
-  // Trim trailing zeros/decimal if not needed (e.g., 1.2300 -> 1.23, 2.0 -> 2)
-  const trimmed = text.replace(/\.0+$/, '').replace(/(\.\d*?)0+$/, '$1');
-  return `$${trimmed}`;
-};
-
-const formatUSDFixed = (value?: number, decimals = 5) => {
-  if (value === undefined || !Number.isFinite(value)) return undefined;
-  return `$${Number(value).toFixed(decimals)}`;
-};
-
 type ContextIconProps = {
  percent: number; // 0 - 100
 };
@ -96,7 +38,7 @@ export const ContextIcon = ({ percent }: ContextIconProps) => {

  return (
    <svg
-      aria-label={`${formatPercent(percent)} of model context used`}
+      aria-label={`${percent.toFixed(2)}% of model context used`}
      height="28"
      role="img"
      style={{ color: 'currentcolor' }}
@ -129,22 +71,6 @@ export const ContextIcon = ({ percent }: ContextIconProps) => {
  );
 };

-function TokensWithCost({
-  tokens,
-  costText,
-}: {
-  tokens?: number;
-  costText?: string;
-}) {
-  return (
-    <span>
-      {tokens === undefined ? '—' : formatTokens(tokens)}
-      {costText ? (
-        <span className="ml-2 text-muted-foreground">• {costText}</span>
-      ) : null}
-    </span>
-  );
-}

 function InfoRow({
  label,
@ -156,126 +82,36 @@ function InfoRow({
  costText?: string;
 }) {
  return (
-    <div className="flex justify-between items-center text-xs">
+    <div className="flex items-center justify-between text-xs">
      <span className="text-muted-foreground">{label}</span>
-      <TokensWithCost tokens={tokens} costText={costText} />
+      <div className="flex items-center gap-2 font-mono">
+        <span className="text-right min-w-[4ch]">
+          {tokens === undefined ? '—' : tokens.toLocaleString()}
+        </span>
+        {costText !== undefined && costText !== null && !isNaN(parseFloat(costText)) && (
+          <span className="text-muted-foreground">
+            ${parseFloat(costText).toFixed(6)}
+          </span>
+        )}
+      </div>
    </div>
  );
 }

-export const Context = ({
-  className,
-  maxTokens,
-  usedTokens,
-  usage,
-  modelId,
-  ...props
-}: ContextProps) => {
-  const safeMax = Math.max(0, Number.isFinite(maxTokens) ? maxTokens : 0);
-  const safeUsed = Math.min(
-    Math.max(0, Number.isFinite(usedTokens) ? usedTokens : 0),
-    safeMax,
-  );
-
-  // used percent and used tokens to display (demo-aware)
-  const displayUsedTokens = safeUsed;
-  const usedPercent =
-    safeMax > 0
-      ? Math.min(
-          PERCENT_MAX,
-          Math.max(0, (displayUsedTokens / safeMax) * PERCENT_MAX),
-        )
-      : 0;
-
-  const displayPct = formatPercent(Math.round(usedPercent * 10) / 10);
-
-  const used = formatTokens(displayUsedTokens);
-  const total = formatTokens(safeMax);
-
-  const uNorm = normalizeUsage(usage);
-  const uBreakdown = breakdownTokens(usage);
-
-  const hasUsage =
-    !!usage &&
-    ((uNorm.input ?? 0) > 0 ||
-      (uNorm.output ?? 0) > 0 ||
-      (uBreakdown.cacheReads ?? 0) > 0 ||
-      (uBreakdown.cacheWrites ?? 0) > 0 ||
-      (uBreakdown.reasoningTokens ?? 0) > 0);
-
-  // Values to render in rows (demo or real)
-  const displayInput = uNorm.input;
-  const displayOutput = uNorm.output;
-
-  // Per-segment costs
-  const inputCostText = modelId
-    ? formatUSDFixed(
-        estimateCost({
-          modelId,
-          usage: { input: displayInput ?? 0, output: 0 },
-        }).inputUSD,
-      )
-    : undefined;
-  const outputCostText = modelId
-    ? formatUSDFixed(
-        estimateCost({
-          modelId,
-          usage: { input: 0, output: displayOutput ?? 0 },
-        }).outputUSD,
-      )
-    : undefined;
-  // Not supported by tokenlens pricing hints; leave undefined so no bullet is shown
-  const cacheReadsTokens = uBreakdown.cacheReads ?? 0;
-  const cacheWritesTokens = uBreakdown.cacheWrites ?? 0;
-  const cacheReadsCostText =
-    modelId && cacheReadsTokens > 0
-      ? formatUSDFixed(
-          estimateCost({
-            modelId,
-            // Cast to any to support extended pricing fields provided by tokenlens
-            usage: { cacheReads: cacheReadsTokens } as any,
-          }).totalUSD,
-        )
-      : undefined;
-  const cacheWritesCostText =
-    modelId && cacheWritesTokens > 0
-      ? formatUSDFixed(
-          estimateCost({
-            modelId,
-            usage: { cacheWrites: cacheWritesTokens } as any,
-          }).totalUSD,
-        )
-      : undefined;
-
-  const reasoningTokens = uBreakdown.reasoningTokens ?? 0;
-  let reasoningCostText: string | undefined;
-  if (modelId && reasoningTokens > 0) {
-    const est = estimateCost({
-      modelId,
-      usage: { reasoningTokens },
-    }).totalUSD;
-    // TokenLens does not provide reasoning pricing for some models. Show em dash when unknown.
-    reasoningCostText =
-      est && Number.isFinite(est) && est > 0 ? formatUSDFixed(est) : '—';
-  }
-
-  const costUSD = modelId
-    ? estimateCost({
-        modelId,
-        usage: { input: displayInput ?? 0, output: displayOutput ?? 0 },
-      }).totalUSD
-    : undefined;
-  const costText = formatUSDFixed(costUSD);
-
-  const fmtOrUnknown = (n?: number) =>
-    n === undefined ? '—' : formatTokens(n);
-
+export const Context = ({ className, usage, ...props }: ContextProps) => {
+  const used = usage?.totalTokens ?? 0;
+  const max =
+    usage?.context?.totalMax ??
+    usage?.context?.combinedMax ??
+    usage?.context?.inputMax;
+  const hasMax = typeof max === 'number' && Number.isFinite(max) && max > 0;
+  const usedPercent = hasMax ? Math.min(100, (used / max) * 100) : 0;
  return (
    <DropdownMenu>
      <DropdownMenuTrigger asChild>
        <button
          className={cn(
-            'inline-flex gap-1 items-center text-sm rounded-md select-none',
+            'inline-flex items-center gap-1 select-none rounded-md text-sm',
            'cursor-pointer bg-background text-foreground',
            className,
          )}
@ -283,58 +119,63 @@ export const Context = ({
          {...props}
        >
          <span className="hidden font-medium text-muted-foreground">
-            {displayPct}
+            {usedPercent.toFixed(1)}%
          </span>
          <ContextIcon percent={usedPercent} />
        </button>
      </DropdownMenuTrigger>
-      <DropdownMenuContent align="end" side="top" className="p-3 w-fit">
+      <DropdownMenuContent align="end" side="top" className="w-fit p-3">
        <div className="min-w-[240px] space-y-2">
-          <div className="flex justify-between items-start text-sm">
-            <span>{displayPct}</span>
-            <span className="text-muted-foreground">{used} / {total} tokens</span>
+          <div className="flex items-start justify-between text-sm">
+            <span>{usedPercent.toFixed(1)}%</span>
+            <span className="text-muted-foreground">
+              {hasMax ? `${used} / ${max} tokens` : `${used} tokens`}
+            </span>
          </div>
          <div className="space-y-2">
            <Progress className="h-2 bg-muted" value={usedPercent} />
          </div>
          <div className="mt-1 space-y-1">
-            {hasUsage && uBreakdown.cacheReads && uBreakdown.cacheReads > 0 && (
+            {usage?.cachedInputTokens && usage.cachedInputTokens > 0 && (
              <InfoRow
                label="Cache Hits"
-                tokens={uBreakdown.cacheReads}
-                costText={cacheReadsCostText}
+                tokens={usage?.cachedInputTokens}
+                costText={usage?.costUSD?.cacheReadUSD?.toString()}
              />
            )}
-            {hasUsage &&
-              uBreakdown.cacheWrites &&
-              uBreakdown.cacheWrites > 0 && (
-                <InfoRow
-                  label="Cache Writes"
-                  tokens={uBreakdown.cacheWrites}
-                  costText={cacheWritesCostText}
-                />
-              )}
            <InfoRow
              label="Input"
-              tokens={displayInput}
-              costText={inputCostText}
+              tokens={usage?.inputTokens}
+              costText={usage?.costUSD?.inputUSD?.toString()}
            />
            <InfoRow
              label="Output"
-              tokens={displayOutput}
-              costText={outputCostText}
+              tokens={usage?.outputTokens}
+              costText={usage?.costUSD?.outputUSD?.toString()}
            />
            <InfoRow
              label="Reasoning"
-              tokens={reasoningTokens > 0 ? reasoningTokens : undefined}
-              costText={reasoningCostText}
+              tokens={
+                usage?.reasoningTokens && usage.reasoningTokens > 0
+                  ? usage.reasoningTokens
+                  : undefined
+              }
+              costText={usage?.costUSD?.reasoningUSD?.toString()}
            />
-            {costText && (
+            {usage?.costUSD?.totalUSD !== undefined && (
              <>
                <Separator className="mt-1" />
                <div className="flex justify-between items-center pt-1 text-xs">
                  <span className="text-muted-foreground">Total cost</span>
-                  <span>{costText}</span>
+                  <div className="flex items-center gap-2 font-mono">
+                    <span className="text-right min-w-[4ch]"></span>
+                    <span>
+                      {!isNaN(parseFloat(usage.costUSD.totalUSD.toString())) 
+                        ? `$${parseFloat(usage.costUSD.totalUSD.toString()).toFixed(6)}`
+                        : '—'
+                      }
+                    </span>
+                  </div>
                </div>
              </>
            )}
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@ -1,6 +1,6 @@
 'use client';

-import type { LanguageModelUsage, UIMessage } from 'ai';
+import type { UIMessage } from 'ai';
 import {
  useRef,
  useEffect,
@ -15,7 +15,13 @@ import {
 import { toast } from 'sonner';
 import { useLocalStorage, useWindowSize } from 'usehooks-ts';

-import { ArrowUpIcon, PaperclipIcon, CpuIcon, StopIcon, ChevronDownIcon } from './icons';
+import {
+  ArrowUpIcon,
+  PaperclipIcon,
+  CpuIcon,
+  StopIcon,
+  ChevronDownIcon,
+} from './icons';
 import { PreviewAttachment } from './preview-attachment';
 import { Button } from './ui/button';
 import { SuggestedActions } from './suggested-actions';
@ -37,10 +43,10 @@ import { ArrowDown } from 'lucide-react';
 import { useScrollToBottom } from '@/hooks/use-scroll-to-bottom';
 import type { VisibilityType } from './visibility-selector';
 import type { Attachment, ChatMessage } from '@/lib/types';
+import type { AppUsage } from '@/lib/usage';
 import { chatModels } from '@/lib/ai/models';
 import { saveChatModelAsCookie } from '@/app/(chat)/actions';
 import { startTransition } from 'react';
-import { getContextWindow, normalizeUsage } from 'tokenlens';
 import { Context } from './elements/context';
 import { myProvider } from '@/lib/ai/providers';

@ -73,7 +79,7 @@ function PureMultimodalInput({
  className?: string;
  selectedVisibilityType: VisibilityType;
  selectedModelId: string;
-  usage?: LanguageModelUsage;
+  usage?: AppUsage;
 }) {
  const textareaRef = useRef<HTMLTextAreaElement>(null);
  const { width } = useWindowSize();
@ -193,29 +199,11 @@ function PureMultimodalInput({
    return myProvider.languageModel(selectedModelId);
  }, [selectedModelId]);

-  const contextMax = useMemo(() => {
-    // Resolve from selected model; stable across chunks.
-    const cw = getContextWindow(modelResolver.modelId);
-    return cw.combinedMax ?? cw.inputMax ?? 0;
-  }, [modelResolver]);
-
-  const usedTokens = useMemo(() => {
-    // Prefer explicit usage data part captured via onData
-    if (!usage) return 0; // update only when final usage arrives
-    const n = normalizeUsage(usage);
-    return typeof n.total === 'number'
-      ? n.total
-      : (n.input ?? 0) + (n.output ?? 0);
-  }, [usage]);
-
  const contextProps = useMemo(
    () => ({
-      maxTokens: contextMax,
-      usedTokens,
      usage,
-      modelId: modelResolver.modelId,
    }),
-    [contextMax, usedTokens, usage, modelResolver],
+    [usage],
  );

  const handleFileChange = useCallback(
@ -253,7 +241,7 @@ function PureMultimodalInput({
  }, [status, scrollToBottom]);

  return (
-    <div className='flex relative flex-col gap-4 w-full'>
+    <div className="flex relative flex-col gap-4 w-full">
      <AnimatePresence>
        {!isAtBottom && (
          <motion.div
@ -261,7 +249,7 @@ function PureMultimodalInput({
            animate={{ opacity: 1, y: 0 }}
            exit={{ opacity: 0, y: 10 }}
            transition={{ type: 'spring', stiffness: 300, damping: 20 }}
-            className='absolute -top-12 left-1/2 z-50 -translate-x-1/2'
+            className="absolute -top-12 left-1/2 z-50 -translate-x-1/2"
          >
            <Button
              data-testid="scroll-to-bottom-button"
@ -299,7 +287,7 @@ function PureMultimodalInput({
      />

      <PromptInput
-        className='p-3 rounded-xl border transition-all duration-200 border-border bg-background shadow-xs focus-within:border-border hover:border-muted-foreground/50'
+        className="p-3 rounded-xl border transition-all duration-200 border-border bg-background shadow-xs focus-within:border-border hover:border-muted-foreground/50"
        onSubmit={(event) => {
          event.preventDefault();
          if (status !== 'ready') {
@ -312,7 +300,7 @@ function PureMultimodalInput({
        {(attachments.length > 0 || uploadQueue.length > 0) && (
          <div
            data-testid="attachments-preview"
-            className='flex overflow-x-scroll flex-row gap-2 items-end'
+            className="flex overflow-x-scroll flex-row gap-2 items-end"
          >
            {attachments.map((attachment) => (
              <PreviewAttachment
@ -342,7 +330,7 @@ function PureMultimodalInput({
            ))}
          </div>
        )}
-        <div className='flex flex-row gap-1 items-start sm:gap-2'>
+        <div className="flex flex-row gap-1 items-start sm:gap-2">
          <PromptInputTextarea
            data-testid="multimodal-input"
            ref={textareaRef}
@ -352,13 +340,13 @@ function PureMultimodalInput({
            minHeight={44}
            maxHeight={200}
            disableAutoResize={true}
-            className='grow resize-none border-0! p-2 border-none! bg-transparent text-sm outline-none ring-0 [-ms-overflow-style:none] [scrollbar-width:none] placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-0 focus-visible:ring-offset-0 [&::-webkit-scrollbar]:hidden'
+            className="grow resize-none border-0! p-2 border-none! bg-transparent text-sm outline-none ring-0 [-ms-overflow-style:none] [scrollbar-width:none] placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-0 focus-visible:ring-offset-0 [&::-webkit-scrollbar]:hidden"
            rows={1}
            autoFocus
          />{' '}
          <Context {...contextProps} />
        </div>
-        <PromptInputToolbar className='!border-top-0 border-t-0! p-0 shadow-none dark:border-0 dark:border-transparent!'>
+        <PromptInputToolbar className="!border-top-0 border-t-0! p-0 shadow-none dark:border-0 dark:border-transparent!">
          <PromptInputTools className="gap-0 sm:gap-0.5">
            <AttachmentsButton
              fileInputRef={fileInputRef}
@ -413,7 +401,7 @@ function PureAttachmentsButton({
  return (
    <Button
      data-testid="attachments-button"
-      className='p-1 h-8 rounded-lg transition-colors aspect-square hover:bg-accent'
+      className="p-1 h-8 rounded-lg transition-colors aspect-square hover:bg-accent"
      onClick={(event) => {
        event.preventDefault();
        fileInputRef.current?.click();
@ -454,26 +442,30 @@ function PureModelSelectorCompact({
    >
      <SelectPrimitive.Trigger
        type="button"
-        className='flex gap-2 items-center px-2 h-8 rounded-lg border-0 shadow-none transition-colors bg-background text-foreground hover:bg-accent focus:outline-none focus:ring-0 focus-visible:ring-0 focus-visible:ring-offset-0'
+        className="flex gap-2 items-center px-2 h-8 rounded-lg border-0 shadow-none transition-colors bg-background text-foreground hover:bg-accent focus:outline-none focus:ring-0 focus-visible:ring-0 focus-visible:ring-offset-0"
      >
        <CpuIcon size={16} />
-        <span className="hidden text-xs font-medium sm:block">{selectedModel?.name}</span>
+        <span className="hidden text-xs font-medium sm:block">
+          {selectedModel?.name}
+        </span>
        <ChevronDownIcon size={16} />
      </SelectPrimitive.Trigger>
      <PromptInputModelSelectContent className="min-w-[260px] p-0">
        <div className="flex flex-col gap-px">
-        {chatModels.map((model) => (
-          <SelectItem key={model.id} value={model.name} className="px-3 py-2 text-xs">
-            <div className="flex flex-col flex-1 gap-1 min-w-0">
-              <div className="text-xs font-medium truncate">
-                {model.name}
+          {chatModels.map((model) => (
+            <SelectItem
+              key={model.id}
+              value={model.name}
+              className="px-3 py-2 text-xs"
+            >
+              <div className="flex flex-col flex-1 gap-1 min-w-0">
+                <div className="text-xs font-medium truncate">{model.name}</div>
+                <div className="text-[10px] text-muted-foreground truncate leading-tight">
+                  {model.description}
+                </div>
              </div>
-              <div className="text-[10px] text-muted-foreground truncate leading-tight">
-                {model.description}
-              </div>
-            </div>
-          </SelectItem>
-        ))}
+            </SelectItem>
+          ))}
        </div>
      </PromptInputModelSelectContent>
    </PromptInputModelSelect>
--- a/lib/db/queries.ts
+++ b/lib/db/queries.ts
@ -33,7 +33,7 @@ import { generateUUID } from '../utils';
 import { generateHashedPassword } from './utils';
 import type { VisibilityType } from '@/components/visibility-selector';
 import { ChatSDKError } from '../errors';
-import type { LanguageModelV2Usage } from '@ai-sdk/provider';
+import type { AppUsage } from '../usage';

 // Optionally, if not using email/pass login, you can
 // use the Drizzle adapter for Auth.js / NextAuth
@ -479,8 +479,8 @@ export async function updateChatLastContextById({
  context,
 }: {
  chatId: string;
-  // Store raw LanguageModelUsage to keep it simple
-  context: LanguageModelV2Usage;
+  // Store merged server-enriched usage object
+  context: AppUsage;
 }) {
  try {
    return await db
--- a/lib/db/schema.ts
+++ b/lib/db/schema.ts
@ -11,7 +11,7 @@ import {
  foreignKey,
  boolean,
 } from 'drizzle-orm/pg-core';
-import type { LanguageModelV2Usage } from '@ai-sdk/provider';
+import type { AppUsage } from '../usage';

 export const user = pgTable('User', {
  id: uuid('id').primaryKey().notNull().defaultRandom(),
@ -31,7 +31,7 @@ export const chat = pgTable('Chat', {
  visibility: varchar('visibility', { enum: ['public', 'private'] })
    .notNull()
    .default('private'),
-  lastContext: jsonb('lastContext').$type<LanguageModelV2Usage | null>(),
+  lastContext: jsonb('lastContext').$type<AppUsage | null>(),
 });

 export type Chat = InferSelectModel<typeof chat>;
--- a/lib/types.ts
+++ b/lib/types.ts
@ -3,7 +3,8 @@ import type { getWeather } from './ai/tools/get-weather';
 import type { createDocument } from './ai/tools/create-document';
 import type { updateDocument } from './ai/tools/update-document';
 import type { requestSuggestions } from './ai/tools/request-suggestions';
-import type { InferUITool, LanguageModelUsage, UIMessage } from 'ai';
+import type { InferUITool, UIMessage } from 'ai';
+import type { AppUsage } from './usage';

 import type { ArtifactKind } from '@/components/artifact';
 import type { Suggestion } from './db/schema';
@ -42,7 +43,7 @@ export type CustomUIDataTypes = {
  kind: ArtifactKind;
  clear: null;
  finish: null;
-  usage: LanguageModelUsage;
+  usage: AppUsage;
 };

 export type ChatMessage = UIMessage<
--- a/lib/usage.ts
+++ b/lib/usage.ts
@ -0,0 +1,5 @@
+import type { LanguageModelUsage } from 'ai';
+import type { UsageData } from 'tokenlens/helpers';
+
+// Server-merged usage: base usage + TokenLens summary + optional modelId
+export type AppUsage = LanguageModelUsage & UsageData & { modelId?: string };
--- a/package.json
+++ b/package.json
@ -96,7 +96,7 @@
    "swr": "^2.2.5",
    "tailwind-merge": "^2.5.2",
    "tailwindcss-animate": "^1.0.7",
-    "tokenlens": "^1.1.2",
+    "tokenlens": "1.3.0-canary.5",
    "use-stick-to-bottom": "^1.1.1",
    "usehooks-ts": "^3.1.0",
    "zod": "^3.25.76"
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -240,8 +240,8 @@ importers:
        specifier: ^1.0.7
        version: 1.0.7(tailwindcss@4.1.13)
      tokenlens:
-        specifier: ^1.1.2
-        version: 1.1.2
+        specifier: 1.3.0-canary.5
+        version: 1.3.0-canary.5
      use-stick-to-bottom:
        specifier: ^1.1.1
        version: 1.1.1(react@19.0.0-rc-45804af1-20241021)
@ -2042,6 +2042,18 @@ packages:
    peerDependencies:
      tailwindcss: '>=3.0.0 || insiders || >=4.0.0-alpha.20 || >=4.0.0-beta.1'

+  '@tokenlens/core@1.0.0-beta.2':
+    resolution: {integrity: sha512-LDOicWOG0xrgWDII6gwE+Y89JjdNw8dpI0J/d/bXuJlqW8eopYhWJ4lPqePrLP+e3Q8bfpXrie2bN1/taRCozg==}
+
+  '@tokenlens/fetch@1.0.0-beta.1':
+    resolution: {integrity: sha512-YTTg7+9u5SZGPq6peeOHXZ+DFN4p5Jfs59SgUesd1i9RwCU7lXIalAW5/APCEIlfDNK8Gzk361Nq96hTs7t7CQ==}
+
+  '@tokenlens/helpers@1.0.0-beta.2':
+    resolution: {integrity: sha512-w0vdNAh3EtrPGFxTTjHWW1E+iJAvp+JwK2kghktQmqlSYk8N10La8lLb4trh3MM42gx6ytlyub+Zf0q6/txNpg==}
+
+  '@tokenlens/models@1.0.0-beta.2':
+    resolution: {integrity: sha512-QX8iTgrWb+bmjoormO1Zd2cK04bUp5ExOZ1UIAZfIICy9z8h/J1phVV9sN4767Y924HjQa6IKLW9jO+EBfp08A==}
+
  '@types/cookie@0.6.0':
    resolution: {integrity: sha512-4Kh9a6B2bQciAhf7FSuMRRkUWecJgJu9nPnx3yzpsfXX/c50REIqpHY4C82bXP90qrLtXtkDxTZosYO3UpOwlA==}

@ -4762,8 +4774,8 @@ packages:
    resolution: {integrity: sha512-65P7iz6X5yEr1cwcgvQxbbIw7Uk3gOy5dIdtZ4rDveLqhrdJP+Li/Hx6tyK0NEb+2GCyneCMJiGqrADCSNk8sQ==}
    engines: {node: '>=8.0'}

-  tokenlens@1.1.2:
-    resolution: {integrity: sha512-PICXahAiVS/2nyc30SijuYZj5wIm2G9+2U7fdvQbl5pa/hIUpQ8H+RdYOzOrSvva6GMqPM5+6cWd9Hqi81Yv3g==}
+  tokenlens@1.3.0-canary.5:
+    resolution: {integrity: sha512-NeQgyfuAIPyyaO/aVYZnnYd+lxBlpwyB19I3QQaXZMT0WHJsrXXA4tYVIjn4B2jYrXxuO7CZ7MTxk7a4UnuQ5w==}

  trim-lines@3.0.1:
    resolution: {integrity: sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg==}
@ -6461,6 +6473,21 @@ snapshots:
      postcss-selector-parser: 6.0.10
      tailwindcss: 4.1.13

+  '@tokenlens/core@1.0.0-beta.2': {}
+
+  '@tokenlens/fetch@1.0.0-beta.1':
+    dependencies:
+      '@tokenlens/core': 1.0.0-beta.2
+
+  '@tokenlens/helpers@1.0.0-beta.2':
+    dependencies:
+      '@tokenlens/core': 1.0.0-beta.2
+      '@tokenlens/fetch': 1.0.0-beta.1
+
+  '@tokenlens/models@1.0.0-beta.2':
+    dependencies:
+      '@tokenlens/core': 1.0.0-beta.2
+
  '@types/cookie@0.6.0': {}

  '@types/d3-array@3.2.1': {}
@ -9788,7 +9815,12 @@ snapshots:
    dependencies:
      is-number: 7.0.0

-  tokenlens@1.1.2: {}
+  tokenlens@1.3.0-canary.5:
+    dependencies:
+      '@tokenlens/core': 1.0.0-beta.2
+      '@tokenlens/fetch': 1.0.0-beta.1
+      '@tokenlens/helpers': 1.0.0-beta.2
+      '@tokenlens/models': 1.0.0-beta.2

  trim-lines@3.0.1: {}