fix (core,streams): support ResponseInit variants (#1766)

lgrammel · sshader · web-flow · commit 213f2411fb9d · 2024-05-30T19:10:25.000+02:00
Co-authored-by: Sarah Shader &lt;sarahshader@gmail.com&gt;
diff --git a/.changeset/eighty-kings-attend.md b/.changeset/eighty-kings-attend.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (core,streams): support ResponseInit variants
diff --git a/packages/core/core/generate-text/stream-text.ts b/packages/core/core/generate-text/stream-text.ts
@@ -22,6 +22,7 @@ import {
   AsyncIterableStream,
   createAsyncIterableStream,
 } from '../util/async-iterable-stream';
+import { prepareResponseHeaders } from '../util/prepare-response-headers';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { runToolsTransformation } from './run-tools-transformation';
 import { TokenUsage } from './token-usage';
@@ -600,10 +601,9 @@ Non-text-delta events are ignored.
   toTextStreamResponse(init?: ResponseInit): Response {
     return new Response(this.textStream.pipeThrough(new TextEncoderStream()), {
       status: init?.status ?? 200,
-      headers: {
-        'Content-Type': 'text/plain; charset=utf-8',
-        ...init?.headers,
-      },
+      headers: prepareResponseHeaders(init, {
+        contentType: 'text/plain; charset=utf-8',
+      }),
     });
   }
 }
diff --git a/packages/core/core/util/prepare-response-headers.test.ts b/packages/core/core/util/prepare-response-headers.test.ts
@@ -0,0 +1,53 @@
+import { expect, it } from 'vitest';
+import { prepareResponseHeaders } from './prepare-response-headers';
+
+it('should set Content-Type header if not present', () => {
+  const headers = prepareResponseHeaders(
+    {},
+    { contentType: 'application/json' },
+  );
+
+  expect(headers.get('Content-Type')).toBe('application/json');
+});
+
+it('should not overwrite existing Content-Type header', () => {
+  const headers = prepareResponseHeaders(
+    { headers: { 'Content-Type': 'text/html' } },
+    { contentType: 'application/json' },
+  );
+
+  expect(headers.get('Content-Type')).toBe('text/html');
+});
+
+it('should handle undefined init', () => {
+  const headers = prepareResponseHeaders(undefined, {
+    contentType: 'application/json',
+  });
+
+  expect(headers.get('Content-Type')).toBe('application/json');
+});
+
+it('should handle init headers as Headers object', () => {
+  const headers = prepareResponseHeaders(
+    { headers: new Headers({ init: 'foo' }) },
+    { contentType: 'application/json' },
+  );
+
+  expect(headers.get('init')).toBe('foo');
+  expect(headers.get('Content-Type')).toBe('application/json');
+});
+
+it('should handle Response object headers', () => {
+  const initHeaders = { init: 'foo' };
+  const response = new Response(null, {
+    headers: { ...initHeaders, extra: 'bar' },
+  });
+
+  const headers = prepareResponseHeaders(response, {
+    contentType: 'application/json',
+  });
+
+  expect(headers.get('init')).toBe('foo');
+  expect(headers.get('extra')).toBe('bar');
+  expect(headers.get('Content-Type')).toBe('application/json');
+});
diff --git a/packages/core/core/util/prepare-response-headers.ts b/packages/core/core/util/prepare-response-headers.ts
@@ -0,0 +1,12 @@
+export function prepareResponseHeaders(
+  init: ResponseInit | undefined,
+  { contentType }: { contentType: string },
+) {
+  const headers = new Headers(init?.headers ?? {});
+
+  if (!headers.has('Content-Type')) {
+    headers.set('Content-Type', contentType);
+  }
+
+  return headers;
+}
diff --git a/packages/core/core/util/split-array.test.ts b/packages/core/core/util/split-array.test.ts
@@ -1,62 +1,60 @@
-import { describe, it, expect } from 'vitest';
+import { expect, it } from 'vitest';
 import { splitArray } from './split-array';
 
