From ff0571007046050fe15a0bc89d3afa0a56921a40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sema@fb.com>
Date: Sat, 27 Mar 2021 13:50:38 -0400
Subject: [PATCH] [Fizz] Support special HTML/SVG/MathML tags to suspend
 (#21113)

* Encode tables as a special insertion mode

The table modes are special in that its children can't be created outside
a table context so we need the segment container to be wrapped in a table.

* Move formatContext from Task to Segment

It works the same otherwise. It's just that this context needs to outlive
the task so that I can use it when writing the segment.

* Use template tag for placeholders and inserted dummy nodes with IDs

These can be used in any parent. At least outside IE11. Not sure yet what
happens in IE11 to these.

Not sure if these are bad for perf since they're special nodes.

* Add special wrappers around inserted segments depending on their insertion mode

* Allow the root namespace to be configured

This allows us to insert the correct wrappers when streaming into an
existing non-HTML tree.

* Add comment
---
 .../src/__tests__/ReactDOMFizzServer-test.js  | 249 +++++++++++++++++-
 .../src/server/ReactDOMFizzServerBrowser.js   |   3 +-
 .../src/server/ReactDOMFizzServerNode.js      |   3 +-
 .../src/server/ReactDOMServerFormatConfig.js  | 220 +++++++++++++---
 .../server/ReactNativeServerFormatConfig.js   |   6 +-
 .../src/ReactNoopServer.js                    |   3 +-
 packages/react-server/src/ReactFizzServer.js  |  47 ++--
 scripts/error-codes/codes.json                |   3 +-
 8 files changed, 475 insertions(+), 59 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index 9021c23c6ce74..31b1f437c9871 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -97,10 +97,22 @@ describe('ReactDOMFizzServer', () => {
     let node = element.firstChild;
     while (node) {
       if (node.nodeType === 1) {
-        if (node.tagName !== 'SCRIPT' && !node.hasAttribute('hidden')) {
+        if (
+          node.tagName !== 'SCRIPT' &&
+          node.tagName !== 'TEMPLATE' &&
+          !node.hasAttribute('hidden') &&
+          !node.hasAttribute('aria-hidden')
+        ) {
           const props = {};
           const attributes = node.attributes;
           for (let i = 0; i < attributes.length; i++) {
+            if (
+              attributes[i].name === 'id' &&
+              attributes[i].value.includes(':')
+            ) {
+              // We assume this is a React added ID that's a non-visual implementation detail.
+              continue;
+            }
             props[attributes[i].name] = attributes[i].value;
           }
           props.children = getVisibleChildren(node);
@@ -112,7 +124,7 @@ describe('ReactDOMFizzServer', () => {
       node = node.nextSibling;
     }
     return children.length === 0
-      ? null
+      ? undefined
       : children.length === 1
       ? children[0]
       : children;
@@ -408,4 +420,237 @@ describe('ReactDOMFizzServer', () => {
       </div>,
     ]);
   });
+
+  // @gate experimental
+  it('can resolve async content in esoteric parents', async () => {
+    function AsyncOption({text}) {
+      return <option>{readText(text)}</option>;
+    }
+
+    function AsyncCol({className}) {
+      return <col className={readText(className)}>{[]}</col>;
+    }
+
+    function AsyncPath({id}) {
+      return <path id={readText(id)}>{[]}</path>;
+    }
+
+    function AsyncMi({id}) {
+      return <mi id={readText(id)}>{[]}</mi>;
+    }
+
+    function App() {
+      return (
+        <div>
+          <select>
+            <Suspense fallback="Loading...">
+              <AsyncOption text="Hello" />
+            </Suspense>
+          </select>
+          <Suspense fallback="Loading...">
+            <table>
+              <colgroup>
+                <AsyncCol className="World" />
+              </colgroup>
+            </table>
+            <svg>
+              <g>
+                <AsyncPath id="my-path" />
+              </g>
+            </svg>
+            <math>
+              <AsyncMi id="my-mi" />
+            </math>
+          </Suspense>
+        </div>
+      );
+    }
+
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <App />,
+        writable,
+      );
+      startWriting();
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <select>Loading...</select>Loading...
+      </div>,
+    );
+
+    await act(async () => {
+      resolveText('Hello');
+    });
+
+    await act(async () => {
+      resolveText('World');
+    });
+
+    await act(async () => {
+      resolveText('my-path');
+      resolveText('my-mi');
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <select>
+          <option>Hello</option>
+        </select>
+        <table>
+          <colgroup>
+            <col class="World" />
+          </colgroup>
+        </table>
+        <svg>
+          <g>
+            <path id="my-path" />
+          </g>
+        </svg>
+        <math>
+          <mi id="my-mi" />
+        </math>
+      </div>,
+    );
+
+    expect(container.querySelector('#my-path').namespaceURI).toBe(
+      'http://www.w3.org/2000/svg',
+    );
+    expect(container.querySelector('#my-mi').namespaceURI).toBe(
+      'http://www.w3.org/1998/Math/MathML',
+    );
+  });
+
+  // @gate experimental
+  it('can resolve async content in table parents', async () => {
+    function AsyncTableBody({className, children}) {
+      return <tbody className={readText(className)}>{children}</tbody>;
+    }
+
+    function AsyncTableRow({className, children}) {
+      return <tr className={readText(className)}>{children}</tr>;
+    }
+
+    function AsyncTableCell({text}) {
+      return <td>{readText(text)}</td>;
+    }
+
+    function App() {
+      return (
+        <table>
+          <Suspense
+            fallback={
+              <tbody>
+                <tr>
+                  <td>Loading...</td>
+                </tr>
+              </tbody>
+            }>
+            <AsyncTableBody className="A">
+              <AsyncTableRow className="B">
+                <AsyncTableCell text="C" />
+              </AsyncTableRow>
+            </AsyncTableBody>
+          </Suspense>
+        </table>
+      );
+    }
+
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <App />,
+        writable,
+      );
+      startWriting();
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <table>
+        <tbody>
+          <tr>
+            <td>Loading...</td>
+          </tr>
+        </tbody>
+      </table>,
+    );
+
+    await act(async () => {
+      resolveText('A');
+    });
+
+    await act(async () => {
+      resolveText('B');
+    });
+
+    await act(async () => {
+      resolveText('C');
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <table>
+        <tbody class="A">
+          <tr class="B">
+            <td>C</td>
+          </tr>
+        </tbody>
+      </table>,
+    );
+  });
+
+  // @gate experimental
+  it('can stream into an SVG container', async () => {
+    function AsyncPath({id}) {
+      return <path id={readText(id)}>{[]}</path>;
+    }
+
+    function App() {
+      return (
+        <g>
+          <Suspense fallback={<text>Loading...</text>}>
+            <AsyncPath id="my-path" />
+          </Suspense>
+        </g>
+      );
+    }
+
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <App />,
+        writable,
+        {
+          namespaceURI: 'http://www.w3.org/2000/svg',
+          onReadyToStream() {
+            writable.write('<svg>');
+            startWriting();
+            writable.write('</svg>');
+          },
+        },
+      );
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <svg>
+        <g>
+          <text>Loading...</text>
+        </g>
+      </svg>,
+    );
+
+    await act(async () => {
+      resolveText('my-path');
+    });
+
+    expect(getVisibleChildren(container)).toEqual(
+      <svg>
+        <g>
+          <path id="my-path" />
+        </g>
+      </svg>,
+    );
+
+    expect(container.querySelector('#my-path').namespaceURI).toBe(
+      'http://www.w3.org/2000/svg',
+    );
+  });
 });
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
index 254d2380d5ab2..8d5c59a2d6862 100644
--- a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
+++ b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
@@ -23,6 +23,7 @@ import {
 
 type Options = {
   identifierPrefix?: string,
+  namespaceURI?: string,
   progressiveChunkSize?: number,
   signal?: AbortSignal,
   onReadyToStream?: () => void,
@@ -49,7 +50,7 @@ function renderToReadableStream(
         children,
         controller,
         createResponseState(options ? options.identifierPrefix : undefined),
-        createRootFormatContext(), // We call this here in case we need options to initialize it.
+        createRootFormatContext(options ? options.namespaceURI : undefined),
         options ? options.progressiveChunkSize : undefined,
         options ? options.onError : undefined,
         options ? options.onCompleteAll : undefined,
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerNode.js b/packages/react-dom/src/server/ReactDOMFizzServerNode.js
index c9d0cbc303148..1d53e6fda472e 100644
--- a/packages/react-dom/src/server/ReactDOMFizzServerNode.js
+++ b/packages/react-dom/src/server/ReactDOMFizzServerNode.js
@@ -28,6 +28,7 @@ function createDrainHandler(destination, request) {
 
 type Options = {
   identifierPrefix?: string,
+  namespaceURI?: string,
   progressiveChunkSize?: number,
   onReadyToStream?: () => void,
   onCompleteAll?: () => void,
@@ -49,7 +50,7 @@ function pipeToNodeWritable(
     children,
     destination,
     createResponseState(options ? options.identifierPrefix : undefined),
-    createRootFormatContext(), // We call this here in case we need options to initialize it.
+    createRootFormatContext(options ? options.namespaceURI : undefined),
     options ? options.progressiveChunkSize : undefined,
     options ? options.onError : undefined,
     options ? options.onCompleteAll : undefined,
diff --git a/packages/react-dom/src/server/ReactDOMServerFormatConfig.js b/packages/react-dom/src/server/ReactDOMServerFormatConfig.js
index f7ed7ccf93971..d42e92aa1e2b4 100644
--- a/packages/react-dom/src/server/ReactDOMServerFormatConfig.js
+++ b/packages/react-dom/src/server/ReactDOMServerFormatConfig.js
@@ -50,33 +50,45 @@ export function createResponseState(
   };
 }
 
-// Constants for the namespace we use. We don't actually provide the namespace but conditionally
-// use different segment parents based on namespace. Therefore we use constants instead of the string.
-const ROOT_NAMESPACE = 0; // At the root we don't need to know which namespace it is. We just need to know that it's already the right one.
-const HTML_NAMESPACE = 1;
-const SVG_NAMESPACE = 2;
-const MATHML_NAMESPACE = 3;
-
-type NamespaceFlag = 0 | 1 | 2 | 3;
+// Constants for the insertion mode we're currently writing in. We don't encode all HTML5 insertion
+// modes. We only include the variants as they matter for the sake of our purposes.
+// We don't actually provide the namespace therefore we use constants instead of the string.
+const HTML_MODE = 0;
+const SVG_MODE = 1;
+const MATHML_MODE = 2;
+const HTML_TABLE_MODE = 4;
+const HTML_TABLE_BODY_MODE = 5;
+const HTML_TABLE_ROW_MODE = 6;
+const HTML_COLGROUP_MODE = 7;
+// We have a greater than HTML_TABLE_MODE check elsewhere. If you add more cases here, make sure it
+// still makes sense
+
+type InsertionMode = 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7;
 
 // Lets us keep track of contextual state and pick it back up after suspending.
 export type FormatContext = {
-  namespace: NamespaceFlag, // root/svg/html/mathml
+  insertionMode: InsertionMode, // root/svg/html/mathml/table
   selectedValue: null | string, // the selected value(s) inside a <select>, or null outside <select>
 };
 
 function createFormatContext(
-  namespace: NamespaceFlag,
+  insertionMode: InsertionMode,
   selectedValue: null | string,
 ): FormatContext {
   return {
-    namespace,
+    insertionMode,
     selectedValue,
   };
 }
 
-export function createRootFormatContext(): FormatContext {
-  return createFormatContext(ROOT_NAMESPACE, null);
+export function createRootFormatContext(namespaceURI?: string): FormatContext {
+  const insertionMode =
+    namespaceURI === 'http://www.w3.org/2000/svg'
+      ? SVG_MODE
+      : namespaceURI === 'http://www.w3.org/1998/Math/MathML'
+      ? MATHML_MODE
+      : HTML_MODE;
+  return createFormatContext(insertionMode, null);
 }
 
 export function getChildFormatContext(
@@ -87,15 +99,32 @@ export function getChildFormatContext(
   switch (type) {
     case 'select':
       return createFormatContext(
-        parentContext.namespace,
+        HTML_MODE,
         props.value != null ? props.value : props.defaultValue,
       );
     case 'svg':
-      return createFormatContext(SVG_NAMESPACE, null);
+      return createFormatContext(SVG_MODE, null);
     case 'math':
-      return createFormatContext(MATHML_NAMESPACE, null);
+      return createFormatContext(MATHML_MODE, null);
     case 'foreignObject':
-      return createFormatContext(HTML_NAMESPACE, null);
+      return createFormatContext(HTML_MODE, null);
+    // Table parents are special in that their children can only be created at all if they're
+    // wrapped in a table parent. So we need to encode that we're entering this mode.
+    case 'table':
+      return createFormatContext(HTML_TABLE_MODE, null);
+    case 'thead':
+    case 'tbody':
+    case 'tfoot':
+      return createFormatContext(HTML_TABLE_BODY_MODE, null);
+    case 'colgroup':
+      return createFormatContext(HTML_COLGROUP_MODE, null);
+    case 'tr':
+      return createFormatContext(HTML_TABLE_ROW_MODE, null);
+  }
+  if (parentContext.insertionMode >= HTML_TABLE_MODE) {
+    // Whatever tag this was, it wasn't a table parent or other special parent, so we must have
+    // entered plain HTML again.
+    return createFormatContext(HTML_MODE, null);
   }
   return parentContext;
 }
@@ -132,8 +161,8 @@ function assignAnID(
   ));
 }
 
-const dummyNode1 = stringToPrecomputedChunk('<span hidden id="');
-const dummyNode2 = stringToPrecomputedChunk('"></span>');
+const dummyNode1 = stringToPrecomputedChunk('<template id="');
+const dummyNode2 = stringToPrecomputedChunk('"></template>');
 
 function pushDummyNodeWithID(
   target: Array<Chunk | PrecomputedChunk>,
@@ -206,7 +235,20 @@ export function pushStartInstance(
       startTag2,
     );
   } else {
-    target.push(startTag1, stringToChunk(type), startTag2);
+    target.push(startTag1, stringToChunk(type));
+    if (props.className) {
+      target.push(
+        stringToChunk(
+          ' class="' + encodeHTMLIDAttribute(props.className) + '"',
+        ),
+      );
+    }
+    if (props.id) {
+      target.push(
+        stringToChunk(' id="' + encodeHTMLIDAttribute(props.id) + '"'),
+      );
+    }
+    target.push(startTag2);
   }
 }
 
@@ -225,16 +267,15 @@ export function pushEndInstance(
 // Structural Nodes
 
 // A placeholder is a node inside a hidden partial tree that can be filled in later, but before
-// display. It's never visible to users.
-const placeholder1 = stringToPrecomputedChunk('<span id="');
-const placeholder2 = stringToPrecomputedChunk('"></span>');
+// display. It's never visible to users. We use the template tag because it can be used in every
+// type of parent. <script> tags also work in every other tag except <colgroup>.
+const placeholder1 = stringToPrecomputedChunk('<template id="');
+const placeholder2 = stringToPrecomputedChunk('"></template>');
 export function writePlaceholder(
   destination: Destination,
   responseState: ResponseState,
   id: number,
 ): boolean {
-  // TODO: This needs to be contextually aware and switch tag since not all parents allow for spans like
-  // <select> or <tbody>. E.g. suspending a component that renders a table row.
   writeChunk(destination, placeholder1);
   writeChunk(destination, responseState.placeholderPrefix);
   const formattedID = stringToChunk(id.toString(16));
@@ -272,23 +313,130 @@ export function writeEndSuspenseBoundary(destination: Destination): boolean {
   return writeChunk(destination, endSuspenseBoundary);
 }
 
-const startSegment = stringToPrecomputedChunk('<div hidden id="');
-const startSegment2 = stringToPrecomputedChunk('">');
-const endSegment = stringToPrecomputedChunk('</div>');
+const startSegmentHTML = stringToPrecomputedChunk('<div hidden id="');
+const startSegmentHTML2 = stringToPrecomputedChunk('">');
+const endSegmentHTML = stringToPrecomputedChunk('</div>');
+
+const startSegmentSVG = stringToPrecomputedChunk(
+  '<svg aria-hidden="true" style="display:none" id="',
+);
+const startSegmentSVG2 = stringToPrecomputedChunk('">');
+const endSegmentSVG = stringToPrecomputedChunk('</svg>');
+
+const startSegmentMathML = stringToPrecomputedChunk(
+  '<math aria-hidden="true" style="display:none" id="',
+);
+const startSegmentMathML2 = stringToPrecomputedChunk('">');
+const endSegmentMathML = stringToPrecomputedChunk('</math>');
+
+const startSegmentTable = stringToPrecomputedChunk('<table hidden id="');
+const startSegmentTable2 = stringToPrecomputedChunk('">');
+const endSegmentTable = stringToPrecomputedChunk('</table>');
+
+const startSegmentTableBody = stringToPrecomputedChunk(
+  '<table hidden><tbody id="',
+);
+const startSegmentTableBody2 = stringToPrecomputedChunk('">');
+const endSegmentTableBody = stringToPrecomputedChunk('</tbody></table>');
+
+const startSegmentTableRow = stringToPrecomputedChunk('<table hidden><tr id="');
+const startSegmentTableRow2 = stringToPrecomputedChunk('">');
+const endSegmentTableRow = stringToPrecomputedChunk('</tr></table>');
+
+const startSegmentColGroup = stringToPrecomputedChunk(
+  '<table hidden><colgroup id="',
+);
+const startSegmentColGroup2 = stringToPrecomputedChunk('">');
+const endSegmentColGroup = stringToPrecomputedChunk('</colgroup></table>');
+
 export function writeStartSegment(
   destination: Destination,
   responseState: ResponseState,
+  formatContext: FormatContext,
   id: number,
 ): boolean {
-  // TODO: What happens with special children like <tr> if they're inserted in a div? Maybe needs contextually aware containers.
-  writeChunk(destination, startSegment);
-  writeChunk(destination, responseState.segmentPrefix);
-  const formattedID = stringToChunk(id.toString(16));
-  writeChunk(destination, formattedID);
-  return writeChunk(destination, startSegment2);
+  switch (formatContext.insertionMode) {
+    case HTML_MODE: {
+      writeChunk(destination, startSegmentHTML);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentHTML2);
+    }
+    case SVG_MODE: {
+      writeChunk(destination, startSegmentSVG);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentSVG2);
+    }
+    case MATHML_MODE: {
+      writeChunk(destination, startSegmentMathML);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentMathML2);
+    }
+    case HTML_TABLE_MODE: {
+      writeChunk(destination, startSegmentTable);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentTable2);
+    }
+    // TODO: For the rest of these, there will be extra wrapper nodes that never
+    // get deleted from the document. We need to delete the table too as part
+    // of the injected scripts. They are invisible though so it's not too terrible
+    // and it's kind of an edge case to suspend in a table. Totally supported though.
+    case HTML_TABLE_BODY_MODE: {
+      writeChunk(destination, startSegmentTableBody);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentTableBody2);
+    }
+    case HTML_TABLE_ROW_MODE: {
+      writeChunk(destination, startSegmentTableRow);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentTableRow2);
+    }
+    case HTML_COLGROUP_MODE: {
+      writeChunk(destination, startSegmentColGroup);
+      writeChunk(destination, responseState.segmentPrefix);
+      writeChunk(destination, stringToChunk(id.toString(16)));
+      return writeChunk(destination, startSegmentColGroup2);
+    }
+    default: {
+      invariant(false, 'Unknown insertion mode. This is a bug in React.');
+    }
+  }
 }
