test(ai): address message ID review feedback

Author: mikemikimike

packages/ai/tests/message-converters.test.ts

@@ -403,6 +403,18 @@ describe('Message Converters', () => {
       ])
     })
 
+    it('should preserve the UI message id on an empty assistant fallback', () => {
+      const uiMessage: UIMessage = {
+        id: 'assistant-empty',
+        role: 'assistant',
+        parts: [],
+      }
+
+      expect(uiMessageToModelMessages(uiMessage)).toEqual([
+        { id: uiMessage.id, role: 'assistant', content: null },
+      ])
+    })
+
     it('should preserve interleaving of text, tool calls, and tool results', () => {
       const uiMessage: UIMessage = {
         id: 'msg-1',

testing/e2e/src/routes/api.message-ids.ts

@@ -1,6 +1,110 @@
 import { createFileRoute } from '@tanstack/react-router'
 import { convertMessagesToModelMessages } from '@tanstack/ai'
-import type { ModelMessage, UIMessage } from '@tanstack/ai'
+import { z } from 'zod'
+import type { UIMessage } from '@tanstack/ai'
+
+const sourceSchema = z.discriminatedUnion('type', [
+  z.object({
+    type: z.literal('data'),
+    value: z.string(),
+    mimeType: z.string(),
+  }),
+  z.object({
+    type: z.literal('url'),
+    value: z.string(),
+    mimeType: z.string().optional(),
+  }),
+])
+
+const contentPartSchema = z.discriminatedUnion('type', [
+  z.object({
+    type: z.literal('text'),
+    content: z.string(),
+    metadata: z.unknown().optional(),
+  }),
+  ...(['image', 'audio', 'video', 'document'] as const).map((type) =>
+    z.object({
+      type: z.literal(type),
+      source: sourceSchema,
+      metadata: z.unknown().optional(),
+    }),
+  ),
+])
+
+const messagePartSchema = z.discriminatedUnion('type', [
+  ...contentPartSchema.options,
+  z.object({
+    type: z.literal('tool-call'),
+    id: z.string(),
+    name: z.string(),
+    arguments: z.string(),
+    input: z.unknown().optional(),
+    state: z.enum([
+      'awaiting-input',
+      'input-streaming',
+      'input-complete',
+      'approval-requested',
+      'approval-responded',
+      'complete',
+      'error',
+    ]),
+    approval: z
+      .object({
+        id: z.string(),
+        needsApproval: z.boolean(),
+        approved: z.boolean().optional(),
+      })
+      .optional(),
+    output: z.unknown().optional(),
+    metadata: z.unknown().optional(),
+  }),
+  z.object({
+    type: z.literal('tool-result'),
+    toolCallId: z.string(),
+    content: z.union([z.string(), z.array(contentPartSchema)]),
+    state: z.enum(['streaming', 'complete', 'error']),
+    error: z.string().optional(),
+  }),
+  z.object({
+    type: z.literal('thinking'),
+    content: z.string(),
+    stepId: z.string().optional(),
+    signature: z.string().optional(),
+  }),
+  z.object({
+    type: z.literal('structured-output'),
+    status: z.enum(['streaming', 'complete', 'error']),
+    partial: z.unknown().optional(),
+    data: z.unknown().optional(),
+    raw: z.string(),
+    reasoning: z.string().optional(),
+    errorMessage: z.string().optional(),
+  }),
+  z.object({
+    type: z.literal('ui-resource'),
+    resource: z.object({
+      uri: z.string(),
+      mimeType: z.string(),
+      text: z.string().optional(),
+      blob: z.string().optional(),
+    }),
+    serverId: z.string().optional(),
+    toolCallId: z.string(),
+    toolName: z.string(),
+    meta: z.record(z.string(), z.unknown()).optional(),
+  }),
+])
+
+const requestBodySchema: z.ZodType<{ messages: Array<UIMessage> }> = z.object({
+  messages: z.array(
+    z.object({
+      id: z.string(),
+      role: z.enum(['system', 'user', 'assistant']),
+      parts: z.array(messagePartSchema),
+      createdAt: z.coerce.date().optional(),
+    }),
+  ),
+})
 
 /**
  * Provider-free harness for the UIMessage -> ModelMessage identity contract.
@@ -11,11 +115,21 @@ export const Route = createFileRoute('/api/message-ids')({
   server: {
     handlers: {
       POST: async ({ request }) => {
-        const body = (await request.json()) as {
-          messages: Array<UIMessage | ModelMessage>
+        let body: unknown
+        try {
+          body = await request.json()
+        } catch {
+          return new Response('Invalid JSON request body', { status: 400 })
+        }
+
+        const parsed = requestBodySchema.safeParse(body)
+        if (!parsed.success) {
+          return new Response('Invalid message data', { status: 400 })
         }
 
-        return Response.json(convertMessagesToModelMessages(body.messages))
+        return Response.json(
+          convertMessagesToModelMessages(parsed.data.messages),
+        )
       },
     },
   },

testing/e2e/tests/chat.spec.ts

@@ -7,6 +7,9 @@ import {
 } from './helpers'
 import { providersFor } from './test-matrix'
 
+// The server conversion test in this spec does not call a provider HTTP
+// endpoint, so it intentionally does not configure aimock.
+
 for (const provider of providersFor('chat')) {
   test.describe(`${provider} — chat`, () => {
     test('sends a message and receives a streaming response', async ({
@@ -114,6 +117,35 @@ test('preserves UI message IDs at the server conversion boundary', async ({
   ])
 })
 
+test('rejects malformed JSON at the server conversion boundary', async ({
+  request,
+}) => {
+  const response = await request.post('/api/message-ids', {
+    data: '{',
+    headers: { 'Content-Type': 'application/json' },
+  })
+
+  expect(response.status()).toBe(400)
+})
+
+test('rejects invalid message parts at the server conversion boundary', async ({
+  request,
+}) => {
+  const response = await request.post('/api/message-ids', {
+    data: {
+      messages: [
+        {
+          id: 'user-1',
+          role: 'user',
+          parts: [{ type: 'text', content: 42 }],
+        },
+      ],
+    },
+  })
+
+  expect(response.status()).toBe(400)
+})
+
 test.describe('openai chat persistence', () => {
   test('persists chat messages across browser reload with localStorage', async ({
     page,