ai 6.0.39 → 6.0.40

package/src/ui/chat.ts

@@ -5,6 +5,7 @@ import {
   InferSchema,
 } from '@ai-sdk/provider-utils';
 import { FinishReason } from '../types/language-model';
+import { LanguageModelUsage } from '../types/usage';
 import { UIMessageChunk } from '../ui-message-stream/ui-message-chunks';
 import { consumeStream } from '../util/consume-stream';
 import { SerialJobExecutor } from '../util/serial-job-executor';
@@ -124,6 +125,7 @@ export type ChatOnDataCallback<UI_MESSAGE extends UIMessage> = (
  * @param isDisconnect Indicates whether the request has been ended by a network error.
  * @param isError Indicates whether the request has been ended by an error.
  * @param finishReason The reason why the generation finished.
+ * @param usage Token usage information for the response.
  */
 export type ChatOnFinishCallback<UI_MESSAGE extends UIMessage> = (options: {
   message: UI_MESSAGE;
@@ -132,6 +134,7 @@ export type ChatOnFinishCallback<UI_MESSAGE extends UIMessage> = (options: {
   isDisconnect: boolean;
   isError: boolean;
   finishReason?: FinishReason;
+  usage?: LanguageModelUsage;
 }) => void;
 
 export interface ChatInit<UI_MESSAGE extends UIMessage> {
@@ -691,6 +694,9 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
           isDisconnect,
           isError,
           finishReason: this.activeResponse?.state.finishReason,
+          ...(this.activeResponse?.state.usage != null && {
+            usage: this.activeResponse.state.usage,
+          }),
         });
       } catch (err) {
         console.error(err);

package/src/ui/process-ui-message-stream.ts

@@ -2,6 +2,7 @@ import { FlexibleSchema, validateTypes } from '@ai-sdk/provider-utils';
 import { UIMessageStreamError } from '../error/ui-message-stream-error';
 import { ProviderMetadata } from '../types';
 import { FinishReason } from '../types/language-model';
+import { LanguageModelUsage } from '../types/usage';
 import {
   DataUIMessageChunk,
   InferUIMessageChunk,
@@ -44,6 +45,7 @@ export type StreamingUIMessageState<UI_MESSAGE extends UIMessage> = {
     }
   >;
   finishReason?: FinishReason;
+  usage?: LanguageModelUsage;
 };
 
 export function createStreamingUIMessageState<UI_MESSAGE extends UIMessage>({
@@ -686,6 +688,9 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               if (chunk.finishReason != null) {
                 state.finishReason = chunk.finishReason;
               }
+              if (chunk.usage != null) {
+                state.usage = chunk.usage;
+              }
               await updateMessageMetadata(chunk.messageMetadata);
               if (chunk.messageMetadata != null) {
                 write();

package/src/ui-message-stream/ui-message-chunks.ts

@@ -4,6 +4,7 @@ import {
   providerMetadataSchema,
 } from '../types/provider-metadata';
 import { FinishReason } from '../types/language-model';
+import { LanguageModelUsage, languageModelUsageSchema } from '../types/usage';
 import {
   InferUIMessageData,
   InferUIMessageMetadata,
@@ -165,6 +166,7 @@ export const uiMessageChunkSchema = lazySchema(() =>
             'other',
           ] as const satisfies readonly FinishReason[])
           .optional(),
+        usage: languageModelUsageSchema.optional(),
         messageMetadata: z.unknown().optional(),
       }),
       z.strictObject({
@@ -323,6 +325,7 @@ export type UIMessageChunk<
   | {
       type: 'finish';
       finishReason?: FinishReason;
+      usage?: LanguageModelUsage;
       messageMetadata?: METADATA;
     }
   | {