Skip to content

chore: Context in message metadata #1181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

chore: Context in message metadata #1181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 13 additions & 19 deletions app/(chat)/api/chat/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ import {
convertToModelMessages,
createUIMessageStream,
JsonToSseTransformStream,
type LanguageModelUsage,
smoothStream,
stepCountIs,
streamText,
Expand All @@ -18,7 +17,6 @@ import {
saveChat,
saveMessages,
} from '@/lib/db/queries';
import { updateChatLastContextById } from '@/lib/db/queries';
import { convertToUIMessages, generateUUID } from '@/lib/utils';
import { generateTitleFromUserMessage } from '../../actions';
import { createDocument } from '@/lib/ai/tools/create-document';
Expand Down Expand Up @@ -143,6 +141,7 @@ export async function POST(request: Request) {
role: 'user',
parts: message.parts,
attachments: [],
lastContext: null,
createdAt: new Date(),
},
],
Expand All @@ -151,8 +150,6 @@ export async function POST(request: Request) {
const streamId = generateUUID();
await createStreamId({ streamId, chatId: id });

let finalUsage: LanguageModelUsage | undefined;

const stream = createUIMessageStream({
execute: ({ writer: dataStream }) => {
const result = streamText({
Expand Down Expand Up @@ -183,16 +180,22 @@ export async function POST(request: Request) {
isEnabled: isProductionEnvironment,
functionId: 'stream-text',
},
onFinish: ({ usage }) => {
finalUsage = usage;
dataStream.write({ type: 'data-usage', data: usage });
},
});

result.consumeStream();

dataStream.merge(
result.toUIMessageStream({
messageMetadata: ({ part }) => {
// send custom information to the client on start:
// when the message is finished, send additional information:
if (part.type === 'finish') {
return {
createdAt: new Date().toISOString(),
usage: part.totalUsage,
};
}
},
sendReasoning: true,
}),
);
Expand All @@ -207,20 +210,11 @@ export async function POST(request: Request) {
createdAt: new Date(),
attachments: [],
chatId: id,
lastContext: null,
Copy link

@vercel vercel bot Sep 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The lastContext field is hardcoded to null when saving messages, but the client receives usage data through messageMetadata, creating a disconnect between stored and displayed data.

View Details
📝 Patch Details 
diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
index 05df486..fad9af8 100644
--- a/app/(chat)/api/chat/route.ts
+++ b/app/(chat)/api/chat/route.ts
@@ -150,6 +150,9 @@ export async function POST(request: Request) {
     const streamId = generateUUID();
     await createStreamId({ streamId, chatId: id });
 
+    // Store usage data to be captured from streamText
+    let totalUsage: any = null;
+
     const stream = createUIMessageStream({
       execute: ({ writer: dataStream }) => {
         const result = streamText({
@@ -180,6 +183,10 @@ export async function POST(request: Request) {
             isEnabled: isProductionEnvironment,
             functionId: 'stream-text',
           },
+          onFinish: ({ totalUsage: usage }) => {
+            // Capture usage data from streamText for later storage
+            totalUsage = usage;
+          },
         });
 
         result.consumeStream();
@@ -210,7 +217,7 @@ export async function POST(request: Request) {
             createdAt: new Date(),
             attachments: [],
             chatId: id,
-            lastContext: null,
+            lastContext: message.role === 'assistant' ? totalUsage : null,
           })),
         });
       },

Analysis

Usage Data Loss Bug in Chat API

Summary

The chat API is losing usage data (token counts) when messages are persisted to the database, creating a disconnect between what users see during active sessions and what's available after page refreshes.

Technical Details

The Problem

In app/(chat)/api/chat/route.ts, the onFinish callback of createUIMessageStream hardcodes lastContext: null when saving messages to the database (line 213), despite having access to usage data from the underlying streamText operation.

onFinish: async ({ messages }) => {
  await saveMessages({
    messages: messages.map((message) => ({
      id: message.id,
      role: message.role,
      parts: message.parts,
      createdAt: new Date(),
      attachments: [],
      chatId: id,
      lastContext: null, // ❌ Hardcoded null loses usage data
    })),
  });
},

Data Flow Analysis

  1. During streaming: The messageMetadata function (lines 189-198) correctly sends usage data to clients via part.totalUsage
  2. During persistence: lastContext is hardcoded to null, losing usage information
  3. On reload: convertToUIMessages maps message.lastContext to metadata.usage, resulting in undefined instead of actual usage data

Database Schema Evidence

The database schema confirms lastContext is designed to store usage data:

lastContext: jsonb('lastContext').$type<LanguageModelV2Usage | null>()

API Documentation Verification

According to the AI SDK documentation, streamText's onFinish callback provides a totalUsage parameter containing token usage information. The current implementation fails to capture this data.

Impact

  • User Experience: Usage information disappears after page refresh/reload
  • Analytics: Lost token usage data prevents accurate billing and usage tracking
  • Debugging: Difficulty troubleshooting expensive API calls without persistent usage data

Solution Implemented

The fix captures usage data from streamText's onFinish callback and stores it in the database:

  1. Added a variable to capture totalUsage from the streamText result
  2. Used onFinish callback on streamText to capture the usage data
  3. Modified the message persistence to store usage data for assistant messages
// Capture usage data from streamText
onFinish: ({ totalUsage: usage }) => {
  totalUsage = usage;
},

// Store usage data for assistant messages
lastContext: message.role === 'assistant' ? totalUsage : null,

Verification

  • ✅ TypeScript compilation passes
  • ✅ Next.js build succeeds
  • ✅ Code follows existing patterns in the codebase
  • ✅ Database schema supports the fix
  • ✅ API documentation confirms usage data availability

This fix ensures usage data persists across sessions while maintaining backward compatibility with existing null values for user messages.

})),
});

