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

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

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

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