feat(agent): add message metadata support when inferring UI messages

.changeset/red-wasps-know.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat(agent): add message metadata support when inferring UI messages

content/docs/07-reference/01-ai-sdk-core/16-tool-loop-agent.mdx

@@ -285,6 +285,8 @@ The `stream()` method returns a `StreamTextResult` object (see [`streamText`](/d
 
 Infers the UI message type for the given agent instance. Useful for type-safe UI and message exchanges.
 
+#### Basic Example
+
 ```ts
 import { ToolLoopAgent, InferAgentUIMessage } from 'ai';
 
@@ -296,6 +298,36 @@ const weatherAgent = new ToolLoopAgent({
 type WeatherAgentUIMessage = InferAgentUIMessage<typeof weatherAgent>;
 ```
 
+#### Example with Message Metadata
+
+You can provide a second type argument to customize the metadata for each message. This is useful for tracking rich metadata returned by the agent (such as createdAt, tokens, finish reason, etc.).
+
+```ts
+import { ToolLoopAgent, InferAgentUIMessage } from 'ai';
+import { z } from 'zod';
+
+// Example schema for message metadata
+const exampleMetadataSchema = z.object({
+  createdAt: z.number().optional(),
+  model: z.string().optional(),
+  totalTokens: z.number().optional(),
+  finishReason: z.string().optional(),
+});
+type ExampleMetadata = z.infer<typeof exampleMetadataSchema>;
+
+// Define agent as usual
+const metadataAgent = new ToolLoopAgent({
+  model: 'openai/gpt-4o',
+  // ...other options
+});
+
+// Type-safe UI message type with custom metadata
+type MetadataAgentUIMessage = InferAgentUIMessage<
+  typeof metadataAgent,
+  ExampleMetadata
+>;
+```
+
 ## Examples
 
 ### Basic Agent with Tools

examples/next-openai/app/api/use-chat-message-metadata/example-metadata-schema.ts → examples/next-openai/agent/openai-metadata-agent.ts

@@ -1,3 +1,6 @@
+import { openai } from '@ai-sdk/openai';
+import { ToolLoopAgent, InferAgentUIMessage } from 'ai';
+
 import { z } from 'zod';
 
 export const exampleMetadataSchema = z.object({
@@ -9,3 +12,12 @@ export const exampleMetadataSchema = z.object({
 });
 
 export type ExampleMetadata = z.infer<typeof exampleMetadataSchema>;
+
+export const openaiMetadataAgent = new ToolLoopAgent({
+  model: openai('gpt-4o'),
+});
+
+export type OpenAIMetadataMessage = InferAgentUIMessage<
+  typeof openaiMetadataAgent,
+  ExampleMetadata
+>;

examples/next-openai/app/api/use-chat-message-metadata/route.ts

@@ -1,16 +1,15 @@
-import { openai } from '@ai-sdk/openai';
-import { convertToModelMessages, streamText, UIMessage } from 'ai';
-import { ExampleMetadata } from './example-metadata-schema';
+import { createAgentUIStreamResponse, UIMessage } from 'ai';
+import {
+  ExampleMetadata,
+  openaiMetadataAgent,
+} from '@/agent/openai-metadata-agent';
 
 export async function POST(req: Request) {
   const { messages }: { messages: UIMessage[] } = await req.json();
 
-  const result = streamText({
-    model: openai('gpt-4o'),
-    prompt: convertToModelMessages(messages),
-  });
-
-  return result.toUIMessageStreamResponse({
+  return createAgentUIStreamResponse({
+    agent: openaiMetadataAgent,
+    messages,
     messageMetadata: ({ part }): ExampleMetadata | undefined => {
       // send custom information to the client on start:
       if (part.type === 'start') {

examples/next-openai/app/use-chat-message-metadata/page.tsx

@@ -3,7 +3,7 @@
 import ChatInput from '@/components/chat-input';
 import { useChat } from '@ai-sdk/react';
 import { DefaultChatTransport, UIMessage } from 'ai';
-import { ExampleMetadata } from '../api/use-chat-message-metadata/example-metadata-schema';
+import { ExampleMetadata } from '@/agent/openai-metadata-agent';
 
 type MyMessage = UIMessage<ExampleMetadata>;
 
@@ -16,7 +16,7 @@ export default function Chat() {
     });
 
   return (
-    <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
+    <div className="flex flex-col py-24 mx-auto w-full max-w-md stretch">
       {messages.map(message => (
         <div key={message.id} className="whitespace-pre-wrap">
           {message.role === 'user' ? 'User: ' : 'AI: '}
@@ -46,7 +46,7 @@ export default function Chat() {
           {status === 'submitted' && <div>Loading...</div>}
           <button
             type="button"
-            className="px-4 py-2 mt-4 text-blue-500 border border-blue-500 rounded-md"
+            className="px-4 py-2 mt-4 text-blue-500 rounded-md border border-blue-500"
             onClick={stop}
           >
             Stop
@@ -59,7 +59,7 @@ export default function Chat() {
           <div className="text-red-500">An error occurred.</div>
           <button
             type="button"
-            className="px-4 py-2 mt-4 text-blue-500 border border-blue-500 rounded-md"
+            className="px-4 py-2 mt-4 text-blue-500 rounded-md border border-blue-500"
             onClick={() => regenerate()}
           >
             Retry

packages/ai/src/agent/infer-agent-ui-message.test-d.ts

@@ -15,14 +15,14 @@ import { InferAgentUIMessage } from './infer-agent-ui-message';
 
 describe('InferAgentUIMessage', () => {
   it('should not contain arbitrary static tools when no tools are provided', () => {
-    const baseAgent = new ToolLoopAgent({
+    const agent = new ToolLoopAgent({
       model: 'openai/gpt-4o',
       // no tools
     });
 
-    type Message = InferAgentUIMessage<typeof baseAgent>;
+    type Message = InferAgentUIMessage<typeof agent>;
 
-    expectTypeOf<Message>().toMatchTypeOf<UIMessage<never, never, {}>>();
+    expectTypeOf<Message>().toMatchTypeOf<UIMessage<unknown, never, {}>>();
 
     type MessagePart = Message['parts'][number];
 
@@ -38,4 +38,17 @@ describe('InferAgentUIMessage', () => {
       | StepStartUIPart
     >();
   });
+
+  it('should include metadata when provided', () => {
+    const agent = new ToolLoopAgent({
+      model: 'openai/gpt-4o',
+      // no tools
+    });
+
+    type Message = InferAgentUIMessage<typeof agent, { foo: string }>;
+
+    expectTypeOf<Message>().toMatchTypeOf<
+      UIMessage<{ foo: string }, never, {}>
+    >();
+  });
 });

packages/ai/src/agent/infer-agent-ui-message.ts

@@ -4,8 +4,8 @@ import { InferAgentTools } from './infer-agent-tools';
 /**
  * Infer the UI message type of an agent.
  */
-export type InferAgentUIMessage<AGENT> = UIMessage<
-  never,
+export type InferAgentUIMessage<AGENT, MESSAGE_METADATA = unknown> = UIMessage<
+  MESSAGE_METADATA,
   never,
   InferUITools<InferAgentTools<AGENT>>
 >;