fix(metro): patch every loaded Metro graph for uncached modules

CONTEXT.md

@@ -119,7 +119,7 @@ Compilation flow:
 
 Metro integration:
 
-- `withUniwindConfig(config, uniwindConfig)` patches Metro graph support for uncached modules.
+- `withUniwindConfig(config, uniwindConfig)` patches Metro graph support for uncached modules. The patches apply to the Metro this package resolves and to every Metro graph module already loaded, because the CLI running the server can bring its own Metro copy (React Native CLI's community plugin and `@expo/metro` do) and loads it before the config.
 - Metro adds `css` as source extension and removes it from asset extensions.
 - Metro transformer handles the configured CSS entry file specially. In development, native entries declare imported local CSS files as Metro dependencies, including nested imports and workspace files resolved outside `node_modules`, so token-only edits trigger recompilation. Dependencies are collected afresh on each compile.
 - Non-entry native CSS is an empty module in plain Metro; Expo keeps its own CSS handling. Web CSS handling is unchanged.

packages/uniwind/src/bundler/adapters/metro/patches.ts

@@ -29,58 +29,79 @@ interface InitialTraverseDependencies {
     __uniwindLazyCssEntryPatched?: boolean
 }
 
-export const patchMetroGraphToIncludeCssInLazyGraphs = (cssEntryPath: string) => {
-    const { Graph } = require('metro/private/DeltaBundler/Graph') as typeof MetroGraphModule
+type MetroGraph = typeof MetroGraphModule.Graph
 
-    // oxlint-disable-next-line @typescript-eslint/unbound-method
-    const original_initialTraverseDependencies = Graph.prototype.initialTraverseDependencies as unknown as InitialTraverseDependencies
+const graphModuleSuffix = path.join('metro', 'src', 'DeltaBundler', 'Graph.js')
 
-    if (original_initialTraverseDependencies.__uniwindLazyCssEntryPatched) {
-        return
-    }
+// The Metro this package resolves isn't always the one running the server: a CLI can depend on its own copy, as React
+// Native CLI's community plugin and `@expo/metro` do, while the workspace hoists another. Both load their Metro graph
+// before they evaluate the config that applies these patches, so every Metro graph module already loaded is patched too.
+const getMetroGraphs = () => {
+    const graphs = new Set<MetroGraph>([(require('metro/private/DeltaBundler/Graph') as typeof MetroGraphModule).Graph])
 
-    async function initialTraverseDependencies(this: Graph, options: GraphOptions<any>) {
-        if (options.lazy && options.transformOptions.dev) {
-            const entryPoints = this.entryPoints as Set<string>
-            entryPoints.add(cssEntryPath)
-        }
+    Object.entries(require.cache).forEach(([filename, loadedModule]) => {
+        const graph = (loadedModule?.exports as Partial<typeof MetroGraphModule> | undefined)?.Graph
 
-        return original_initialTraverseDependencies.call(this, options)
-    }
+        if (filename.endsWith(`${path.sep}${graphModuleSuffix}`) && typeof graph === 'function') {
+            graphs.add(graph)
+        }
+    })
 
-    // @ts-expect-error patch Graph initialTraverseDependencies method
-    Graph.prototype.initialTraverseDependencies = initialTraverseDependencies
-    initialTraverseDependencies.__uniwindLazyCssEntryPatched = true
+    return graphs
 }
 
-export const patchMetroGraphToSupportUncachedModules = () => {
-    const { Graph } = require('metro/private/DeltaBundler/Graph') as typeof MetroGraphModule
-
-    // oxlint-disable-next-line @typescript-eslint/unbound-method
-    const original_traverseDependencies = Graph.prototype.traverseDependencies as unknown as TraverseDependencies
+export const patchMetroGraphToIncludeCssInLazyGraphs = (cssEntryPath: string) => {
+    getMetroGraphs().forEach(GraphClass => {
+        // oxlint-disable-next-line @typescript-eslint/unbound-method
+        const original_initialTraverseDependencies = GraphClass.prototype.initialTraverseDependencies as unknown as InitialTraverseDependencies
 
-    if (original_traverseDependencies.__patched) {
-        return
-    }
+        if (original_initialTraverseDependencies.__uniwindLazyCssEntryPatched) {
+            return
+        }
 
-    original_traverseDependencies.__patched = true
-
-    function traverseDependencies(this: Graph, paths: Array<string>, options: GraphOptions<any>) {
-        this.dependencies.forEach(dependency => {
-            if (
-                dependency.output.find(file => (file.data as any)?.css?.skipCache === true)
-                && !paths.includes(dependency.path)
-            ) {
-                // @ts-expect-error Hidden property
-                dependency.unstable_transformResultKey = `${dependency.unstable_transformResultKey}.`
-                paths.push(dependency.path)
+        async function initialTraverseDependencies(this: Graph, options: GraphOptions<any>) {
+            if (options.lazy && options.transformOptions.dev) {
+                const entryPoints = this.entryPoints as Set<string>
+                entryPoints.add(cssEntryPath)
             }
-        })
 
-        return original_traverseDependencies.call(this, paths, options)
-    }
+            return original_initialTraverseDependencies.call(this, options)
+        }
+
+        // @ts-expect-error patch Graph initialTraverseDependencies method
+        GraphClass.prototype.initialTraverseDependencies = initialTraverseDependencies
+        initialTraverseDependencies.__uniwindLazyCssEntryPatched = true
+    })
+}
+
+export const patchMetroGraphToSupportUncachedModules = () => {
+    getMetroGraphs().forEach(GraphClass => {
+        // oxlint-disable-next-line @typescript-eslint/unbound-method
+        const original_traverseDependencies = GraphClass.prototype.traverseDependencies as unknown as TraverseDependencies
+
+        if (original_traverseDependencies.__patched) {
+            return
+        }
+
+        original_traverseDependencies.__patched = true
+
+        function traverseDependencies(this: Graph, paths: Array<string>, options: GraphOptions<any>) {
+            this.dependencies.forEach(dependency => {
+                if (
+                    dependency.output.find(file => (file.data as any)?.css?.skipCache === true)
+                    && !paths.includes(dependency.path)
+                ) {
+                    // @ts-expect-error Hidden property
+                    dependency.unstable_transformResultKey = `${dependency.unstable_transformResultKey}.`
+                    paths.push(dependency.path)
+                }
+            })
+
+            return original_traverseDependencies.call(this, paths, options)
+        }
 
-    // @ts-expect-error patch Graph traverseDependencies method
-    Graph.prototype.traverseDependencies = traverseDependencies
-    traverseDependencies.__patched = true
+        // @ts-expect-error patch Graph traverseDependencies method
+        GraphClass.prototype.traverseDependencies = traverseDependencies
+        traverseDependencies.__patched = true
+    })
 }

packages/uniwind/tests/native/bundler/metro-patches.test.ts

@@ -0,0 +1,89 @@
+import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { patchMetroGraphToIncludeCssInLazyGraphs, patchMetroGraphToSupportUncachedModules } from '../../../src/bundler/adapters/metro/patches'
+
+type FakeModule = {
+    path: string
+    output: Array<{ data: Record<string, unknown> }>
+    unstable_transformResultKey: string
+}
+
+type FakeGraph = {
+    dependencies: Map<string, FakeModule>
+    entryPoints: Set<string>
+    traverseDependencies: (paths: Array<string>, options: unknown) => Promise<{ paths: Array<string> }>
+    initialTraverseDependencies: (options: unknown) => Promise<{ entryPoints: Array<string> }>
+}
+
+// Another Metro copy than the one this package resolves, already loaded the way a CLI loads its own Metro
+// (React Native CLI's community plugin depends on one) before it evaluates the config.
+const fakeMetroGraphSource = `
+class Graph {
+    constructor() {
+        this.dependencies = new Map()
+        this.entryPoints = new Set(['/app/index.js'])
+    }
+
+    async traverseDependencies(paths) {
+        return { paths: [...paths] }
+    }
+
+    async initialTraverseDependencies() {
+        return { entryPoints: [...this.entryPoints] }
+    }
+}
+
+module.exports = { Graph }
+`
+
+let root: string
+let CliGraph: new () => FakeGraph
+
+beforeAll(() => {
+    root = mkdtempSync(join(tmpdir(), 'uniwind-metro-patches-'))
+    const graphModulePath = join(root, 'node_modules', 'metro', 'src', 'DeltaBundler', 'Graph.js')
+
+    mkdirSync(join(graphModulePath, '..'), { recursive: true })
+    writeFileSync(graphModulePath, fakeMetroGraphSource)
+    CliGraph = require(graphModulePath).Graph
+
+    patchMetroGraphToSupportUncachedModules()
+    patchMetroGraphToIncludeCssInLazyGraphs('/app/global.css')
+})
+
+afterAll(() => {
+    rmSync(root, { recursive: true, force: true })
+})
+
+test('re-traverses uncached modules in every loaded Metro graph', async () => {
+    const graph = new CliGraph()
+    const cssEntry: FakeModule = { path: '/app/global.css', output: [{ data: { css: { skipCache: true } } }], unstable_transformResultKey: 'key' }
+    const cachedModule: FakeModule = { path: '/app/App.tsx', output: [{ data: {} }], unstable_transformResultKey: 'key' }
+
+    graph.dependencies.set(cssEntry.path, cssEntry)
+    graph.dependencies.set(cachedModule.path, cachedModule)
+
+    await expect(graph.traverseDependencies(['/app/tokens.css'], {})).resolves.toEqual({
+        paths: ['/app/tokens.css', '/app/global.css'],
+    })
+    expect(cssEntry.unstable_transformResultKey).toBe('key.')
+    expect(cachedModule.unstable_transformResultKey).toBe('key')
+})
+
+test('adds the CSS entry to lazy development graphs of every loaded Metro graph', async () => {
+    await expect(new CliGraph().initialTraverseDependencies({ lazy: true, transformOptions: { dev: true } })).resolves.toEqual({
+        entryPoints: ['/app/index.js', '/app/global.css'],
+    })
+    await expect(new CliGraph().initialTraverseDependencies({ lazy: false, transformOptions: { dev: true } })).resolves.toEqual({
+        entryPoints: ['/app/index.js'],
+    })
+})
+
+test('patches each Metro graph once', () => {
+    const { traverseDependencies } = CliGraph.prototype
+
+    patchMetroGraphToSupportUncachedModules()
+
+    expect(CliGraph.prototype.traverseDependencies).toBe(traverseDependencies)
+})