if (finalUsage) {
try {
await updateChatLastContextById({
chatId: id,
context: finalUsage,
});
} catch (err) {
console.warn('Unable to persist last usage for chat', id, err);
}
}
},

onError: () => {
return 'Oops, an error occurred!';
},
Expand Down
2 changes: 0 additions & 2 deletions app/(chat)/chat/[id]/page.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,6 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
isReadonly={session?.user?.id !== chat.userId}
session={session}
autoResume={true}
initialLastContext={chat.lastContext ?? undefined}
/>
<DataStreamHandler />
</>
Expand All @@ -70,7 +69,6 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
isReadonly={session?.user?.id !== chat.userId}
session={session}
autoResume={true}
initialLastContext={chat.lastContext ?? undefined}
/>
<DataStreamHandler />
</>
Expand Down
7 changes: 7 additions & 0 deletions components/artifact.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import {
type SetStateAction,
useCallback,
useEffect,
useMemo,
useState,
} from 'react';
import useSWR, { useSWRConfig } from 'swr';
Expand Down Expand Up @@ -254,6 +255,11 @@ function PureArtifact({
}
}, [artifact.documentId, artifactDefinition, setMetadata]);

const lastAssistantMessage = useMemo(
() => messages.findLast((message) => message.role === 'assistant'),
[messages],
);

return (
<AnimatePresence>
{artifact.isVisible && (
Expand Down Expand Up @@ -339,6 +345,7 @@ function PureArtifact({
setMessages={setMessages}
selectedVisibilityType={selectedVisibilityType}
selectedModelId={selectedModelId}
usage={lastAssistantMessage?.metadata?.usage}
/>
</div>
</div>
Expand Down
19 changes: 8 additions & 11 deletions components/chat.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
'use client';

import { DefaultChatTransport, type LanguageModelUsage } from 'ai';
import { DefaultChatTransport } from 'ai';
import { useChat } from '@ai-sdk/react';
import { useEffect, useState } from 'react';
import { useEffect, useMemo, useState } from 'react';
import useSWR, { useSWRConfig } from 'swr';
import { ChatHeader } from '@/components/chat-header';
import type { Vote } from '@/lib/db/schema';
Expand Down Expand Up @@ -31,7 +31,6 @@ export function Chat({
isReadonly,
session,
autoResume,
initialLastContext,
}: {
id: string;
initialMessages: ChatMessage[];
Expand All @@ -40,7 +39,6 @@ export function Chat({
isReadonly: boolean;
session: Session;
autoResume: boolean;
initialLastContext?: LanguageModelUsage;
}) {
const { visibilityType } = useChatVisibility({
chatId: id,
Expand All @@ -51,9 +49,6 @@ export function Chat({
const { setDataStream } = useDataStream();

const [input, setInput] = useState<string>('');
const [usage, setUsage] = useState<LanguageModelUsage | undefined>(
initialLastContext,
);

const {
messages,
Expand Down Expand Up @@ -85,9 +80,6 @@ export function Chat({
}),
onData: (dataPart) => {
setDataStream((ds) => (ds ? [...ds, dataPart] : []));
if (dataPart.type === 'data-usage') {
setUsage(dataPart.data);
}
},
onFinish: () => {
mutate(unstable_serialize(getChatHistoryPaginationKey));
Expand Down Expand Up @@ -134,6 +126,11 @@ export function Chat({
setMessages,
});

const lastAssistantMessage = useMemo(
() => messages.findLast((message) => message.role === 'assistant'),
[messages],
);

return (
<>
<div className="overscroll-behavior-contain flex h-dvh min-w-0 touch-pan-y flex-col bg-background">
Expand Down Expand Up @@ -171,7 +168,7 @@ export function Chat({
sendMessage={sendMessage}
selectedVisibilityType={visibilityType}
selectedModelId={initialChatModel}
usage={usage}
usage={lastAssistantMessage?.metadata?.usage}
/>
)}
</div>
Expand Down
2 changes: 2 additions & 0 deletions lib/db/migrations/0008_smart_king_cobra.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
ALTER TABLE "Message_v2" ADD COLUMN "lastContext" jsonb;--> statement-breakpoint
ALTER TABLE "Chat" DROP COLUMN IF EXISTS "lastContext";
Loading
Loading