Merge pull request #7157 from Shopify/revert-7153-rcb/rendering-changes-react-fix

Revert "use central unmount for ink render tree; stop messing with node's event loop"

Author: isaacroldan

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

@@ -22,18 +22,9 @@ 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,
-    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 {}
-    }),
-  }
+  return {...realModule, ...mockModule}
 })
 
 describe('setupPreviewThemeAppExtensionsProcess', () => {

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

@@ -1,5 +1,5 @@
-import {Stdout, InkLifecycleRoot} from '../ui.js'
-import React, {ReactElement} from 'react'
+import {Stdout} from '../ui.js'
+import {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(React.createElement(InkLifecycleRoot, null, tree), {
+  const instance = inkRender(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: (tree: ReactElement) => instance.rerender(React.createElement(InkLifecycleRoot, null, tree)),
+    rerender: instance.rerender,
     unmount: instance.unmount,
     cleanup: instance.cleanup,
     waitUntilExit: () => trackPromise(instance.waitUntilExit().then(() => {})),

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

@@ -3,49 +3,11 @@ 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 React, {ReactElement, createContext, useCallback, useContext, useEffect, useState} from 'react'
-import {Key, render as inkRender, RenderOptions, useApp} from 'ink'
+import {ReactElement} from 'react'
+import {Key, render as inkRender, RenderOptions} 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
@@ -65,8 +27,10 @@ export function renderOnce(element: JSX.Element, {logLevel = 'info', renderOptio
 }
 
 export async function render(element: JSX.Element, options?: RenderOptions) {
-  const {waitUntilExit} = inkRender(<InkLifecycleRoot>{element}</InkLifecycleRoot>, options)
+  const {waitUntilExit} = inkRender(element, 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,11 +5,10 @@ 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} from 'ink'
+import {Box, useApp} from 'ink'
 
 export interface SearchResults<T> {
   data: SelectItem<T>[]
@@ -43,7 +42,7 @@ function AutocompletePrompt<T>({
   infoMessage,
   groupOrder,
 }: React.PropsWithChildren<AutocompletePromptProps<T>>): ReactElement | null {
-  const complete = useComplete()
+  const {exit: unmountInk} = useApp()
   const [searchTerm, setSearchTerm] = useState('')
   const [searchResults, setSearchResults] = useState<SelectItem<T>[]>(choices)
   const canSearch = choices.length > MIN_NUMBER_OF_ITEMS_FOR_SEARCH
@@ -73,10 +72,10 @@ function AutocompletePrompt<T>({
   useEffect(() => {
     if (promptState === PromptState.Submitted && answer) {
       setSearchTerm('')
+      unmountInk()
       onSubmit(answer.value)
-      complete()
     }
-  }, [answer, onSubmit, promptState, complete])
+  }, [answer, onSubmit, promptState, unmountInk])
 
   const setLoadingWhenSlow = useRef<NodeJS.Timeout>()
 

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

@@ -28,7 +28,6 @@ describe('ConcurrentOutput', () => {
     // Given
     const backendSync = new Synchronizer()
     const frontendSync = new Synchronizer()
-    const gate = new Synchronizer()
 
     const backendProcess = {
       prefix: 'backend',
@@ -38,7 +37,6 @@ describe('ConcurrentOutput', () => {
         stdout.write('third backend message')
 
         backendSync.resolve()
-        await gate.promise
       },
     }
 
@@ -52,7 +50,6 @@ describe('ConcurrentOutput', () => {
         stdout.write('third frontend message')
 
         frontendSync.resolve()
-        await gate.promise
       },
     }
     // When
@@ -75,23 +72,19 @@ 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
         },
       },
     ]
@@ -105,15 +98,13 @@ 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',
@@ -122,7 +113,6 @@ describe('ConcurrentOutput', () => {
             stdout.write(output)
           })
           processSync.resolve()
-          await gate.promise
         },
       },
     ]
@@ -136,13 +126,11 @@ 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 = [
       {
@@ -152,7 +140,6 @@ describe('ConcurrentOutput', () => {
             stdout.write('foo bar')
           })
           processSync.resolve()
-          await gate.promise
         },
       },
     ]
@@ -174,14 +161,12 @@ 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 = [
@@ -190,15 +175,13 @@ 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
         },
       },
     ]
@@ -223,25 +206,22 @@ 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 () => gate.promise},
-      {prefix: '123', action: async () => gate.promise},
-      {prefix: '1234', action: async () => gate.promise},
+      {prefix: '12', action: async () => {}},
+      {prefix: '123', action: async () => {}},
+      {prefix: '1234', action: async () => {}},
     ]
 
     // When
@@ -254,23 +234,20 @@ 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 () => gate.promise},
+      {prefix: new Array(26).join('0'), action: async () => {}},
     ]
 
     // When
@@ -283,7 +260,6 @@ 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,8 +1,7 @@
 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} from 'ink'
+import {Box, Static, Text, TextProps, useApp} from 'ink'
 import figures from 'figures'
 import stripAnsi from 'strip-ansi'
 
@@ -93,8 +92,7 @@ const ConcurrentOutput: FunctionComponent<ConcurrentOutputProps> = ({
   useAlternativeColorPalette = false,
 }) => {
   const [processOutput, setProcessOutput] = useState<Chunk[]>([])
-  const [completionResult, setCompletionResult] = useState<{error?: Error} | null>(null)
-  const complete = useComplete()
+  const {exit: unmountInk} = useApp()
   const concurrentColors: TextProps['color'][] = useMemo(
     () =>
       useAlternativeColorPalette
@@ -181,25 +179,24 @@ const ConcurrentOutput: FunctionComponent<ConcurrentOutputProps> = ({
           }),
         )
         if (!keepRunningAfterProcessesResolve) {
-          setCompletionResult({})
+          // 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))
         }
         // eslint-disable-next-line no-catch-all/no-catch-all
       } catch (error: unknown) {
         if (!keepRunningAfterProcessesResolve) {
-          setCompletionResult({error: error as Error})
+          setImmediate(() => setTimeout(() => unmountInk(error as Error | undefined), 0))
         }
       }
     }
 
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
     runProcesses()
-  }, [abortSignal, processes, writableStream, keepRunningAfterProcessesResolve])
-
-  useEffect(() => {
-    if (completionResult !== null) {
-      complete(completionResult.error)
-    }
-  }, [completionResult, complete])
+  }, [abortSignal, processes, writableStream, unmountInk, keepRunningAfterProcessesResolve])
 
   const {lineVertical} = figures