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

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

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>
@@ -504,4 +496,4 @@ function PureStopButton({
   );
 }
 
-const StopButton = memo(PureStopButton);
+const StopButton = memo(PureStopButton);