fix race condition in createInjectionTransformStream (#533)

phryneas · Copilot · web-flow · commit 0385d57d184c · 2025-12-18T11:13:20.000+01:00
* fix race condition in `createInjectionTransformStream`

* specify test title

* Update packages/client-react-streaming/src/stream-utils/createInjectionTransformStream.test.tsx

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;

* increase timing for flaky tests

---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/.changeset/quick-numbers-chew.md b/.changeset/quick-numbers-chew.md
@@ -0,0 +1,5 @@
+---
+"@apollo/client-react-streaming": patch
+---
+
+fix race condition in `createInjectionTransformStream`
diff --git a/packages/client-react-streaming/src/stream-utils/createInjectionTransformStream.test.tsx b/packages/client-react-streaming/src/stream-utils/createInjectionTransformStream.test.tsx
@@ -87,7 +87,7 @@ test("if the head has been flushed, just insert wherever", async () => {
   const { resultPromise, controller } = setupStream(transformStream);
 
   controller.enqueue("<html><head><title>Test</title></head><body>");
-  await scheduler.wait(1);
+  await scheduler.wait(10);
   injectIntoStream(() => <TestComponent />);
   controller.enqueue("</body></html>");
   controller.close();
@@ -107,6 +107,36 @@ test("if the head has been flushed, just insert wherever", async () => {
   );
 });
 
+test("race condition: injection queued while currentlyStreaming is true, no async follow-up chunks", async () => {
+  const { transformStream, injectIntoStream } =
+    createInjectionTransformStream();
+  const { resultPromise, controller } = setupStream(transformStream);
+
+  controller.enqueue("<html><head><title>Test</title></head><body>");
+  await scheduler.wait(10);
+  controller.enqueue("<di");
+  injectIntoStream(() => <TestComponent />);
+  controller.enqueue("v>unrelated content</div>");
+  controller.enqueue("</body></html>");
+  // TODO: it would be better if the script chunk ended up inside the `html` tag, but right now it's better here than nowhere. Browsers will be okay with it.
+  controller.close();
+
+  assert.equal(
+    await resultPromise,
+    html`<html>
+        <head>
+          <title>Test</title>
+        </head>
+        <body>
+          <div>unrelated content</div>
+        </body>
+      </html>
+      <script>
+        console.log("test");
+      </script>`
+  );
+});
+
 test("if the head is chopped into pieces, take the last chunk into account when looking for the end of the head", async () => {
   const { transformStream, injectIntoStream } =
     createInjectionTransformStream();
diff --git a/packages/client-react-streaming/src/stream-utils/createInjectionTransformStream.tsx b/packages/client-react-streaming/src/stream-utils/createInjectionTransformStream.tsx
@@ -83,7 +83,9 @@ export function createInjectionTransformStream(): {
 
   const transformStream = new TransformStream({
     async transform(chunk, controller) {
-      // While react is flushing chunks, we don't apply insertions
+      // While react is flushing chunks, we don't apply insertions.
+      // Injecting during streaming could break HTML by inserting content
+      // at arbitrary chunk boundaries (e.g., inside partial tags like <di${injected}v>).
       if (currentlyStreaming) {
         controller.enqueue(chunk);
         return;
@@ -109,14 +111,21 @@ export function createInjectionTransformStream(): {
           controller.enqueue(textEncoder.encode(content.slice(0, -KEEP_BYTES)));
         }
       } else {
-        controller.enqueue(textEncoder.encode(await renderInjectedHtml()));
+        if (queuedInjections.length > 0) {
+          controller.enqueue(textEncoder.encode(await renderInjectedHtml()));
+        }
         controller.enqueue(chunk);
         currentlyStreaming = true;
         setImmediate(() => {
           currentlyStreaming = false;
         });
       }
     },
+    async flush(controller) {
+      if (queuedInjections.length > 0) {
+        controller.enqueue(textEncoder.encode(await renderInjectedHtml()));
+      }
+    },
   });
 
   return {

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"@apollo/client-react-streaming": patch
 +---
++
 +fix race condition in `createInjectionTransformStream`