Fix SSR truncated response (#4196)

.changeset/lemon-ducks-smile.md

@@ -0,0 +1,5 @@
+---
+'@lit-labs/ssr': patch
+---
+
+Fix SSR returning a truncated response

packages/labs/ssr/src/lib/render-result-readable.ts

@@ -23,6 +23,7 @@ export class RenderResultReadable extends Readable {
    * completion in any one loop.
    */
   private _iterators: Array<RenderResultIterator>;
+  private _currentIterator?: RenderResultIterator;
 
   constructor(result: RenderResult) {
     super();
@@ -56,14 +57,15 @@ export class RenderResultReadable extends Readable {
     // result, because we must be able to return in the middle of the loop
     // and resume on the next call to _read().
 
-    // Get the current iterator
-    let iterator = this._iterators.pop();
+    // Get the current iterator, only if we don't already have one from the
+    // previous call to _read()
+    this._currentIterator ??= this._iterators.pop();
 
-    while (iterator !== undefined) {
-      const next = iterator.next();
+    while (this._currentIterator !== undefined) {
+      const next = this._currentIterator.next();
       if (next.done === true) {
         // Restore the outer iterator
-        iterator = this._iterators.pop();
+        this._currentIterator = this._iterators.pop();
         continue;
       }
 
@@ -77,8 +79,10 @@ export class RenderResultReadable extends Readable {
         }
       } else {
         // Must be a Promise
-        this._iterators.push(iterator);
-        iterator = (await value)[Symbol.iterator]() as RenderResultIterator;
+        this._iterators.push(this._currentIterator);
+        this._currentIterator = (await value)[
+          Symbol.iterator
+        ]() as RenderResultIterator;
       }
     }
     // Pushing `null` ends the stream

packages/labs/ssr/src/test/lib/render-result-readable_test.ts

@@ -40,6 +40,27 @@ test('RenderResultReadable collects strings and nested Promises of iterables', a
   assert.equal(s, 'abcd');
 });
 
+test('RenderResultReadable collects all iterables when stream is back pressured', async () => {
+  class TestRenderResultReadable extends RenderResultReadable {
+    override push(value: any) {
+      super.push(value);
+
+      // for the value different from "__" we return true to indicate that we can accept more data
+      return value !== '__';
+    }
+  }
+
+  const readable = new TestRenderResultReadable([
+    'a',
+    Promise.resolve([Promise.resolve(['__', 'b'])]),
+    '__',
+    'c',
+  ]);
+
+  const s = await collectReadable(readable);
+  assert.equal(s, 'a__b__c');
+});
+
 test.run();
 
 const collectReadable = async (r: Readable) => {