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

Co-authored-by: josh <josh@afterima.ge>

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 }) => {
+          onFinish: async ({ usage }) => {
-            finalUsage = usage;
+            try {
-            dataStream.write({ type: 'data-usage', data: usage });
+              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);

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>(
+  const [usage, setUsage] = useState<AppUsage | undefined>(initialLastContext);
-    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') {
+      if (dataPart.type === 'data-usage') setUsage(dataPart.data);
-        setUsage(dataPart.data);
-      }
     },
     onFinish: () => {
       mutate(unstable_serialize(getChatHistoryPaginationKey));

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;
+  usage?: AppUsage;
-  /** Optional model id (canonical or alias) to compute cost */
-  modelId?: string;
 };
 
 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 = ({
+export const Context = ({ className, usage, ...props }: ContextProps) => {
-  className,
+  const used = usage?.totalTokens ?? 0;
-  maxTokens,
+  const max =
-  usedTokens,
+    usage?.context?.totalMax ??
-  usage,
+    usage?.context?.combinedMax ??
-  modelId,
+    usage?.context?.inputMax;
-  ...props
+  const hasMax = typeof max === 'number' && Number.isFinite(max) && max > 0;
-}: ContextProps) => {
+  const usedPercent = hasMax ? Math.min(100, (used / max) * 100) : 0;
-  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);
-
   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">
+          <div className="flex items-start justify-between text-sm">
-            <span>{displayPct}</span>
+            <span>{usedPercent.toFixed(1)}%</span>
-            <span className="text-muted-foreground">{used} / {total} tokens</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}
+                tokens={usage?.cachedInputTokens}
-                costText={cacheReadsCostText}
+                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}
+              tokens={usage?.inputTokens}
-              costText={inputCostText}
+              costText={usage?.costUSD?.inputUSD?.toString()}
             />
             <InfoRow
               label="Output"
-              tokens={displayOutput}
+              tokens={usage?.outputTokens}
-              costText={outputCostText}
+              costText={usage?.costUSD?.outputUSD?.toString()}
             />
             <InfoRow
               label="Reasoning"
-              tokens={reasoningTokens > 0 ? reasoningTokens : undefined}
+              tokens={
-              costText={reasoningCostText}
+                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>
               </>
             )}
@@ -343,4 +184,4 @@ export const Context = ({
       </DropdownMenuContent>
     </DropdownMenu>
   );
-};
+};

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) => (
+          {chatModels.map((model) => (
-          <SelectItem key={model.id} value={model.name} className="px-3 py-2 text-xs">
+            <SelectItem
-            <div className="flex flex-col flex-1 gap-1 min-w-0">
+              key={model.id}
-              <div className="text-xs font-medium truncate">
+              value={model.name}
-                {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">
+            </SelectItem>
-                {model.description}
+          ))}
-              </div>
-            </div>
-          </SelectItem>
-        ))}
         </div>
       </PromptInputModelSelectContent>
     </PromptInputModelSelect>
@@ -504,4 +496,4 @@ function PureStopButton({
   );
 }
 
-const StopButton = memo(PureStopButton);
+const StopButton = memo(PureStopButton);

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
+  // Store merged server-enriched usage object
-  context: LanguageModelV2Usage;
+  context: AppUsage;
 }) {
   try {
     return await db

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>;

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<

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 };

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"

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
+        specifier: 1.3.0-canary.5
-        version: 1.1.2
+        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:
+  tokenlens@1.3.0-canary.5:
-    resolution: {integrity: sha512-PICXahAiVS/2nyc30SijuYZj5wIm2G9+2U7fdvQbl5pa/hIUpQ8H+RdYOzOrSvva6GMqPM5+6cWd9Hqi81Yv3g==}
+    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: {}