Merge pull request #7217 from Shopify/rcb/rendering-changes-react-fix-v2

Reapply: central unmount for ink render tree + fix render lifecycle race

Author: ryancbahan

packages/app/src/cli/services/dev/processes/theme-app-extension.test.ts

@@ -22,9 +22,18 @@ vi.mock('@shopify/cli-kit/node/themes/api')
 vi.mock('@shopify/cli-kit/node/context/fqdn')
 vi.mock('@shopify/cli-kit/node/ui', async (realImport) => {
   const realModule = await realImport<typeof import('@shopify/cli-kit/node/ui')>()
-  const mockModule = {renderInfo: vi.fn()}
 
-  return {...realModule, ...mockModule}
+  return {
+    ...realModule,
+    renderInfo: vi.fn(),
+    renderTasks: vi.fn(async (tasks: any[]) => {
+      for (const task of tasks) {
+        // eslint-disable-next-line no-await-in-loop
+        await task.task({}, task)
+      }
+      return {}
+    }),
+  }
 })
 
 describe('setupPreviewThemeAppExtensionsProcess', () => {

packages/cli-kit/src/private/node/testing/ui.ts

@@ -1,5 +1,5 @@
-import {Stdout} from '../ui.js'
-import {ReactElement} from 'react'
+import {Stdout, InkLifecycleRoot} from '../ui.js'
+import React, {ReactElement} from 'react'
 import {render as inkRender} from 'ink'
 
 import {EventEmitter} from 'events'
@@ -66,7 +66,7 @@ export const render = (tree: ReactElement, options: RenderOptions = {}): Instanc
   const stderr = new Stderr()
   const stdin = new Stdin()
 
-  const instance = inkRender(tree, {
+  const instance = inkRender(React.createElement(InkLifecycleRoot, null, tree), {
     stdout: options.stdout ?? (stdout as any),
 
     stderr: options.stderr ?? (stderr as any),
@@ -78,7 +78,7 @@ export const render = (tree: ReactElement, options: RenderOptions = {}): Instanc
   })
 
   return {
-    rerender: instance.rerender,
+    rerender: (tree: ReactElement) => instance.rerender(React.createElement(InkLifecycleRoot, null, tree)),
     unmount: instance.unmount,
     cleanup: instance.cleanup,
     waitUntilExit: () => trackPromise(instance.waitUntilExit().then(() => {})),

packages/cli-kit/src/private/node/ui.tsx

@@ -3,11 +3,49 @@ import {Logger, LogLevel} from '../../public/node/output.js'
 import {isUnitTest} from '../../public/node/context/local.js'
 import {treeKill} from '../../public/node/tree-kill.js'
 
-import {ReactElement} from 'react'
-import {Key, render as inkRender, RenderOptions} from 'ink'
+import React, {ReactElement, createContext, useCallback, useContext, useEffect, useState} from 'react'
+import {Key, render as inkRender, RenderOptions, useApp} from 'ink'
 
 import {EventEmitter} from 'events'
 
+const CompletionContext = createContext<((error?: Error) => void) | null>(null)
+
+/**
+ * Signal that the current Ink tree is done. Must be called within an
+ * InkLifecycleRoot — throws if the provider is missing so lifecycle
+ * bugs surface immediately instead of silently hanging.
+ */
+export function useComplete(): (error?: Error) => void {
+  const complete = useContext(CompletionContext)
+  if (!complete) {
+    throw new Error('useComplete() called outside InkLifecycleRoot')
+  }
+  return complete
+}
+
+/**
+ * Root wrapper for Ink trees. Owns the single `exit()` call site — children
+ * signal completion via `useComplete()`, which sets state here. The `useEffect`
+ * fires post-render, guaranteeing all batched state updates have been flushed
+ * before the tree is torn down.
+ */
+export function InkLifecycleRoot({children}: {children: React.ReactNode}) {
+  const {exit} = useApp()
+  const [exitResult, setExitResult] = useState<{error?: Error} | null>(null)
+
+  const complete = useCallback((error?: Error) => {
+    setExitResult({error})
+  }, [])
+
+  useEffect(() => {
+    if (exitResult !== null) {
+      exit(exitResult.error)
+    }
+  }, [exitResult, exit])
+
+  return <CompletionContext.Provider value={complete}>{children}</CompletionContext.Provider>
+}
+
 interface RenderOnceOptions {
   logLevel?: LogLevel
   logger?: Logger
@@ -27,10 +65,8 @@ export function renderOnce(element: JSX.Element, {logLevel = 'info', renderOptio
 }
 
 export async function render(element: JSX.Element, options?: RenderOptions) {
-  const {waitUntilExit} = inkRender(element, options)
+  const {waitUntilExit} = inkRender(<InkLifecycleRoot>{element}</InkLifecycleRoot>, options)
   await waitUntilExit()
-  // We need to wait for other pending tasks -- unmounting of the ink component -- to complete
-  return new Promise((resolve) => setImmediate(resolve))
 }
 
 interface Instance {

packages/cli-kit/src/private/node/ui/components/AutocompletePrompt.tsx

@@ -5,10 +5,11 @@ import {InfoMessageProps} from './Prompts/InfoMessage.js'
 import {Message, PromptLayout} from './Prompts/PromptLayout.js'
 import {throttle} from '../../../../public/common/function.js'
 import {AbortSignal} from '../../../../public/node/abort.js'
+import {useComplete} from '../../ui.js'
 import usePrompt, {PromptState} from '../hooks/use-prompt.js'
 
 import React, {ReactElement, useCallback, useEffect, useRef, useState} from 'react'
-import {Box, useApp} from 'ink'
+import {Box} from 'ink'
 
 export interface SearchResults<T> {
   data: SelectItem<T>[]
@@ -42,7 +43,7 @@ function AutocompletePrompt<T>({
   infoMessage,
   groupOrder,
 }: React.PropsWithChildren<AutocompletePromptProps<T>>): ReactElement | null {
-  const {exit: unmountInk} = useApp()
+  const complete = useComplete()
   const [searchTerm, setSearchTerm] = useState('')
   const [searchResults, setSearchResults] = useState<SelectItem<T>[]>(choices)
   const canSearch = choices.length > MIN_NUMBER_OF_ITEMS_FOR_SEARCH
@@ -72,10 +73,10 @@ function AutocompletePrompt<T>({
   useEffect(() => {
     if (promptState === PromptState.Submitted && answer) {
       setSearchTerm('')
-      unmountInk()
       onSubmit(answer.value)
+      complete()
     }
-  }, [answer, onSubmit, promptState, unmountInk])
+  }, [answer, onSubmit, promptState, complete])
 
   const setLoadingWhenSlow = useRef<NodeJS.Timeout>()
 

packages/cli-kit/src/private/node/ui/components/ConcurrentOutput.test.tsx

@@ -28,6 +28,7 @@ describe('ConcurrentOutput', () => {
     // Given
     const backendSync = new Synchronizer()
     const frontendSync = new Synchronizer()
+    const gate = new Synchronizer()
 
     const backendProcess = {
       prefix: 'backend',
@@ -37,6 +38,7 @@ describe('ConcurrentOutput', () => {
         stdout.write('third backend message')
 
         backendSync.resolve()
+        await gate.promise
       },
     }
 
@@ -50,6 +52,7 @@ describe('ConcurrentOutput', () => {
         stdout.write('third frontend message')
 
         frontendSync.resolve()
+        await gate.promise
       },
     }
     // When
@@ -72,19 +75,23 @@ describe('ConcurrentOutput', () => {
       00:00:00 │ frontend │ third frontend message
       "
     `)
+
+    gate.resolve()
   })
 
   test('strips ansi codes from the output by default', async () => {
     const output = 'foo'
 
     // Given
     const processSync = new Synchronizer()
+    const gate = new Synchronizer()
     const processes = [
       {
         prefix: '1',
         action: async (stdout: Writable, _stderr: Writable, _signal: AbortSignal) => {
           stdout.write(`\u001b[32m${output}\u001b[39m`)
           processSync.resolve()
+          await gate.promise
         },
       },
     ]
@@ -98,13 +105,15 @@ describe('ConcurrentOutput', () => {
     const logColumns = renderInstance.lastFrame()!.split('│')
     expect(logColumns.length).toBe(3)
     expect(logColumns[2]?.trim()).toEqual(output)
+    gate.resolve()
   })
 
   test('does not strip ansi codes from the output when stripAnsi is false', async () => {
     const output = '\u001b[32mfoo\u001b[39m'
 
     // Given
     const processSync = new Synchronizer()
+    const gate = new Synchronizer()
     const processes = [
       {
         prefix: '1',
@@ -113,6 +122,7 @@ describe('ConcurrentOutput', () => {
             stdout.write(output)
           })
           processSync.resolve()
+          await gate.promise
         },
       },
     ]
@@ -126,11 +136,13 @@ describe('ConcurrentOutput', () => {
     const logColumns = renderInstance.lastFrame()!.split('│')
     expect(logColumns.length).toBe(3)
     expect(logColumns[2]?.trim()).toEqual(output)
+    gate.resolve()
   })
 
   test('renders custom prefixes on log lines', async () => {
     // Given
     const processSync = new Synchronizer()
+    const gate = new Synchronizer()
     const extensionName = 'my-extension'
     const processes = [
       {
@@ -140,6 +152,7 @@ describe('ConcurrentOutput', () => {
             stdout.write('foo bar')
           })
           processSync.resolve()
+          await gate.promise
         },
       },
     ]
@@ -161,12 +174,14 @@ describe('ConcurrentOutput', () => {
     const logColumns = unstyled(renderInstance.lastFrame()!).split('│')
     expect(logColumns.length).toBe(3)
     expect(logColumns[1]?.trim()).toEqual(extensionName)
+    gate.resolve()
   })
 
   test('renders prefix column width based on prefixColumnSize', async () => {
     // Given
     const processSync1 = new Synchronizer()
     const processSync2 = new Synchronizer()
+    const gate = new Synchronizer()
 
     const columnSize = 5
     const processes = [
@@ -175,13 +190,15 @@ describe('ConcurrentOutput', () => {
         action: async (stdout: Writable, _stderr: Writable, _signal: AbortSignal) => {
           stdout.write('foo')
           processSync1.resolve()
+          await gate.promise
         },
       },
       {
         prefix: '1',
         action: async (stdout: Writable, _stderr: Writable, _signal: AbortSignal) => {
           stdout.write('bar')
           processSync2.resolve()
+          await gate.promise
         },
       },
     ]
@@ -206,22 +223,25 @@ describe('ConcurrentOutput', () => {
       // Including spacing
       expect(logColumns[1]?.length).toBe(columnSize + 2)
     })
+    gate.resolve()
   })
 
   test('renders prefix column width based on processes by default', async () => {
     // Given
     const processSync = new Synchronizer()
+    const gate = new Synchronizer()
     const processes = [
       {
         prefix: '1',
         action: async (stdout: Writable, _stderr: Writable, _signal: AbortSignal) => {
           stdout.write('foo')
           processSync.resolve()
+          await gate.promise
         },
       },
-      {prefix: '12', action: async () => {}},
-      {prefix: '123', action: async () => {}},
-      {prefix: '1234', action: async () => {}},
+      {prefix: '12', action: async () => gate.promise},
+      {prefix: '123', action: async () => gate.promise},
+      {prefix: '1234', action: async () => gate.promise},
     ]
 
     // When
@@ -234,20 +254,23 @@ describe('ConcurrentOutput', () => {
     expect(logColumns.length).toBe(3)
     // 4 is largest prefix, plus spacing
     expect(logColumns[1]?.length).toBe(4 + 2)
+    gate.resolve()
   })
 
   test('does not render prefix column larger than max', async () => {
     // Given
     const processSync = new Synchronizer()
+    const gate = new Synchronizer()
     const processes = [
       {
         prefix: '1',
         action: async (stdout: Writable, _stderr: Writable, _signal: AbortSignal) => {
           stdout.write('foo')
           processSync.resolve()
+          await gate.promise
         },
       },
-      {prefix: new Array(26).join('0'), action: async () => {}},
+      {prefix: new Array(26).join('0'), action: async () => gate.promise},
     ]
 
     // When
@@ -260,6 +283,7 @@ describe('ConcurrentOutput', () => {
     expect(logColumns.length).toBe(3)
     // 25 is largest column allowed, plus spacing
     expect(logColumns[1]?.length).toBe(25 + 2)
+    gate.resolve()
   })
 
   test('rejects with the error thrown inside one of the processes', async () => {

packages/cli-kit/src/private/node/ui/components/ConcurrentOutput.tsx

@@ -1,7 +1,8 @@
 import {OutputProcess} from '../../../../public/node/output.js'
 import {AbortSignal} from '../../../../public/node/abort.js'
+import {useComplete} from '../../ui.js'
 import React, {FunctionComponent, useCallback, useEffect, useMemo, useState} from 'react'
-import {Box, Static, Text, TextProps, useApp} from 'ink'
+import {Box, Static, Text, TextProps} from 'ink'
 import figures from 'figures'
 import stripAnsi from 'strip-ansi'
 
@@ -92,7 +93,8 @@ const ConcurrentOutput: FunctionComponent<ConcurrentOutputProps> = ({
   useAlternativeColorPalette = false,
 }) => {
   const [processOutput, setProcessOutput] = useState<Chunk[]>([])
-  const {exit: unmountInk} = useApp()
+  const [completionResult, setCompletionResult] = useState<{error?: Error} | null>(null)
+  const complete = useComplete()
   const concurrentColors: TextProps['color'][] = useMemo(
     () =>
       useAlternativeColorPalette
@@ -179,24 +181,25 @@ const ConcurrentOutput: FunctionComponent<ConcurrentOutputProps> = ({
           }),
         )
         if (!keepRunningAfterProcessesResolve) {
-          // Defer unmount so React 19 can flush batched setProcessOutput
-          // state updates before the component tree is torn down.
-          // Use setImmediate → setTimeout(0) to span two event-loop phases,
-          // giving React's scheduler (which uses setImmediate in Node.js)
-          // a full cycle to flush before we tear down the component tree.
-          setImmediate(() => setTimeout(() => unmountInk(), 0))
+          setCompletionResult({})
         }
         // eslint-disable-next-line no-catch-all/no-catch-all
       } catch (error: unknown) {
         if (!keepRunningAfterProcessesResolve) {
-          setImmediate(() => setTimeout(() => unmountInk(error as Error | undefined), 0))
+          setCompletionResult({error: error as Error})
         }
       }
     }
 
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
     runProcesses()
-  }, [abortSignal, processes, writableStream, unmountInk, keepRunningAfterProcessesResolve])
+  }, [abortSignal, processes, writableStream, keepRunningAfterProcessesResolve])
+
+  useEffect(() => {
+    if (completionResult !== null) {
+      complete(completionResult.error)
+    }
+  }, [completionResult, complete])
 
   const {lineVertical} = figures