-export function writeEndSegment(destination: Destination): boolean {
-  return writeChunk(destination, endSegment);
+export function writeEndSegment(
+  destination: Destination,
+  formatContext: FormatContext,
+): boolean {
+  switch (formatContext.insertionMode) {
+    case HTML_MODE: {
+      return writeChunk(destination, endSegmentHTML);
+    }
+    case SVG_MODE: {
+      return writeChunk(destination, endSegmentSVG);
+    }
+    case MATHML_MODE: {
+      return writeChunk(destination, endSegmentMathML);
+    }
+    case HTML_TABLE_MODE: {
+      return writeChunk(destination, endSegmentTable);
+    }
+    case HTML_TABLE_BODY_MODE: {
+      return writeChunk(destination, endSegmentTableBody);
+    }
+    case HTML_TABLE_ROW_MODE: {
+      return writeChunk(destination, endSegmentTableRow);
+    }
+    case HTML_COLGROUP_MODE: {
+      return writeChunk(destination, endSegmentColGroup);
+    }
+    default: {
+      invariant(false, 'Unknown insertion mode. This is a bug in React.');
+    }
+  }
 }
 
 // Instruction Set
diff --git a/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js b/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js
index 3e6b37e76c18a..102ee87f05d5a 100644
--- a/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js
+++ b/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js
@@ -208,12 +208,16 @@ export function writeEndSuspenseBoundary(destination: Destination): boolean {
 export function writeStartSegment(
   destination: Destination,
   responseState: ResponseState,
+  formatContext: FormatContext,
   id: number,
 ): boolean {
   writeChunk(destination, SEGMENT);
   return writeChunk(destination, formatID(id));
 }
-export function writeEndSegment(destination: Destination): boolean {
+export function writeEndSegment(
+  destination: Destination,
+  formatContext: FormatContext,
+): boolean {
   return writeChunk(destination, END);
 }
 
diff --git a/packages/react-noop-renderer/src/ReactNoopServer.js b/packages/react-noop-renderer/src/ReactNoopServer.js
index f384baa8d9f9e..548a004bd7aeb 100644
--- a/packages/react-noop-renderer/src/ReactNoopServer.js
+++ b/packages/react-noop-renderer/src/ReactNoopServer.js
@@ -161,6 +161,7 @@ const ReactNoopServer = ReactFizzServer({
   writeStartSegment(
     destination: Destination,
     responseState: ResponseState,
+    formatContext: null,
     id: number,
   ): boolean {
     const segment = {
@@ -172,7 +173,7 @@ const ReactNoopServer = ReactFizzServer({
     }
     destination.stack.push(segment);
   },
-  writeEndSegment(destination: Destination): boolean {
+  writeEndSegment(destination: Destination, formatContext: null): boolean {
     destination.stack.pop();
   },
 
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index c4b795d24ff17..739ba0c436aa6 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -71,7 +71,6 @@ type Task = {
   blockedBoundary: Root | SuspenseBoundary,
   blockedSegment: Segment, // the segment we'll write to
   abortSet: Set<Task>, // the abortable set that this task belongs to
-  formatContext: FormatContext,
   assignID: null | SuspenseBoundaryID, // id to assign to the content
 };
 
@@ -90,6 +89,8 @@ type Segment = {
   +index: number, // the index within the parent's chunks or 0 at the root
   +chunks: Array<Chunk | PrecomputedChunk>,
   +children: Array<Segment>,
+  // The context that this segment was created in.
+  formatContext: FormatContext,
   // If this segment represents a fallback, this is the content that will replace that fallback.
   +boundary: null | SuspenseBoundary,
 };
@@ -172,7 +173,7 @@ export function createRequest(
     onReadyToStream,
   };
   // This segment represents the root fallback.
-  const rootSegment = createPendingSegment(request, 0, null);
+  const rootSegment = createPendingSegment(request, 0, null, rootContext);
   // There is no parent so conceptually, we're unblocked to flush this segment.
   rootSegment.parentFlushed = true;
   const rootTask = createTask(
@@ -181,7 +182,6 @@ export function createRequest(
     null,
     rootSegment,
     abortSet,
-    rootContext,
     null,
   );
   pingedTasks.push(rootTask);
@@ -218,7 +218,6 @@ function createTask(
   blockedBoundary: Root | SuspenseBoundary,
   blockedSegment: Segment,
   abortSet: Set<Task>,
-  formatContext: FormatContext,
   assignID: null | SuspenseBoundaryID,
 ): Task {
   request.allPendingTasks++;
@@ -233,7 +232,6 @@ function createTask(
     blockedBoundary,
     blockedSegment,
     abortSet,
-    formatContext,
     assignID,
   };
   abortSet.add(task);
@@ -244,6 +242,7 @@ function createPendingSegment(
   request: Request,
   index: number,
   boundary: null | SuspenseBoundary,
+  formatContext: FormatContext,
 ): Segment {
   return {
     status: PENDING,
@@ -252,6 +251,7 @@ function createPendingSegment(
     parentFlushed: false,
     chunks: [],
     children: [],
+    formatContext,
     boundary,
   };
 }
@@ -317,7 +317,12 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
         // Something suspended, we'll need to create a new segment and resolve it later.
         const segment = task.blockedSegment;
         const insertionIndex = segment.chunks.length;
-        const newSegment = createPendingSegment(request, insertionIndex, null);
+        const newSegment = createPendingSegment(
+          request,
+          insertionIndex,
+          null,
+          segment.formatContext,
+        );
         segment.children.push(newSegment);
         const newTask = createTask(
           request,
@@ -325,7 +330,6 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
           task.blockedBoundary,
           newSegment,
           task.abortSet,
-          task.formatContext,
           task.assignID,
         );
         // We've delegated the assignment.
@@ -338,8 +342,9 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
       }
     }
   } else if (typeof type === 'string') {
+    const segment = task.blockedSegment;
     pushStartInstance(
-      task.blockedSegment.chunks,
+      segment.chunks,
       type,
       props,
       request.responseState,
@@ -347,13 +352,13 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
     );
     // We must have assigned it already above so we don't need this anymore.
     task.assignID = null;
-    const prevContext = task.formatContext;
-    task.formatContext = getChildFormatContext(prevContext, type, props);
+    const prevContext = segment.formatContext;
+    segment.formatContext = getChildFormatContext(prevContext, type, props);
     renderNode(request, task, props.children);
     // We expect that errors will fatal the whole task and that we don't need
     // the correct context. Therefore this is not in a finally.
-    task.formatContext = prevContext;
-    pushEndInstance(task.blockedSegment.chunks, type, props);
+    segment.formatContext = prevContext;
+    pushEndInstance(segment.chunks, type, props);
   } else if (type === REACT_SUSPENSE_TYPE) {
     const parentBoundary = task.blockedBoundary;
     const parentSegment = task.blockedSegment;
@@ -376,11 +381,17 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
       request,
       insertionIndex,
       newBoundary,
+      parentSegment.formatContext,
     );
     parentSegment.children.push(boundarySegment);
 
     // This segment is the actual child content. We can start rendering that immediately.
-    const contentRootSegment = createPendingSegment(request, 0, null);
+    const contentRootSegment = createPendingSegment(
+      request,
+      0,
+      null,
+      parentSegment.formatContext,
+    );
     // We mark the root segment as having its parent flushed. It's not really flushed but there is
     // no parent segment so there's nothing to wait on.
     contentRootSegment.parentFlushed = true;
@@ -425,7 +436,6 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
       parentBoundary,
       boundarySegment,
       fallbackAbortSet,
-      task.formatContext,
       newBoundary.id, // This is the ID we want to give this fallback so we can replace it later.
     );
     // TODO: This should be queued at a separate lower priority queue so that we only task
@@ -776,9 +786,14 @@ function flushSegmentContainer(
   destination: Destination,
   segment: Segment,
 ): boolean {
-  writeStartSegment(destination, request.responseState, segment.id);
+  writeStartSegment(
+    destination,
+    request.responseState,
+    segment.formatContext,
+    segment.id,
+  );
   flushSegment(request, destination, segment);
-  return writeEndSegment(destination);
+  return writeEndSegment(destination, segment.formatContext);
 }
 
 function flushCompletedBoundary(
diff --git a/scripts/error-codes/codes.json b/scripts/error-codes/codes.json
index 180106a9831c7..3c5fd25853f7c 100644
--- a/scripts/error-codes/codes.json
+++ b/scripts/error-codes/codes.json
@@ -384,5 +384,6 @@
   "393": "Cache cannot be refreshed during server rendering.",
   "394": "startTransition cannot be called during server rendering.",
   "395": "An ID must have been assigned before we can complete the boundary.",
-  "396": "More boundaries or placeholders than we expected to ever emit."
+  "396": "More boundaries or placeholders than we expected to ever emit.",
+  "397": "Unknown insertion mode. This is a bug in React."
 }