ai 6.0.38 → 6.0.40

package/src/generate-text/stream-text.ts

@@ -2229,6 +2229,9 @@ However, the LLM results are expected to be small enough to not cause issues.
                 ...(part.providerExecuted != null
                   ? { providerExecuted: part.providerExecuted }
                   : {}),
+                ...(part.providerMetadata != null
+                  ? { providerMetadata: part.providerMetadata }
+                  : {}),
                 ...(dynamic != null ? { dynamic } : {}),
                 ...(part.title != null ? { title: part.title } : {}),
               });
@@ -2371,6 +2374,7 @@ However, the LLM results are expected to be small enough to not cause issues.
                 controller.enqueue({
                   type: 'finish',
                   finishReason: part.finishReason,
+                  usage: part.totalUsage,
                   ...(messageMetadataValue != null
                     ? { messageMetadata: messageMetadataValue }
                     : {}),

package/src/types/usage.ts

@@ -3,6 +3,8 @@ import {
   JSONObject,
   LanguageModelV3Usage,
 } from '@ai-sdk/provider';
+import { z } from 'zod/v4';
+import { jsonValueSchema } from './json-value';
 
 /**
  * Represents the number of tokens used in a prompt and completion.
@@ -77,6 +79,28 @@ export type LanguageModelUsage = {
   raw?: JSONObject;
 };
 
+// Helper for required properties that can be undefined (key: Type | undefined)
+const numberOrUndefined = z.union([z.number(), z.undefined()]);
+
+export const languageModelUsageSchema: z.ZodType<LanguageModelUsage> =
+  z.strictObject({
+    inputTokens: numberOrUndefined,
+    inputTokenDetails: z.strictObject({
+      noCacheTokens: numberOrUndefined,
+      cacheReadTokens: numberOrUndefined,
+      cacheWriteTokens: numberOrUndefined,
+    }),
+    outputTokens: numberOrUndefined,
+    outputTokenDetails: z.strictObject({
+      textTokens: numberOrUndefined,
+      reasoningTokens: numberOrUndefined,
+    }),
+    totalTokens: numberOrUndefined,
+    reasoningTokens: z.number().optional(),
+    cachedInputTokens: z.number().optional(),
+    raw: z.record(z.string(), jsonValueSchema.optional()).optional(),
+  });
+
 /**
 Represents the number of tokens used in an embedding.
  */

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.test.ts

@@ -6171,6 +6171,85 @@ describe('processUIMessageStream', () => {
     });
   });
 
