fix: align canvas media batch ids

app/(chat)/api/chat/_lib/stream.ts

@@ -73,6 +73,98 @@ function normalizeStreamIdList(value: unknown) {
     : [];
 }
 
+function normalizeStreamMediaUrls(value: unknown) {
+  return Array.isArray(value)
+    ? value.map((url) => String(url || "").trim())
+    : [];
+}
+
+function isValidStreamMediaUrl(url: string) {
+  return /^https?:\/\//i.test(url);
+}
+
+export function normalizeCanvasMediaStreamPayload(params: {
+  imageUrls?: unknown;
+  videoUrls?: unknown;
+  mediaItemIds?: unknown;
+  handledCanvasMediaItemIds?: unknown;
+}) {
+  const rawImageUrls = normalizeStreamMediaUrls(params.imageUrls);
+  const rawVideoUrls = normalizeStreamMediaUrls(params.videoUrls);
+  const rawEntries = [
+    ...rawImageUrls.map((url, index) => ({
+      type: "image" as const,
+      url,
+      originalIndex: index,
+    })),
+    ...rawVideoUrls.map((url, index) => ({
+      type: "video" as const,
+      url,
+      originalIndex: rawImageUrls.length + index,
+    })),
+  ];
+  const rawMediaItemIds = normalizeStreamIdList(params.mediaItemIds);
+  const validEntries = rawEntries.filter((entry) =>
+    isValidStreamMediaUrl(entry.url)
+  );
+
+  const idByValidIndex = new Map<number, string>();
+  if (rawMediaItemIds.length === rawEntries.length) {
+    validEntries.forEach((entry, validIndex) => {
+      const id = rawMediaItemIds[entry.originalIndex];
+      if (id) idByValidIndex.set(validIndex, id);
+    });
+  } else if (rawMediaItemIds.length === validEntries.length) {
+    rawMediaItemIds.forEach((id, index) => {
+      if (id) idByValidIndex.set(index, id);
+    });
+  }
+
+  const dedupedEntries: Array<{
+    type: "image" | "video";
+    url: string;
+    mediaItemId?: string;
+  }> = [];
+  const seenMediaKeys = new Set<string>();
+  validEntries.forEach((entry, validIndex) => {
+    const key = `${entry.type}:${entry.url}`;
+    if (seenMediaKeys.has(key)) return;
+    seenMediaKeys.add(key);
+    dedupedEntries.push({
+      type: entry.type,
+      url: entry.url,
+      mediaItemId: idByValidIndex.get(validIndex),
+    });
+  });
+
+  if (
+    dedupedEntries.length > 0 &&
+    dedupedEntries.some((entry) => !entry.mediaItemId) &&
+    rawMediaItemIds.length === dedupedEntries.length
+  ) {
+    dedupedEntries.forEach((entry, index) => {
+      entry.mediaItemId = rawMediaItemIds[index];
+    });
+  }
+
+  const mediaItemIds = dedupedEntries.every((entry) => entry.mediaItemId)
+    ? dedupedEntries.map((entry) => entry.mediaItemId!)
+    : [];
+
+  return {
+    imageUrls: dedupedEntries
+      .filter((entry) => entry.type === "image")
+      .map((entry) => entry.url),
+    videoUrls: dedupedEntries
+      .filter((entry) => entry.type === "video")
+      .map((entry) => entry.url),
+    mediaItemIds,
+    handledCanvasMediaItemIds: normalizeStreamIdList(
+      params.handledCanvasMediaItemIds
+    ),
+  };
+}
+
 export function writeProtocolDebugChunk(
   writer: StreamWriter,
   enabled: boolean,
@@ -209,26 +301,20 @@ export function writeCanvasMediaChunk(
   },
   label: string
 ) {
-  const imageUrls = Array.isArray(params.imageUrls)
+  const inputImageUrls = Array.isArray(params.imageUrls)
     ? params.imageUrls
     : extractImageUrlsFromText(String(params.text || ""));
-  const videoUrls = Array.isArray(params.videoUrls) ? params.videoUrls : [];
-  const normalizedImageUrls = Array.from(
-    new Set(
-      imageUrls
-        .map((url) => String(url || "").trim())
-        .filter((url) => /^https?:\/\//i.test(url))
-    )
-  );
-  const normalizedVideoUrls = Array.from(
-    new Set(
-      videoUrls
-        .map((url) => String(url || "").trim())
-        .filter((url) => /^https?:\/\//i.test(url))
-    )
-  );
+  const normalizedPayload = normalizeCanvasMediaStreamPayload({
+    imageUrls: inputImageUrls,
+    videoUrls: params.videoUrls,
+    mediaItemIds: params.mediaItemIds,
+    handledCanvasMediaItemIds: params.handledCanvasMediaItemIds,
+  });
 
-  if (normalizedImageUrls.length === 0 && normalizedVideoUrls.length === 0) {
+  if (
+    normalizedPayload.imageUrls.length === 0 &&
+    normalizedPayload.videoUrls.length === 0
+  ) {
     return;
   }
 
@@ -237,8 +323,8 @@ export function writeCanvasMediaChunk(
     "data-canvasMedia",
     {
       messageId: params.messageId,
-      imageUrls: normalizedImageUrls,
-      videoUrls: normalizedVideoUrls,
+      imageUrls: normalizedPayload.imageUrls,
+      videoUrls: normalizedPayload.videoUrls,
       prompt: String(params.prompt || "").trim() || undefined,
       sourceImageUrls: Array.isArray(params.sourceImageUrls)
         ? params.sourceImageUrls
@@ -251,10 +337,8 @@ export function writeCanvasMediaChunk(
       agentCategoryName: String(params.agentCategoryName || "").trim() || undefined,
       agentExecutionSnapshotId:
         String(params.agentExecutionSnapshotId || "").trim() || undefined,
-      mediaItemIds: normalizeStreamIdList(params.mediaItemIds),
-      handledCanvasMediaItemIds: normalizeStreamIdList(
-        params.handledCanvasMediaItemIds
-      ),
+      mediaItemIds: normalizedPayload.mediaItemIds,
+      handledCanvasMediaItemIds: normalizedPayload.handledCanvasMediaItemIds,
     },
     label
   );

tests/routes/canvas-chat-helpers.test.ts

@@ -29,6 +29,7 @@ import {
   type CanvasAgentPendingRequestWindow,
 } from "@/lib/canvas/agent-pending-requests";
 import {
+  normalizeCanvasMediaStreamPayload,
   writeCanvasMediaChunk,
   writeCanvasOperationChunk,
 } from "@/app/(chat)/api/chat/_lib/stream";
@@ -2422,6 +2423,54 @@ test.describe("canvas chat pure helpers", () => {
     });
   });
 
+  test("keeps canvasMedia ids aligned with normalized batch media urls", () => {
+    expect(
+      normalizeCanvasMediaStreamPayload({
+        imageUrls: [
+          "not-a-url",
+          "https://cdn.example.com/a.png",
+          "https://cdn.example.com/a.png",
+          "https://cdn.example.com/b.png",
+        ],
+        videoUrls: ["https://cdn.example.com/c.mp4"],
+        mediaItemIds: [
+          "media-invalid",
+          "media-a",
+          "media-a-duplicate",
+          "media-b",
+          "media-c",
+        ],
+        handledCanvasMediaItemIds: [" handled-1 ", ""],
+      })
+    ).toEqual({
+      imageUrls: [
+        "https://cdn.example.com/a.png",
+        "https://cdn.example.com/b.png",
+      ],
+      videoUrls: ["https://cdn.example.com/c.mp4"],
+      mediaItemIds: ["media-a", "media-b", "media-c"],
+      handledCanvasMediaItemIds: ["handled-1"],
+    });
+
+    expect(
+      normalizeCanvasMediaStreamPayload({
+        imageUrls: [
+          "https://cdn.example.com/a.png",
+          "https://cdn.example.com/b.png",
+        ],
+        mediaItemIds: ["only-one-id"],
+      })
+    ).toEqual({
+      imageUrls: [
+        "https://cdn.example.com/a.png",
+        "https://cdn.example.com/b.png",
+      ],
+      videoUrls: [],
+      mediaItemIds: [],
+      handledCanvasMediaItemIds: [],
+    });
+  });
+
   test("writes applied, failed, and shadowed canvas operation stream chunks", () => {
     const chunks: Array<{ type: string; data: string }> = [];
     const writer = {