-describe('splitArray', () => {
-  it('should split an array into chunks of the specified size', () => {
-    const array = [1, 2, 3, 4, 5];
-    const size = 2;
-    const result = splitArray(array, size);
-    expect(result).toEqual([[1, 2], [3, 4], [5]]);
-  });
-
-  it('should return an empty array when the input array is empty', () => {
-    const array: number[] = [];
-    const size = 2;
-    const result = splitArray(array, size);
-    expect(result).toEqual([]);
-  });
-
-  it('should return the original array when the chunk size is greater than the array length', () => {
-    const array = [1, 2, 3];
-    const size = 5;
-    const result = splitArray(array, size);
-    expect(result).toEqual([[1, 2, 3]]);
-  });
-
-  it('should return the original array when the chunk size is equal to the array length', () => {
-    const array = [1, 2, 3];
-    const size = 3;
-    const result = splitArray(array, size);
-    expect(result).toEqual([[1, 2, 3]]);
-  });
-
-  it('should handle chunk size of 1 correctly', () => {
-    const array = [1, 2, 3];
-    const size = 1;
-    const result = splitArray(array, size);
-    expect(result).toEqual([[1], [2], [3]]);
-  });
-
-  it('should throw an error for chunk size of 0', () => {
-    const array = [1, 2, 3];
-    const size = 0;
-    expect(() => splitArray(array, size)).toThrow(
-      'chunkSize must be greater than 0',
-    );
-  });
-
-  it('should throw an error for negative chunk size', () => {
-    const array = [1, 2, 3];
-    const size = -1;
-    expect(() => splitArray(array, size)).toThrow(
-      'chunkSize must be greater than 0',
-    );
-  });
-
-  it('should handle non-integer chunk size by flooring the size', () => {
-    const array = [1, 2, 3, 4, 5];
-    const size = 2.5;
-    const result = splitArray(array, Math.floor(size));
-    expect(result).toEqual([[1, 2], [3, 4], [5]]);
-  });
+it('should split an array into chunks of the specified size', () => {
+  const array = [1, 2, 3, 4, 5];
+  const size = 2;
+  const result = splitArray(array, size);
+  expect(result).toEqual([[1, 2], [3, 4], [5]]);
+});
+
+it('should return an empty array when the input array is empty', () => {
+  const array: number[] = [];
+  const size = 2;
+  const result = splitArray(array, size);
+  expect(result).toEqual([]);
+});
+
+it('should return the original array when the chunk size is greater than the array length', () => {
+  const array = [1, 2, 3];
+  const size = 5;
+  const result = splitArray(array, size);
+  expect(result).toEqual([[1, 2, 3]]);
+});
+
+it('should return the original array when the chunk size is equal to the array length', () => {
+  const array = [1, 2, 3];
+  const size = 3;
+  const result = splitArray(array, size);
+  expect(result).toEqual([[1, 2, 3]]);
+});
+
+it('should handle chunk size of 1 correctly', () => {
+  const array = [1, 2, 3];
+  const size = 1;
+  const result = splitArray(array, size);
+  expect(result).toEqual([[1], [2], [3]]);
+});
+
+it('should throw an error for chunk size of 0', () => {
+  const array = [1, 2, 3];
+  const size = 0;
+  expect(() => splitArray(array, size)).toThrow(
+    'chunkSize must be greater than 0',
+  );
+});
+
+it('should throw an error for negative chunk size', () => {
+  const array = [1, 2, 3];
+  const size = -1;
+  expect(() => splitArray(array, size)).toThrow(
+    'chunkSize must be greater than 0',
+  );
+});
+
+it('should handle non-integer chunk size by flooring the size', () => {
+  const array = [1, 2, 3, 4, 5];
+  const size = 2.5;
+  const result = splitArray(array, Math.floor(size));
+  expect(result).toEqual([[1, 2], [3, 4], [5]]);
 });
diff --git a/packages/core/streams/streaming-text-response.ts b/packages/core/streams/streaming-text-response.ts
@@ -1,4 +1,5 @@
 import { mergeStreams } from '../core/util/merge-streams';
+import { prepareResponseHeaders } from '../core/util/prepare-response-headers';
 import { StreamData } from './stream-data';
 
 /**
@@ -15,10 +16,9 @@ export class StreamingTextResponse extends Response {
     super(processedStream as any, {
       ...init,
       status: 200,
-      headers: {
-        'Content-Type': 'text/plain; charset=utf-8',
-        ...init?.headers,
-      },
+      headers: prepareResponseHeaders(init, {
+        contentType: 'text/plain; charset=utf-8',
+      }),
     });
   }
 }

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'ai': patch
 +---
++
 +fix (core,streams): support ResponseInit variants