ai/core: allow reading the stream from streamText multiple times. (#1510)

Author: lgrammel

.changeset/late-toys-perform.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ai/core): allow reading streams in streamText result multiple times

packages/core/core/generate-text/stream-text.test.ts

@@ -439,3 +439,56 @@ describe('result.toTextStreamResponse', () => {
     assert.deepStrictEqual(chunks, ['Hello', ', ', 'world!']);
   });
 });
+
+describe('multiple stream consumption', () => {
+  it('should support text stream, ai stream, full stream on single result object', async () => {
+    const result = await streamText({
+      model: new MockLanguageModelV1({
+        doStream: async () => {
+          return {
+            stream: convertArrayToReadableStream([
+              { type: 'text-delta', textDelta: 'Hello' },
+              { type: 'text-delta', textDelta: ', ' },
+              { type: 'text-delta', textDelta: 'world!' },
+              {
+                type: 'finish',
+                finishReason: 'stop',
+                logprobs: undefined,
+                usage: { completionTokens: 10, promptTokens: 3 },
+              },
+            ]),
+            rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+          };
+        },
+      }),
+      prompt: 'test-input',
+    });
+
+    assert.deepStrictEqual(
+      await convertAsyncIterableToArray(result.textStream),
+      ['Hello', ', ', 'world!'],
+    );
+
+    assert.deepStrictEqual(
+      await convertReadableStreamToArray(
+        result.toAIStream().pipeThrough(new TextDecoderStream()),
+      ),
+      ['0:"Hello"\n', '0:", "\n', '0:"world!"\n'],
+    );
+
+    assert.deepStrictEqual(
+      await convertAsyncIterableToArray(result.fullStream),
+      [
+        { type: 'text-delta', textDelta: 'Hello' },
+        { type: 'text-delta', textDelta: ', ' },
+        { type: 'text-delta', textDelta: 'world!' },
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3, totalTokens: 13 },
+        },
+      ],
+    );
+  });
+});

packages/core/core/generate-text/stream-text.ts

@@ -140,7 +140,7 @@ export type TextStreamPart<TOOLS extends Record<string, CoreTool>> =
 A result object for accessing different stream types and additional information.
  */
 export class StreamTextResult<TOOLS extends Record<string, CoreTool>> {
-  private readonly originalStream: ReadableStream<TextStreamPart<TOOLS>>;
+  private originalStream: ReadableStream<TextStreamPart<TOOLS>>;
 
   /**
 Warnings from the model provider (e.g. unsupported settings)
@@ -173,13 +173,27 @@ Response headers.
     this.rawResponse = rawResponse;
   }
 
+  /**
+Split out a new stream from the original stream.
+The original stream is replaced to allow for further splitting,
+since we do not know how many times the stream will be split.
+
+Note: this leads to buffering the stream content on the server.
+However, the LLM results are expected to be small enough to not cause issues.
+   */
+  private teeStream() {
+    const [stream1, stream2] = this.originalStream.tee();
+    this.originalStream = stream2;
+    return stream1;
+  }
+
   /**
 A text stream that returns only the generated text deltas. You can use it
 as either an AsyncIterable or a ReadableStream. When an error occurs, the
 stream will throw the error.
    */
   get textStream(): AsyncIterableStream<string> {
-    return createAsyncIterableStream(this.originalStream, {
+    return createAsyncIterableStream(this.teeStream(), {
       transform(chunk, controller) {
         if (chunk.type === 'text-delta') {
           // do not stream empty text deltas:
@@ -200,7 +214,7 @@ You can use it as either an AsyncIterable or a ReadableStream. When an error occ
 stream will throw the error.
    */
   get fullStream(): AsyncIterableStream<TextStreamPart<TOOLS>> {
-    return createAsyncIterableStream(this.originalStream, {
+    return createAsyncIterableStream(this.teeStream(), {
       transform(chunk, controller) {
         if (chunk.type === 'text-delta') {
           // do not stream empty text deltas: