refactor: improve streaming implementation

Author: stainless-bot

src/streaming.ts

@@ -13,61 +13,6 @@ type ServerSentEvent = {
   raw: string[];
 };
 
-class SSEDecoder {
-  private data: string[];
-  private event: string | null;
-  private chunks: string[];
-
-  constructor() {
-    this.event = null;
-    this.data = [];
-    this.chunks = [];
-  }
-
-  decode(line: string) {
-    if (line.endsWith('\r')) {
-      line = line.substring(0, line.length - 1);
-    }
-
-    if (!line) {
-      // empty line and we didn't previously encounter any messages
-      if (!this.event && !this.data.length) return null;
-
-      const sse: ServerSentEvent = {
-        event: this.event,
-        data: this.data.join('\n'),
-        raw: this.chunks,
-      };
-
-      this.event = null;
-      this.data = [];
-      this.chunks = [];
-
-      return sse;
-    }
-
-    this.chunks.push(line);
-
-    if (line.startsWith(':')) {
-      return null;
-    }
-
-    let [fieldname, _, value] = partition(line, ':');
-
-    if (value.startsWith(' ')) {
-      value = value.substring(1);
-    }
-
-    if (fieldname === 'event') {
-      this.event = value;
-    } else if (fieldname === 'data') {
-      this.data.push(value);
-    }
-
-    return null;
-  }
-}
-
 export class Stream<Item> implements AsyncIterable<Item>, APIResponse<Stream<Item>> {
   /** @deprecated - please use the async iterator instead. We plan to add additional helper methods shortly. */
   response: Response;
@@ -93,9 +38,7 @@ export class Stream<Item> implements AsyncIterable<Item>, APIResponse<Stream<Ite
 
     const iter = readableStreamAsyncIterable<Bytes>(this.response.body);
     for await (const chunk of iter) {
-      const text = decodeText(chunk);
-
-      for (const line of lineDecoder.decode(text)) {
+      for (const line of lineDecoder.decode(chunk)) {
         const sse = this.decoder.decode(line);
         if (sse) yield sse;
       }
@@ -143,7 +86,60 @@ export class Stream<Item> implements AsyncIterable<Item>, APIResponse<Stream<Ite
   }
 }
 
-const NEWLINE_CHARS = '\n\r\x0b\x0c\x1c\x1d\x1e\x85\u2028\u2029';
+class SSEDecoder {
+  private data: string[];
+  private event: string | null;
+  private chunks: string[];
+
+  constructor() {
+    this.event = null;
+    this.data = [];
+    this.chunks = [];
+  }
+
+  decode(line: string) {
+    if (line.endsWith('\r')) {
+      line = line.substring(0, line.length - 1);
+    }
+
+    if (!line) {
+      // empty line and we didn't previously encounter any messages
+      if (!this.event && !this.data.length) return null;
+
+      const sse: ServerSentEvent = {
+        event: this.event,
+        data: this.data.join('\n'),
+        raw: this.chunks,
+      };
+
+      this.event = null;
+      this.data = [];
+      this.chunks = [];
+
+      return sse;
+    }
+
+    this.chunks.push(line);
+
+    if (line.startsWith(':')) {
+      return null;
+    }
+
+    let [fieldname, _, value] = partition(line, ':');
+
+    if (value.startsWith(' ')) {
+      value = value.substring(1);
+    }
+
+    if (fieldname === 'event') {
+      this.event = value;
+    } else if (fieldname === 'data') {
+      this.data.push(value);
+    }
+
+    return null;
+  }
+}
 
 /**
  * A re-implementation of httpx's `LineDecoder` in Python that handles incrementally
@@ -152,15 +148,22 @@ const NEWLINE_CHARS = '\n\r\x0b\x0c\x1c\x1d\x1e\x85\u2028\u2029';
  * https://github.com/encode/httpx/blob/920333ea98118e9cf617f246905d7b202510941c/httpx/_decoders.py#L258
  */
 class LineDecoder {
+  // prettier-ignore
+  static NEWLINE_CHARS = new Set(['\n', '\r', '\x0b', '\x0c', '\x1c', '\x1d', '\x1e', '\x85', '\u2028', '\u2029']);
+  static NEWLINE_REGEXP = /\r\n|[\n\r\x0b\x0c\x1c\x1d\x1e\x85\u2028\u2029]/g;
+
   buffer: string[];
   trailingCR: boolean;
+  textDecoder: any; // TextDecoder found in browsers; not typed to avoid pulling in either "dom" or "node" types.
 
   constructor() {
     this.buffer = [];
     this.trailingCR = false;
   }
 
-  decode(text: string): string[] {
+  decode(chunk: Bytes): string[] {
+    let text = this.decodeText(chunk);
+
     if (this.trailingCR) {
       text = '\r' + text;
       this.trailingCR = false;
@@ -174,10 +177,10 @@ class LineDecoder {
       return [];
     }
 
-    const trailing_newline = NEWLINE_CHARS.includes(text.slice(-1));
-    let lines = text.split(/\r\n|[\n\r\x0b\x0c\x1c\x1d\x1e\x85\u2028\u2029]/g);
+    const trailingNewline = LineDecoder.NEWLINE_CHARS.has(text[text.length - 1] || '');
+    let lines = text.split(LineDecoder.NEWLINE_REGEXP);
 
-    if (lines.length === 1 && !trailing_newline) {
+    if (lines.length === 1 && !trailingNewline) {
       this.buffer.push(lines[0]!);
       return [];
     }
@@ -187,13 +190,50 @@ class LineDecoder {
       this.buffer = [];
     }
 
-    if (!trailing_newline) {
+    if (!trailingNewline) {
       this.buffer = [lines.pop() || ''];
     }
 
     return lines;
   }
 
+  decodeText(bytes: Bytes): string {
+    if (bytes == null) return '';
+    if (typeof bytes === 'string') return bytes;
+
+    // Node:
+    if (typeof Buffer !== 'undefined') {
+      if (bytes instanceof Buffer) {
+        return bytes.toString();
+      }
+      if (bytes instanceof Uint8Array) {
+        return Buffer.from(bytes).toString();
+      }
+
+      throw new Error(
+        `Unexpected: received non-Uint8Array (${bytes.constructor.name}) stream chunk in an environment with a global "Buffer" defined, which this library assumes to be Node. Please report this error.`,
+      );
+    }
+
+    // Browser
+    if (typeof TextDecoder !== 'undefined') {
+      if (bytes instanceof Uint8Array || bytes instanceof ArrayBuffer) {
+        this.textDecoder ??= new TextDecoder('utf8');
+        return this.textDecoder.decode(bytes);
+      }
+
+      throw new Error(
+        `Unexpected: received non-Uint8Array/ArrayBuffer (${
+          (bytes as any).constructor.name
+        }) in a web platform. Please report this error.`,
+      );
+    }
+
+    throw new Error(
+      `Unexpected: neither Buffer nor TextDecoder are available as globals. Please report this error.`,
+    );
+  }
+
   flush(): string[] {
     if (!this.buffer.length && !this.trailingCR) {
       return [];
@@ -215,61 +255,22 @@ function partition(str: string, delimiter: string): [string, string, string] {
   return [str, '', ''];
 }
 
-let _textDecoder;
-function decodeText(bytes: Bytes): string {
-  if (bytes == null) return '';
-  if (typeof bytes === 'string') return bytes;
-
-  // Node:
-  if (typeof Buffer !== 'undefined') {
-    if (bytes instanceof Buffer) {
-      return bytes.toString();
-    }
-    if (bytes instanceof Uint8Array) {
-      return Buffer.from(bytes).toString();
-    }
-
-    throw new Error(`Unexpected: received non-Uint8Array (${bytes.constructor.name}) in Node.`);
-  }
-
-  // Browser
-  if (typeof TextDecoder !== 'undefined') {
-    if (bytes instanceof Uint8Array || bytes instanceof ArrayBuffer) {
-      _textDecoder ??= new TextDecoder('utf8');
-      return _textDecoder.decode(bytes);
-    }
-
-    throw new Error(
-      `Unexpected: received non-Uint8Array/ArrayBuffer (${
-        (bytes as any).constructor.name
-      }) in a web platform.`,
-    );
-  }
-
-  throw new Error(`Unexpected: neither Buffer nor TextDecoder are available as globals.`);
-}
-
 /**
  * Most browsers don't yet have async iterable support for ReadableStream,
  * and Node has a very different way of reading bytes from its "ReadableStream".
  *
  * This polyfill was pulled from https://github.com/MattiasBuelens/web-streams-polyfill/pull/122#issuecomment-1624185965
- *
- * We make extensive use of "any" here to avoid pulling in either "node" or "dom" types
- * to library users' type scopes.
  */
 function readableStreamAsyncIterable<T>(stream: any): AsyncIterableIterator<T> {
-  if (stream[Symbol.asyncIterator]) {
-    return stream[Symbol.asyncIterator];
-  }
+  if (stream[Symbol.asyncIterator]) return stream[Symbol.asyncIterator];
 
   const reader = stream.getReader();
-
   return {
     next() {
       return reader.read();
     },
     async return() {
+      reader.cancel();
       reader.releaseLock();
       return { done: true, value: undefined };
     },