ai 6.0.232 → 6.0.234

package/src/text-stream/pipe-text-stream-to-response.ts

@@ -13,6 +13,7 @@ import { writeToServerResponse } from '../util/write-to-server-response';
  * @param options.statusText - Optional HTTP status text.
  * @param options.headers - Optional response headers.
  * @param options.textStream - The text stream to pipe.
+ * @returns A promise that resolves when the stream has been written.
  */
 export function pipeTextStreamToResponse({
   response,
@@ -23,8 +24,8 @@ export function pipeTextStreamToResponse({
 }: {
   response: ServerResponse;
   textStream: ReadableStream<string>;
-} & ResponseInit): void {
-  writeToServerResponse({
+} & ResponseInit): Promise<void> {
+  return writeToServerResponse({
     response,
     status,
     statusText,

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

@@ -104,13 +104,43 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
     new TransformStream<UIMessageChunk, InferUIMessageChunk<UI_MESSAGE>>({
       async transform(chunk, controller) {
         await runUpdateMessageJob(async ({ state, write }) => {
+          function getCurrentStepParts() {
+            const parts = state.message.parts;
+            let currentStepStartIndex = parts.length - 1;
+
+            while (
+              currentStepStartIndex >= 0 &&
+              parts[currentStepStartIndex].type !== 'step-start'
+            ) {
+              currentStepStartIndex--;
+            }
+
+            return parts.slice(currentStepStartIndex + 1);
+          }
+
+          function getCurrentStepToolInvocations() {
+            return getCurrentStepParts().filter(isToolUIPart);
+          }
+
           function getToolInvocation(toolCallId: string) {
-            const toolInvocations = state.message.parts.filter(isToolUIPart);
+            const toolInvocations = getCurrentStepToolInvocations();
 
-            const toolInvocation = toolInvocations.find(
+            let toolInvocation = toolInvocations.find(
               invocation => invocation.toolCallId === toolCallId,
             );
 
+            if (toolInvocation == null) {
+              const parts = state.message.parts;
+
+              for (let i = parts.length - 1; i >= 0; i--) {
+                const part = parts[i];
+                if (isToolUIPart(part) && part.toolCallId === toolCallId) {
+                  toolInvocation = part;
+                  break;
+                }
+              }
+            }
+
             if (toolInvocation == null) {
               throw new UIMessageStreamError({
                 chunkType: 'tool-invocation',
@@ -159,12 +189,15 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   providerMetadata?: ProviderMetadata;
                 }
             ),
+            existingPart?: ToolUIPart<InferUIMessageTools<UI_MESSAGE>>,
           ) {
-            const part = state.message.parts.find(
-              part =>
-                isStaticToolUIPart(part) &&
-                part.toolCallId === options.toolCallId,
-            ) as ToolUIPart<InferUIMessageTools<UI_MESSAGE>> | undefined;
+            const part =
+              existingPart ??
+              (getCurrentStepParts().find(
+                part =>
+                  isStaticToolUIPart(part) &&
+                  part.toolCallId === options.toolCallId,
+              ) as ToolUIPart<InferUIMessageTools<UI_MESSAGE>> | undefined);
 
             const anyOptions = options as any;
             const anyPart = part as any;
@@ -266,12 +299,15 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   providerMetadata?: ProviderMetadata;
                 }
             ),
+            existingPart?: DynamicToolUIPart,
           ) {
-            const part = state.message.parts.find(
-              part =>
-                part.type === 'dynamic-tool' &&
-                part.toolCallId === options.toolCallId,
-            ) as DynamicToolUIPart | undefined;
+            const part =
+              existingPart ??
+              (getCurrentStepParts().find(
+                part =>
+                  part.type === 'dynamic-tool' &&
+                  part.toolCallId === options.toolCallId,
+              ) as DynamicToolUIPart | undefined);
 
             const anyOptions = options as any;
             const anyPart = part as any;
@@ -510,7 +546,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
 
             case 'tool-input-start': {
               const toolInvocations =
-                state.message.parts.filter(isStaticToolUIPart);
+                getCurrentStepParts().filter(isStaticToolUIPart);
 
               // add the partial tool call to the map
               state.partialToolCalls[chunk.toolCallId] = {
@@ -635,7 +671,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               // When a part already exists for this toolCallId (e.g. from
               // tool-input-start), honour its type so we update in place
               // instead of creating a duplicate with a mismatched type.
-              const existingPart = state.message.parts
+              const existingPart = getCurrentStepParts()
                 .filter(isToolUIPart)
                 .find(p => p.toolCallId === chunk.toolCallId);
               const isDynamic =
@@ -696,31 +732,37 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               const toolInvocation = getToolInvocation(chunk.toolCallId);
 
               if (toolInvocation.type === 'dynamic-tool') {
-                updateDynamicToolPart({
-                  toolCallId: chunk.toolCallId,
-                  toolName: toolInvocation.toolName,
-                  state: 'output-available',
-                  input: (toolInvocation as any).input,
-                  output: chunk.output,
-                  preliminary: chunk.preliminary,
-                  providerExecuted: chunk.providerExecuted,
-                  providerMetadata: chunk.providerMetadata,
-                  title: toolInvocation.title,
-                  toolMetadata: toolInvocation.toolMetadata,
-                });
+                updateDynamicToolPart(
+                  {
+                    toolCallId: chunk.toolCallId,
+                    toolName: toolInvocation.toolName,
+                    state: 'output-available',
+                    input: (toolInvocation as any).input,
+                    output: chunk.output,
+                    preliminary: chunk.preliminary,
+                    providerExecuted: chunk.providerExecuted,
+                    providerMetadata: chunk.providerMetadata,
+                    title: toolInvocation.title,
+                    toolMetadata: toolInvocation.toolMetadata,
+                  },
+                  toolInvocation,
+                );
               } else {
-                updateToolPart({
-                  toolCallId: chunk.toolCallId,
-                  toolName: getStaticToolName(toolInvocation),
-                  state: 'output-available',
-                  input: (toolInvocation as any).input,
-                  output: chunk.output,
-                  providerExecuted: chunk.providerExecuted,
-                  preliminary: chunk.preliminary,
-                  providerMetadata: chunk.providerMetadata,
-                  title: toolInvocation.title,
-                  toolMetadata: toolInvocation.toolMetadata,
-                });
+                updateToolPart(
+                  {
+                    toolCallId: chunk.toolCallId,
+                    toolName: getStaticToolName(toolInvocation),
+                    state: 'output-available',
+                    input: (toolInvocation as any).input,
+                    output: chunk.output,
+                    providerExecuted: chunk.providerExecuted,
+                    preliminary: chunk.preliminary,
+                    providerMetadata: chunk.providerMetadata,
+                    title: toolInvocation.title,
+                    toolMetadata: toolInvocation.toolMetadata,
+                  },
+                  toolInvocation as ToolUIPart<InferUIMessageTools<UI_MESSAGE>>,
+                );
               }
 
               write();
@@ -731,30 +773,36 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               const toolInvocation = getToolInvocation(chunk.toolCallId);
 
               if (toolInvocation.type === 'dynamic-tool') {
-                updateDynamicToolPart({
-                  toolCallId: chunk.toolCallId,
-                  toolName: toolInvocation.toolName,
-                  state: 'output-error',
-                  input: (toolInvocation as any).input,
-                  errorText: chunk.errorText,
-                  providerExecuted: chunk.providerExecuted,
-                  providerMetadata: chunk.providerMetadata,
-                  title: toolInvocation.title,
-                  toolMetadata: toolInvocation.toolMetadata,
-                });
+                updateDynamicToolPart(
+                  {
+                    toolCallId: chunk.toolCallId,
+                    toolName: toolInvocation.toolName,
+                    state: 'output-error',
+                    input: (toolInvocation as any).input,
+                    errorText: chunk.errorText,
+                    providerExecuted: chunk.providerExecuted,
+                    providerMetadata: chunk.providerMetadata,
+                    title: toolInvocation.title,
+                    toolMetadata: toolInvocation.toolMetadata,
+                  },
+                  toolInvocation,
+                );
               } else {
-                updateToolPart({
-                  toolCallId: chunk.toolCallId,
-                  toolName: getStaticToolName(toolInvocation),
-                  state: 'output-error',
-                  input: (toolInvocation as any).input,
-                  rawInput: (toolInvocation as any).rawInput,
-                  errorText: chunk.errorText,
-                  providerExecuted: chunk.providerExecuted,
-                  providerMetadata: chunk.providerMetadata,
-                  title: toolInvocation.title,
-                  toolMetadata: toolInvocation.toolMetadata,
-                });
+                updateToolPart(
+                  {
+                    toolCallId: chunk.toolCallId,
+                    toolName: getStaticToolName(toolInvocation),
+                    state: 'output-error',
+                    input: (toolInvocation as any).input,
+                    rawInput: (toolInvocation as any).rawInput,
+                    errorText: chunk.errorText,
+                    providerExecuted: chunk.providerExecuted,
+                    providerMetadata: chunk.providerMetadata,
+                    title: toolInvocation.title,
+                    toolMetadata: toolInvocation.toolMetadata,
+                  },
+                  toolInvocation as ToolUIPart<InferUIMessageTools<UI_MESSAGE>>,
+                );
               }
 
               write();

package/src/ui-message-stream/pipe-ui-message-stream-to-response.ts

@@ -16,6 +16,7 @@ import type { UIMessageStreamResponseInit } from './ui-message-stream-response-i
  * @param options.headers - Additional HTTP headers to include in the response.
  * @param options.stream - The UI message chunk stream to send.
  * @param options.consumeSseStream - Optional callback to consume a copy of the SSE stream independently.
+ * @returns A promise that resolves when the stream has been written.
  */
 export function pipeUIMessageStreamToResponse({
   response,
@@ -27,7 +28,7 @@ export function pipeUIMessageStreamToResponse({
 }: {
   response: ServerResponse;
   stream: ReadableStream<UIMessageChunk>;
-} & UIMessageStreamResponseInit): void {
+} & UIMessageStreamResponseInit): Promise<void> {
   let sseStream = stream.pipeThrough(new JsonToSseTransformStream());
 
   // when the consumeSseStream is provided, we need to tee the stream
@@ -39,7 +40,7 @@ export function pipeUIMessageStreamToResponse({
     consumeSseStream({ stream: stream2 }); // no await (do not block the response)
   }
 
-  writeToServerResponse({
+  return writeToServerResponse({
     response,
     status,
     statusText,

package/src/util/detect-media-type.ts

@@ -112,7 +112,16 @@ export const audioMediaTypeSignatures = [
   },
   {
     mediaType: 'audio/mp4' as const,
-    bytesPrefix: [0x66, 0x74, 0x79, 0x70],
+    bytesPrefix: [
+      0x00,
+      0x00,
+      0x00,
+      null,
+      0x66,
+      0x74,
+      0x79,
+      0x70, // ftyp
+    ],
   },
   {
     mediaType: 'audio/webm',
@@ -159,31 +168,54 @@ export const videoMediaTypeSignatures = [
   },
 ] as const;
 
-const stripID3 = (data: Uint8Array | string) => {
-  const bytes =
-    typeof data === 'string' ? convertBase64ToUint8Array(data) : data;
+const DEFAULT_SNIFF_BYTES = 18;
+
+// Longest signature prefix in the tables above (e.g. image/avif = 12 bytes).
+const MAX_SIGNATURE_BYTES = 12;
+
+// Largest ID3v2 tag (10-byte header + body) skipped to reach the audio frame.
+// Covers typical tags including embedded cover art while keeping the decode
+// bounded and O(1) in the attachment size. Exported for boundary tests.
+export const MAX_ID3_TAG_BYTES = 128 * 1024;
+
+// Total prefix decoded when an ID3 tag is present: the tag plus room for the
+// trailing signature, so a tag right at the size limit stays detectable.
+const ID3_SCAN_BYTES = MAX_ID3_TAG_BYTES + MAX_SIGNATURE_BYTES;
+
+// Decode/view exactly the first `maxBytes` bytes from the front of the input.
+// The base64 and raw-byte paths yield the same length, so detection does not
+// depend on the input's representation.
+function decodePrefix(data: Uint8Array | string, maxBytes: number): Uint8Array {
+  if (typeof data !== 'string') {
+    return data.length > maxBytes ? data.subarray(0, maxBytes) : data;
+  }
+  // base64: 4 chars -> 3 bytes. Decode whole 4-char groups, then trim the 0-2
+  // extra bytes so the result matches the raw-byte path exactly.
+  const maxChars = Math.ceil(maxBytes / 3) * 4;
+  const bytes = convertBase64ToUint8Array(
+    data.substring(0, Math.min(data.length, maxChars)),
+  );
+  return bytes.length > maxBytes ? bytes.subarray(0, maxBytes) : bytes;
+}
+
+function hasID3(bytes: Uint8Array): boolean {
+  return (
+    bytes.length > 10 &&
+    bytes[0] === 0x49 && // 'I'
+    bytes[1] === 0x44 && // 'D'
+    bytes[2] === 0x33 // '3'
+  );
+}
+
+const stripID3 = (bytes: Uint8Array): Uint8Array => {
   const id3Size =
     ((bytes[6] & 0x7f) << 21) |
     ((bytes[7] & 0x7f) << 14) |
     ((bytes[8] & 0x7f) << 7) |
     (bytes[9] & 0x7f);
-
-  // The raw MP3 starts here
-  return bytes.slice(id3Size + 10);
+  return bytes.subarray(id3Size + 10);
 };
 
-function stripID3TagsIfPresent(data: Uint8Array | string): Uint8Array | string {
-  const hasId3 =
-    (typeof data === 'string' && data.startsWith('SUQz')) ||
-    (typeof data !== 'string' &&
-      data.length > 10 &&
-      data[0] === 0x49 && // 'I'
-      data[1] === 0x44 && // 'D'
-      data[2] === 0x33); // '3'
-
-  return hasId3 ? stripID3(data) : data;
-}
-
 /**
  * Detect the media IANA media type of a file using a list of signatures.
  *
@@ -201,15 +233,13 @@ export function detectMediaType({
     | typeof imageMediaTypeSignatures
     | typeof videoMediaTypeSignatures;
 }): (typeof signatures)[number]['mediaType'] | undefined {
-  const processedData = stripID3TagsIfPresent(data);
+  let bytes = decodePrefix(data, DEFAULT_SNIFF_BYTES);
 
-  // Convert the first ~18 bytes (24 base64 chars) for consistent detection logic:
-  const bytes =
-    typeof processedData === 'string'
-      ? convertBase64ToUint8Array(
-          processedData.substring(0, Math.min(processedData.length, 24)),
-        )
-      : processedData;
+  // ID3v2-tagged MP3s carry the audio frame after the tag; scan a bounded
+  // prefix past it rather than decoding the whole input.
+  if (hasID3(bytes)) {
+    bytes = stripID3(decodePrefix(data, ID3_SCAN_BYTES));
+  }
 
   for (const signature of signatures) {
     if (

package/src/util/write-to-server-response.ts

@@ -15,7 +15,7 @@ export function writeToServerResponse({
   statusText?: string;
   headers?: Record<string, string | number | string[]>;
   stream: ReadableStream<Uint8Array>;
-}): void {
+}): Promise<void> {
   const statusCode = status ?? 200;
   if (statusText !== undefined) {
     response.writeHead(statusCode, statusText, headers);
@@ -45,5 +45,5 @@ export function writeToServerResponse({
     }
   };
 
-  read();
+  return read();
 }