chore: Optimize streaming

This CL improves the efficiency of the streaming LineDecoder and SSE message iterator, specifically addressing performance degradation when processing large or highly fragmented payloads.

Key changes:
   * Optimized Buffer Scanning: Introduced searchIndex in LineDecoder to track the last searched position. This
     prevents the decoder from rescanning the entire accumulated buffer from the beginning for every new chunk,
     reducing complexity from $O(n^2)$ to $O(n)$.
   * Faster Pattern Matching: Replaced manual byte-by-byte loops with Uint8Array.prototype.indexOf in
     findNewlineIndex and findDoubleNewlineIndex for better performance.
   * Simplified Streaming Logic: Replaced iterSSEChunks with iterBinaryChunks in streaming.ts. This removes
     redundant buffering logic, as the stateful LineDecoder already handles fragmented line assembly.
   * Boundary Safety: Implemented a 1-byte search overlap (length - 1) to correctly detect multi-byte newline
     sequences (like \r\n) when they are split across chunk boundaries.
   * Regression Testing: Added unit tests in api_client_test.ts and interactions_test.ts to verify that large,
     fragmented payloads are handled correctly.

PiperOrigin-RevId: 878636096

Author: MarkDaoust

src/interactions/core/streaming.ts

@@ -7,7 +7,7 @@
 import { GeminiNextGenAPIClientError } from './error.js';
 import { type ReadableStream } from '../internal/shim-types.js';
 import { makeReadableStream } from '../internal/shims.js';
-import { findDoubleNewlineIndex, LineDecoder } from '../internal/decoders/line.js';
+import { LineDecoder } from '../internal/decoders/line.js';
 import { ReadableStreamToAsyncIterable } from '../internal/shims.js';
 import { isAbortError } from '../internal/errors.js';
 import { encodeUTF8 } from '../internal/utils/bytes.js';