+  describe('provider metadata during input-streaming', () => {
+    beforeEach(async () => {
+      const stream = createUIMessageStream([
+        { type: 'start', messageId: 'msg-123' },
+        { type: 'start-step' },
+        {
+          type: 'tool-input-start',
+          toolCallId: 'tool-call-id',
+          toolName: 'tool-name',
+          providerMetadata: { testProvider: { someKey: 'someValue' } },
+        },
+        {
+          type: 'tool-input-delta',
+          toolCallId: 'tool-call-id',
+          inputTextDelta: '{"query":',
+        },
+        {
+          type: 'tool-input-delta',
+          toolCallId: 'tool-call-id',
+          inputTextDelta: '"test"}',
+        },
+        {
+          type: 'tool-input-available',
+          toolCallId: 'tool-call-id',
+          toolName: 'tool-name',
+          input: { query: 'test' },
+        },
+        { type: 'finish-step' },
+        { type: 'finish' },
+      ]);
+
+      state = createStreamingUIMessageState({
+        messageId: 'msg-123',
+        lastMessage: undefined,
+      });
+
+      await consumeStream({
+        stream: processUIMessageStream({
+          stream,
+          runUpdateMessageJob,
+          onError: error => {
+            throw error;
+          },
+        }),
+      });
+    });
+
+    it('should propagate providerMetadata during input-streaming state', async () => {
+      // Find the first update after tool-input-start (when state is input-streaming)
+      const inputStreamingUpdate = writeCalls.find(call =>
+        call.message.parts.some(
+          (p: any) =>
+            p.toolCallId === 'tool-call-id' && p.state === 'input-streaming',
+        ),
+      );
+
+      expect(inputStreamingUpdate).toBeDefined();
+      const toolPart = inputStreamingUpdate!.message.parts.find(
+        (p: any) => p.toolCallId === 'tool-call-id',
+      ) as any;
+
+      // Key assertion: providerMetadata should be available during input-streaming
+      expect(toolPart.callProviderMetadata).toEqual({
+        testProvider: { someKey: 'someValue' },
+      });
+    });
+
+    it('should retain providerMetadata in final input-available state', async () => {
+      const toolPart = state!.message.parts.find(
+        (p: any) => p.toolCallId === 'tool-call-id',
+      ) as any;
+
+      expect(toolPart.state).toBe('input-available');
+      expect(toolPart.callProviderMetadata).toEqual({
+        testProvider: { someKey: 'someValue' },
+      });
+    });
+  });
+
   describe('tool input error', () => {
     beforeEach(async () => {
       const stream = createUIMessageStream([

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>({
@@ -130,6 +132,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   state: 'input-streaming';
                   input: unknown;
                   providerExecuted?: boolean;
+                  providerMetadata?: ProviderMetadata;
                 }
               | {
                   state: 'input-available';
@@ -177,10 +180,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               anyPart.providerExecuted =
                 anyOptions.providerExecuted ?? part.providerExecuted;
 
-              if (
-                anyOptions.providerMetadata != null &&
-                part.state === 'input-available'
-              ) {
+              if (anyOptions.providerMetadata != null) {
                 part.callProviderMetadata = anyOptions.providerMetadata;
               }
             } else {
@@ -212,6 +212,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               | {
                   state: 'input-streaming';
                   input: unknown;
+                  providerMetadata?: ProviderMetadata;
                 }
               | {
                   state: 'input-available';
@@ -256,10 +257,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               anyPart.providerExecuted =
                 anyOptions.providerExecuted ?? part.providerExecuted;
 
-              if (
-                anyOptions.providerMetadata != null &&
-                part.state === 'input-available'
-              ) {
+              if (anyOptions.providerMetadata != null) {
                 part.callProviderMetadata = anyOptions.providerMetadata;
               }
             } else {
@@ -461,6 +459,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   input: undefined,
                   providerExecuted: chunk.providerExecuted,
                   title: chunk.title,
+                  providerMetadata: chunk.providerMetadata,
                 });
               } else {
                 updateToolPart({
@@ -470,6 +469,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   input: undefined,
                   providerExecuted: chunk.providerExecuted,
                   title: chunk.title,
+                  providerMetadata: chunk.providerMetadata,
                 });
               }
 
@@ -688,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/ui-messages.ts

@@ -233,6 +233,7 @@ export type UIToolInvocation<TOOL extends UITool | Tool> = {
       input: DeepPartial<asUITool<TOOL>['input']> | undefined;
       output?: never;
       errorText?: never;
+      callProviderMetadata?: ProviderMetadata;
       approval?: never;
     }
   | {
@@ -337,6 +338,7 @@ export type DynamicToolUIPart = {
       input: unknown | undefined;
       output?: never;
       errorText?: never;
+      callProviderMetadata?: ProviderMetadata;
       approval?: never;
     }
   | {

package/src/ui/validate-ui-messages.ts

@@ -78,6 +78,7 @@ const uiMessagesSchema = lazySchema(() =>
                   state: z.literal('input-streaming'),
                   input: z.unknown().optional(),
                   providerExecuted: z.boolean().optional(),
+                  callProviderMetadata: providerMetadataSchema.optional(),
                   output: z.never().optional(),
                   errorText: z.never().optional(),
                   approval: z.never().optional(),
@@ -185,6 +186,7 @@ const uiMessagesSchema = lazySchema(() =>
                   toolCallId: z.string(),
                   state: z.literal('input-streaming'),
                   providerExecuted: z.boolean().optional(),
+                  callProviderMetadata: providerMetadataSchema.optional(),
                   input: z.unknown().optional(),
                   output: z.never().optional(),
                   errorText: z.never().optional(),

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,
@@ -41,6 +42,7 @@ export const uiMessageChunkSchema = lazySchema(() =>
         toolCallId: z.string(),
         toolName: z.string(),
         providerExecuted: z.boolean().optional(),
+        providerMetadata: providerMetadataSchema.optional(),
         dynamic: z.boolean().optional(),
         title: z.string().optional(),
       }),
@@ -164,6 +166,7 @@ export const uiMessageChunkSchema = lazySchema(() =>
             'other',
           ] as const satisfies readonly FinishReason[])
           .optional(),
+        usage: languageModelUsageSchema.optional(),
         messageMetadata: z.unknown().optional(),
       }),
       z.strictObject({
@@ -277,6 +280,7 @@ export type UIMessageChunk<
       toolCallId: string;
       toolName: string;
       providerExecuted?: boolean;
+      providerMetadata?: ProviderMetadata;
       dynamic?: boolean;
       title?: string;
     }
@@ -321,6 +325,7 @@ export type UIMessageChunk<
   | {
       type: 'finish';
       finishReason?: FinishReason;
+      usage?: LanguageModelUsage;
       messageMetadata?: METADATA;
     }
   | {