@@ -222,7 +222,7 @@ export async function* _iterSSEMessages(
   const lineDecoder = new LineDecoder();
 
   const iter = ReadableStreamToAsyncIterable<Bytes>(response.body);
-  for await (const sseChunk of iterSSEChunks(iter)) {
+  for await (const sseChunk of iterBinaryChunks(iter)) {
     for (const line of lineDecoder.decode(sseChunk)) {
       const sse = sseDecoder.decode(line);
       if (sse) yield sse;
@@ -236,12 +236,10 @@ export async function* _iterSSEMessages(
 }
 
 /**
- * Given an async iterable iterator, iterates over it and yields full
- * SSE chunks, i.e. yields when a double new-line is encountered.
+ * Given an async iterable iterator, normalizes each chunk to a
+ * Uint8Array and yields it.
  */
-async function* iterSSEChunks(iterator: AsyncIterableIterator<Bytes>): AsyncGenerator<Uint8Array> {
-  let data = new Uint8Array();
-
+async function* iterBinaryChunks(iterator: AsyncIterableIterator<Bytes>): AsyncGenerator<Uint8Array> {
   for await (const chunk of iterator) {
     if (chunk == null) {
       continue;
@@ -252,20 +250,7 @@ async function* iterSSEChunks(iterator: AsyncIterableIterator<Bytes>): AsyncGene
       : typeof chunk === 'string' ? encodeUTF8(chunk)
       : chunk;
 
-    let newData = new Uint8Array(data.length + binaryChunk.length);
-    newData.set(data);
-    newData.set(binaryChunk, data.length);
-    data = newData;
-
-    let patternIndex;
-    while ((patternIndex = findDoubleNewlineIndex(data)) !== -1) {
-      yield data.slice(0, patternIndex);
-      data = data.slice(patternIndex);
-    }
-  }
-
-  if (data.length > 0) {
-    yield data;
+    yield binaryChunk;
   }
 }
 

src/interactions/internal/decoders/line.ts

@@ -21,10 +21,12 @@ export class LineDecoder {
 
   private buffer: Uint8Array;
   private carriageReturnIndex: number | null;
+  private searchIndex: number;
 
   constructor() {
     this.buffer = new Uint8Array();
     this.carriageReturnIndex = null;
+    this.searchIndex = 0;
   }
 
   decode(chunk: Bytes): string[] {
@@ -41,7 +43,9 @@ export class LineDecoder {
 
     const lines: string[] = [];
     let patternIndex;
-    while ((patternIndex = findNewlineIndex(this.buffer, this.carriageReturnIndex)) != null) {
+    while (
+      (patternIndex = findNewlineIndex(this.buffer, this.carriageReturnIndex ?? this.searchIndex)) != null
+    ) {
       if (patternIndex.carriage && this.carriageReturnIndex == null) {
         // skip until we either get a corresponding `\n`, a new `\r` or nothing
         this.carriageReturnIndex = patternIndex.index;
@@ -56,6 +60,7 @@ export class LineDecoder {
         lines.push(decodeUTF8(this.buffer.subarray(0, this.carriageReturnIndex - 1)));
         this.buffer = this.buffer.subarray(this.carriageReturnIndex);
         this.carriageReturnIndex = null;
+        this.searchIndex = 0;
         continue;
       }
 
@@ -67,8 +72,11 @@ export class LineDecoder {
 
       this.buffer = this.buffer.subarray(patternIndex.index);
       this.carriageReturnIndex = null;
+      this.searchIndex = 0;
     }
 
+    this.searchIndex = Math.max(0, this.buffer.length - 1);
+
     return lines;
   }
 
@@ -96,45 +104,74 @@ function findNewlineIndex(
   const newline = 0x0a; // \n
   const carriage = 0x0d; // \r
 
-  for (let i = startIndex ?? 0; i < buffer.length; i++) {
-    if (buffer[i] === newline) {
-      return { preceding: i, index: i + 1, carriage: false };
-    }
+  const start = startIndex ?? 0;
+  const nextNewline = buffer.indexOf(newline, start);
+  const nextCarriage = buffer.indexOf(carriage, start);
 
-    if (buffer[i] === carriage) {
-      return { preceding: i, index: i + 1, carriage: true };
-    }
+  if (nextNewline === -1 && nextCarriage === -1) {
+    return null;
+  }
+
+  let i: number;
+  if (nextNewline !== -1 && nextCarriage !== -1) {
+    i = Math.min(nextNewline, nextCarriage);
+  } else {
+    i = nextNewline !== -1 ? nextNewline : nextCarriage;
+  }
+
+  if (buffer[i] === newline) {
+    return { preceding: i, index: i + 1, carriage: false };
   }
 
-  return null;
+  return { preceding: i, index: i + 1, carriage: true };
 }
 
-export function findDoubleNewlineIndex(buffer: Uint8Array): number {
+export function findDoubleNewlineIndex(buffer: Uint8Array, startIndex: number = 0): number {
   // This function searches the buffer for the end patterns (\r\r, \n\n, \r\n\r\n)
   // and returns the index right after the first occurrence of any pattern,
   // or -1 if none of the patterns are found.
   const newline = 0x0a; // \n
   const carriage = 0x0d; // \r
 
-  for (let i = 0; i < buffer.length - 1; i++) {
-    if (buffer[i] === newline && buffer[i + 1] === newline) {
+  let i = startIndex;
+  while (i < buffer.length - 1) {
+    const nextNewline = buffer.indexOf(newline, i);
+    const nextCarriage = buffer.indexOf(carriage, i);
+
+    if (nextNewline === -1 && nextCarriage === -1) {
+      return -1;
+    }
+
+    let index: number;
+    if (nextNewline !== -1 && nextCarriage !== -1) {
+      index = Math.min(nextNewline, nextCarriage);
+    } else {
+      index = nextNewline !== -1 ? nextNewline : nextCarriage;
+    }
+
+    if (index >= buffer.length - 1) {
+      return -1;
+    }
+
+    if (buffer[index] === newline && buffer[index + 1] === newline) {
       // \n\n
-      return i + 2;
+      return index + 2;
     }
-    if (buffer[i] === carriage && buffer[i + 1] === carriage) {
+    if (buffer[index] === carriage && buffer[index + 1] === carriage) {
       // \r\r
-      return i + 2;
+      return index + 2;
     }
     if (
-      buffer[i] === carriage &&
-      buffer[i + 1] === newline &&
-      i + 3 < buffer.length &&
-      buffer[i + 2] === carriage &&
-      buffer[i + 3] === newline
+      buffer[index] === carriage &&
+      buffer[index + 1] === newline &&
+      index + 3 < buffer.length &&
+      buffer[index + 2] === carriage &&
+      buffer[index + 3] === newline
     ) {
       // \r\n\r\n
-      return i + 4;
+      return index + 4;
     }
+    i = index + 1;
   }
 
   return -1;

test/unit/api_client_test.ts

@@ -381,6 +381,38 @@ describe('processStreamResponse', () => {
     const result = await resultHttpResponse.value.json();
     expect(result).toEqual(expectedResponse);
   });
+
+  it('should handle large fragmented payloads correctly (regression for $O(n^2)$)', async () => {
+    const largeData = 'A'.repeat(10 * 1024);
+    const jsonPayload = JSON.stringify({data: largeData});
+    const ssePayload = `data: ${jsonPayload}\n\n`;
+
+    const stream = new Readable();
+    const chunkSize = 1024;
+    for (let i = 0; i < ssePayload.length; i += chunkSize) {
+      stream.push(ssePayload.substring(i, i + chunkSize));
+    }
+    stream.push(null);
+
+    const readableStream = new ReadableStream({
+      start(controller) {
+        stream.on('data', (chunk) =>
+          controller.enqueue(new TextEncoder().encode(chunk)),
+        );
+        stream.on('end', () => controller.close());
+      },
+    });
+    const response = new Response(readableStream);
+
+    const generator = apiClient.processStreamResponse(response);
+    const result = await generator.next();
+    expect(result.done).toBeFalse();
+    const json = await result.value.json();
+    expect(json.data).toBe(largeData);
+
+    const final = await generator.next();
+    expect(final.done).toBeTrue();
+  });
 });
 
 describe('ApiClient', () => {

test/unit/interactions_test.ts

@@ -196,4 +196,56 @@ describe('Interactions resource', () => {
       expect(headers.get('x-goog-api-key')).toBe('some-manual-key');
     });
   });
+
+  describe('streaming regression', () => {
+    let client: GeminiNextGenAPIClient;
+    let fetchSpy: jasmine.Spy<Fetch>;
+
+    beforeEach(() => {
+      client = new GeminiNextGenAPIClient({
+        clientAdapter,
+        apiKey: 'test-api-key',
+      });
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      fetchSpy = spyOn(client as any, 'fetch');
+    });
+
+    it('should handle large fragmented SSE payloads correctly', async () => {
+      const largeData = 'A'.repeat(10 * 1024);
+      const mockSSE =
+        `data: {"event_type": "content.delta", "delta": {"type": "text", "text": "${
+          largeData
+        }"}}\n\n` + `data: [DONE]\n\n`;
+      const sseBytes = new TextEncoder().encode(mockSSE);
+
+      const readableStream = new ReadableStream({
+        start(controller) {
+          const chunkSize = 1024;
+          for (let i = 0; i < sseBytes.length; i += chunkSize) {
+            controller.enqueue(sseBytes.subarray(i, i + chunkSize));
+          }
+          controller.close();
+        },
+      });
+
+      fetchSpy.and.resolveTo(new Response(readableStream));
+
+      const stream = await client.interactions.create({
+        model: 'gemini-pro',
+        input: 'test',
+        stream: true,
+      });
+
+      let received = '';
+      for await (const event of stream) {
+        if (
+          event.event_type === 'content.delta' &&
+          event.delta?.type === 'text'
+        ) {
+          received += event.delta.text;
+        }
+      }
+      expect(received).toBe(largeData);
+    });
+  });
 });

test/unit/models_test.ts

@@ -1402,3 +1402,51 @@ describe('generateContentStream', () => {
     }
   });
 });
+
+describe('generateContentStream regression', () => {
+  it('should handle large fragmented SSE payloads correctly', async () => {
+    const client = new GoogleGenAI({vertexai: false, apiKey: 'fake-api-key'});
+
+    const largeText = 'A'.repeat(10 * 1024);
+    const mockResponse = {
+      candidates: [
+        {
+          content: {
+            parts: [{text: largeText}],
+            role: 'model',
+          },
+        },
+      ],
+    };
+
+    const sseData = `data: ${JSON.stringify(mockResponse)}\n\n`;
+    const encoder = new TextEncoder();
+    const sseBytes = encoder.encode(sseData);
+
+    const readableStream = new ReadableStream({
+      start(controller) {
+        const chunkSize = 1024;
+        for (let i = 0; i < sseBytes.length; i += chunkSize) {
+          controller.enqueue(sseBytes.subarray(i, i + chunkSize));
+        }
+        controller.close();
+      },
+    });
+
+    spyOn(global, 'fetch').and.resolveTo(
+      new Response(readableStream, fetchOkOptions),
+    );
+
+    const stream = await client.models.generateContentStream({
+      model: 'gemini-pro',
+      contents: 'test',
+    });
+
+    let receivedText = '';
+    for await (const chunk of stream) {
+      receivedText += chunk.candidates?.[0]?.content?.parts?.[0]?.text ?? '';
+    }
+
+    expect(receivedText).toBe(largeText);
+  });
+});