@wix/zero-config-implementation 1.77.0 → 1.79.0

package/src/index.ts

@@ -5,6 +5,10 @@ import React, { type ComponentType } from 'react'
 
 import { toEditorReactComponent } from './converters'
 import { BaseError, IoError, type NotFoundError, ParseError } from './errors'
+import { buildContextProviderModules } from './extensions/context-providers/context'
+import { buildContextAwareWrapper, loadMockProviders } from './extensions/context-providers/mock-provider'
+import { buildRefElementContext } from './extensions/ref-elements/context'
+import type { RefElementContext } from './extensions/ref-elements/types'
 import type { ExtractionError } from './extraction-types'
 import type { RunExtractorsOptions } from './information-extractors/react'
 import { extractCssImports, extractDefaultComponentInfo } from './information-extractors/ts'
@@ -12,8 +16,6 @@ import type { ComponentInfo } from './information-extractors/ts/types'
 import { processComponent } from './manifest-pipeline'
 import { findComponent, findDefaultComponent, loadModuleForExtraction } from './module-loader'
 import type { LoadModuleFailure } from './module-loader'
-import { buildRefElementContext } from './ref-elements/context'
-import type { RefElementContext } from './ref-elements/types'
 import { compileTsFile } from './ts-compiler'
 
 const defaultWrapper = (Component: ComponentType<unknown>): ComponentType<unknown> => {
@@ -64,10 +66,6 @@ export function extractComponentManifestResult(
     options?.onError?.(error)
   }
 
-  const optionsWithDefaults: ExtractComponentManifestOptions = {
-    ...options,
-    wrapper: options?.wrapper ?? defaultWrapper,
-  }
   // Step 1: Compile TypeScript (fatal)
   return compileTsFile(componentPath)
     .andThen((program) => {
@@ -91,79 +89,122 @@ export function extractComponentManifestResult(
     .andThen(({ program, componentInfo }) => {
       const { componentName } = componentInfo
 
+      // Step 3a: Build ref-element context
       return ResultAsync.fromPromise(
         buildRefElementContext(program, componentPath, componentInfo, compiledEntryPath),
         (thrown) =>
           new ParseError(
-            `Failed to build the ref-element context for "${componentPath}": ${thrown instanceof Error ? thrown.message : String(thrown)}`,
+            `Failed to build ref-element context for "${componentPath}": ${thrown instanceof Error ? thrown.message : String(thrown)}`,
             { cause: thrown instanceof Error ? thrown : undefined, props: { phase: 'extract' } },
           ),
       ).andThen((refElementContext) =>
-        // Step 3: Load the compiled package module (non-fatal) — done after TS extraction
-        // so componentName is known when reporting failures
-        loadModuleForExtraction(compiledEntryPath, refElementContext.runtimeModules)
-          .map(({ cleanup, moduleExports }) => ({
-            loadComponent: (name: string) => findComponent(moduleExports, name) ?? findDefaultComponent(moduleExports),
-            failure: null as LoadModuleFailure | null,
-            cleanup,
-          }))
-          .orElse((failure) =>
-            okAsync({
-              loadComponent: () => null as ComponentType<unknown> | null,
-              cleanup: async () => {},
-              failure,
-            }),
-          )
-          .andThen(({ cleanup, failure, loadComponent }) => {
-            if (failure) {
-              if (failure.esmError) {
-                report({
-                  componentName,
-                  phase: 'loader',
-                  error: new IoError('ESM import failed', { cause: failure.esmError, props: { phase: 'loader' } }),
-                })
-              }
-              report({
-                componentName,
-                phase: 'loader',
-                error: new IoError('CJS require failed', { cause: failure.cjsError, props: { phase: 'loader' } }),
-              })
+        // Step 3b: Build context provider modules
+        ResultAsync.fromPromise(
+          buildContextProviderModules(program, componentPath, compiledEntryPath),
+          (thrown) =>
+            new ParseError(
+              `Failed to build context provider modules for "${componentPath}": ${thrown instanceof Error ? thrown.message : String(thrown)}`,
+              { cause: thrown instanceof Error ? thrown : undefined, props: { phase: 'extract' } },
+            ),
+        ).andThen((contextProviderModules) => {
+          // Step 4: Build context provider wrapper with mock providers
+          const buildWrapper = async (): Promise<(Component: ComponentType<unknown>) => ComponentType<unknown>> => {
+            const baseWrapper = options?.wrapper ?? defaultWrapper
+            if (contextProviderModules.length === 0) {
+              return baseWrapper
+            }
+            const loadedProviders = await loadMockProviders(
+              contextProviderModules,
+              React,
+              refElementContext.runtimeModules,
+            )
+            if (loadedProviders.length === 0) {
+              return baseWrapper
             }
+            return buildContextAwareWrapper(loadedProviders, baseWrapper, React)
+          }
 
-            // Step 4: Extract CSS imports (non-fatal)
-            let cssImportPaths: string[] = []
-            const cssResult = Result.fromThrowable(extractCssImports, (thrown) => thrown)(program)
-            if (cssResult.isOk()) {
-              cssImportPaths = cssResult.value
-            } else {
-              const thrown = cssResult.error
-              report({
-                componentName,
-                phase: 'css',
-                error: new ParseError(
-                  `Failed to extract CSS imports: ${thrown instanceof Error ? thrown.message : String(thrown)}`,
-                  { cause: thrown instanceof Error ? thrown : undefined, props: { phase: 'css' } },
-                ),
-              })
+          return ResultAsync.fromPromise(
+            buildWrapper(),
+            (thrown) =>
+              new ParseError(
+                `Failed to build context provider wrapper for "${componentPath}": ${thrown instanceof Error ? thrown.message : String(thrown)}`,
+                { cause: thrown instanceof Error ? thrown : undefined, props: { phase: 'extract' } },
+              ),
+          ).andThen((wrapper) => {
+            const optionsWithDefaults: ExtractComponentManifestOptions = {
+              ...options,
+              wrapper,
             }
 
-            // Step 5: Process the default-exported component (non-fatal unless render fails after load)
-            return ResultAsync.fromPromise(
-              processComponentWithCleanup(
-                componentInfo,
-                loadComponent,
-                cssImportPaths,
-                !!failure,
-                report,
-                optionsWithDefaults,
-                compiledEntryPath,
-                refElementContext,
+            // Step 5: Load the compiled package module (non-fatal) - done after TS extraction
+            // so componentName is known when reporting failures
+            return loadModuleForExtraction(compiledEntryPath, refElementContext.runtimeModules)
+              .map(({ cleanup, moduleExports }) => ({
+                loadComponent: (name: string) =>
+                  findComponent(moduleExports, name) ?? findDefaultComponent(moduleExports),
+                failure: null as LoadModuleFailure | null,
                 cleanup,
-                errors,
-              ),
-              normalizeComponentExtractionCleanupFailure,
-            )
-          }),
+              }))
+              .orElse((failure) =>
+                okAsync({
+                  loadComponent: () => null as ComponentType<unknown> | null,
+                  cleanup: async () => {},
+                  failure,
+                }),
+              )
+              .andThen(({ cleanup, failure, loadComponent }) => {
+                if (failure) {
+                  if (failure.esmError) {
+                    report({
+                      componentName,
+                      phase: 'loader',
+                      error: new IoError('ESM import failed', { cause: failure.esmError, props: { phase: 'loader' } }),
+                    })
+                  }
+                  report({
+                    componentName,
+                    phase: 'loader',
+                    error: new IoError('CJS require failed', { cause: failure.cjsError, props: { phase: 'loader' } }),
+                  })
+                }
+
+                // Step 6: Extract CSS imports (non-fatal)
+                let cssImportPaths: string[] = []
+                const cssResult = Result.fromThrowable(extractCssImports, (thrown) => thrown)(program)
+                if (cssResult.isOk()) {
+                  cssImportPaths = cssResult.value
+                } else {
+                  const thrown = cssResult.error
+                  report({
+                    componentName,
+                    phase: 'css',
+                    error: new ParseError(
+                      `Failed to extract CSS imports: ${thrown instanceof Error ? thrown.message : String(thrown)}`,
+                      { cause: thrown instanceof Error ? thrown : undefined, props: { phase: 'css' } },
+                    ),
+                  })
+                }
+
+                // Step 7: Process the default-exported component
+                return ResultAsync.fromPromise(
+                  processComponentWithCleanup(
+                    componentInfo,
+                    loadComponent,
+                    cssImportPaths,
+                    !!failure,
+                    report,
+                    optionsWithDefaults,
+                    compiledEntryPath,
+                    refElementContext,
+                    cleanup,
+                    errors,
+                  ),
+                  normalizeComponentExtractionCleanupFailure,
+                )
+              })
+          })
+        }),
       )
     })
 }

package/src/information-extractors/react/extractors/prop-tracker.test.ts

@@ -1,6 +1,6 @@
 import { describe, expect, it } from 'vitest'
-import { createRefElementMarker } from '../../../ref-elements/path-utils'
-import type { RefElementMatch } from '../../../ref-elements/types'
+import { createRefElementMarker } from '../../../extensions/ref-elements/path-utils'
+import type { RefElementMatch } from '../../../extensions/ref-elements/types'
 import { ExtractorStore } from './core/store'
 import { createPropTrackerExtractor } from './prop-tracker'
 

package/src/information-extractors/react/extractors/prop-tracker.ts

@@ -7,9 +7,9 @@
 
 import type { HTMLAttributes } from 'react'
 import { TRACE_ATTR } from '../../../component-renderer'
-import { readRefElementComponentType } from '../../../ref-elements/component-tag'
-import { parseRefElementMarker } from '../../../ref-elements/path-utils'
-import type { RefElementMatch } from '../../../ref-elements/types'
+import { readRefElementComponentType } from '../../../extensions/ref-elements/component-tag'
+import { parseRefElementMarker } from '../../../extensions/ref-elements/path-utils'
+import type { RefElementMatch } from '../../../extensions/ref-elements/types'
 import { findPreferredSemanticClass, normalizeClassNames } from '../../../utils/css-class'
 import type { PropSpyMeta, TrackingStores } from '../types'
 import { type PropSpyRegistrar, generateMockProps, resetMockCounter } from '../utils/mock-generator'

package/src/information-extractors/react/utils/mock-generator.test.ts

@@ -1,7 +1,7 @@
-import { describe, expect, it } from 'vitest'
-import { createRefElementMarker, parseRefElementMarker } from '../../../ref-elements/path-utils'
+import { beforeEach, describe, expect, it } from 'vitest'
+import { createRefElementMarker, parseRefElementMarker } from '../../../extensions/ref-elements/path-utils'
 import type { ComponentInfo } from '../../ts/types'
-import { generateMockProps, resetMockCounter } from './mock-generator'
+import { generateMockContextValues, generateMockProps, resetMockCounter } from './mock-generator'
 
 describe('generateMockProps', () => {
   it('injects a dedicated ref-element id marker into each eligible elementProps branch', () => {
@@ -129,3 +129,107 @@ describe('generateMockProps', () => {
     expect(parseRefElementMarker(ctaProps.id)).toBe('card.elementProps.cta')
   })
 })
+
+describe('generateMockContextValues', () => {
+  beforeEach(() => {
+    resetMockCounter()
+  })
+
+  it('returns an empty object for an empty schema', () => {
+    expect(generateMockContextValues({ items: {} })).toEqual({})
+  })
+
+  it('maps text to a string', () => {
+    const result = generateMockContextValues({ items: { title: { dataType: 'text' } } })
+    expect(typeof result.title).toBe('string')
+  })
+
+  it('maps number to a number', () => {
+    const result = generateMockContextValues({ items: { count: { dataType: 'number' } } })
+    expect(typeof result.count).toBe('number')
+  })
+
+  it('maps boolean to a boolean', () => {
+    const result = generateMockContextValues({ items: { active: { dataType: 'boolean' } } })
+    expect(typeof result.active).toBe('boolean')
+  })
+
+  it('maps function to a function', () => {
+    const result = generateMockContextValues({ items: { onAction: { dataType: 'function' } } })
+    expect(typeof result.onAction).toBe('function')
+  })
+
+  it('maps richText to an HTML paragraph string', () => {
+    const result = generateMockContextValues({ items: { body: { dataType: 'richText' } } })
+    expect(result.body as string).toMatch(/^<p>/)
+  })
+
+  it('maps image to an image-shaped object', () => {
+    const result = generateMockContextValues({ items: { photo: { dataType: 'image' } } })
+    expect(result.photo).toMatchObject({
+      uri: expect.any(String),
+      url: expect.any(String),
+      width: expect.any(Number),
+      height: expect.any(Number),
+    })
+  })
+
+  it('maps link to a link-shaped object', () => {
+    const result = generateMockContextValues({ items: { destination: { dataType: 'link' } } })
+    expect(result.destination).toMatchObject({ href: expect.any(String), target: expect.any(String) })
+  })
+
+  it('maps data with a nested schema to a nested object', () => {
+    const result = generateMockContextValues({
+      items: {
+        event: {
+          dataType: 'data',
+          data: { items: { name: { dataType: 'text' } } },
+        },
+      },
+    })
+    expect(result.event).toMatchObject({ name: expect.any(String) })
+  })
+
+  it('maps data with an empty schema to a permissive mock that handles Object.keys safely', () => {
+    const result = generateMockContextValues({
+      items: { event: { dataType: 'data', data: { items: {} } } },
+    })
+    expect(() => Object.keys(result.event as object)).not.toThrow()
+    expect(Object.keys(result.event as object)).toHaveLength(0)
+  })
+
+  it('returns a permissive fallback for undeclared context fields so deep access chains do not crash', () => {
+    const result = generateMockContextValues({ items: {} })
+    const undeclared = (result as Record<string, unknown>).anything
+    expect(() => Object.keys(undeclared as object)).not.toThrow()
+    expect(Object.keys(undeclared as object)).toHaveLength(0)
+    const deeplyNested = (undeclared as Record<string, unknown>).deeply
+    expect(() => (deeplyNested as Record<string, unknown>)?.nested).not.toThrow()
+  })
+
+  it('maps arrayItems with a nested schema to a single-element array containing the element object', () => {
+    const result = generateMockContextValues({
+      items: {
+        tickets: {
+          dataType: 'arrayItems',
+          arrayItems: { item: { dataType: 'data', data: { items: { price: { dataType: 'number' } } } } },
+        },
+      },
+    })
+    expect(result.tickets).toHaveLength(1)
+    expect((result.tickets as Record<string, unknown>[])[0]).toMatchObject({ price: expect.any(Number) })
+  })
+
+  it('maps arrayItems with no item schema to an empty array', () => {
+    const result = generateMockContextValues({
+      items: { tickets: { dataType: 'arrayItems', arrayItems: {} } },
+    })
+    expect(result.tickets).toEqual([])
+  })
+
+  it('falls back to a string for unknown dataTypes', () => {
+    const result = generateMockContextValues({ items: { something: { dataType: 'unknown-type' } } })
+    expect(typeof result.something).toBe('string')
+  })
+})

package/src/information-extractors/react/utils/mock-generator.ts

@@ -4,7 +4,11 @@
  */
 
 import { faker } from '@faker-js/faker'
-import { createRefElementMarker, extractElementPropsPathFromPropPath } from '../../../ref-elements/path-utils'
+import type { ContextSchema } from '../../../extensions/context-providers/types'
+import {
+  createRefElementMarker,
+  extractElementPropsPathFromPropPath,
+} from '../../../extensions/ref-elements/path-utils'
 import type { ComponentInfo, DefaultValue, PropInfo, ResolvedType } from '../../ts/types'
 
 export const PRESETS_WRAPPER_CLASS_NAME = 'mock-presets-wrapper-probe'
@@ -484,3 +488,102 @@ function generateMockMenuItems(): unknown[] {
     link: generateMockLink(),
   }))
 }
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Context Schema Mock Generator
+// ─────────────────────────────────────────────────────────────────────────────
+
+/**
+ * A recursive Proxy that safely handles any property access, function call, or string coercion.
+ * Used as a fallback for context fields not declared in the schema.
+ */
+function buildPermissiveMock(): unknown {
+  const mock: unknown = new Proxy((() => {}) as (...args: unknown[]) => unknown, {
+    get(_target, prop) {
+      if (prop === 'then') return undefined // prevent thenable detection
+      if (prop === Symbol.toPrimitive) return (_hint: string) => 'mock'
+      if (prop === Symbol.iterator) return function* () {}
+      if (prop === Symbol.asyncIterator) return async function* () {}
+      if (typeof prop === 'symbol') return undefined
+      return mock
+    },
+    ownKeys() {
+      return []
+    },
+    getOwnPropertyDescriptor() {
+      return undefined
+    },
+    apply() {
+      return mock
+    },
+  })
+  return mock
+}
+
+/**
+ * Wraps schema-derived values in a Proxy so that any property NOT declared in the schema
+ * returns a permissive mock instead of undefined, preventing crashes on undeclared field access.
+ */
+function buildSchemaSeededProxy(schemaValues: Record<string, unknown>): Record<string, unknown> {
+  return new Proxy(schemaValues, {
+    get(target, prop) {
+      if (prop === 'then') return undefined // prevent thenable/Promise detection
+      if (typeof prop === 'symbol') return Reflect.get(target, prop)
+      if (Object.prototype.hasOwnProperty.call(target, prop)) return target[prop as string]
+      return buildPermissiveMock()
+    },
+  })
+}
+
+/** Generate mock context values from a context extension schema. */
+export function generateMockContextValues(schema: ContextSchema): Record<string, unknown> {
+  const result: Record<string, unknown> = {}
+  for (const [key, item] of Object.entries(schema.items)) {
+    result[key] = generateValueFromResolvedType(contextDataTypeToResolvedType(item), key, `context.${key}`)
+  }
+  return buildSchemaSeededProxy(result)
+}
+
+function contextDataTypeToResolvedType(item: ContextSchema['items'][string]): ResolvedType {
+  switch (item.dataType) {
+    case 'text':
+      return { kind: 'primitive', value: 'string' }
+    case 'richText':
+      return { kind: 'semantic', value: 'RichText' }
+    case 'image':
+      return { kind: 'semantic', value: 'Image' }
+    case 'video':
+      return { kind: 'semantic', value: 'Video' }
+    case 'vectorArt':
+      return { kind: 'semantic', value: 'VectorArt' }
+    case 'link':
+      return { kind: 'semantic', value: 'Link' }
+    case 'function':
+      return { kind: 'function' }
+    case 'number':
+      return { kind: 'primitive', value: 'number' }
+    case 'boolean':
+      return { kind: 'primitive', value: 'boolean' }
+    case 'arrayItems': {
+      const itemEntry = item.arrayItems?.item
+      if (itemEntry) {
+        const itemValue = generateValueFromResolvedType(
+          contextDataTypeToResolvedType(itemEntry),
+          'item',
+          'context.arrayItem',
+        )
+        return { kind: 'literal', value: [itemValue] }
+      }
+      return { kind: 'literal', value: [] }
+    }
+    case 'data': {
+      const nestedSchema = item.data
+      if (nestedSchema?.items && Object.keys(nestedSchema.items).length > 0) {
+        return { kind: 'literal', value: generateMockContextValues(nestedSchema) }
+      }
+      return { kind: 'literal', value: buildPermissiveMock() }
+    }
+    default:
+      return { kind: 'primitive', value: 'string' }
+  }
+}

package/src/manifest-pipeline.ts

@@ -15,10 +15,13 @@ import {
 } from './information-extractors/react'
 import type { CoupledComponentInfo, CoupledProp, DOMBinding, TrackingStores } from './information-extractors/react'
 
+import {
+  extractElementPropsPathFromPropPath,
+  extractParentElementPropsPath,
+} from './extensions/ref-elements/path-utils'
 import { matchCssSelectors, parseCss } from './information-extractors/css'
 import type { CSSParserAPI } from './information-extractors/css'
 import { compileSass } from './information-extractors/css/sass-adapter'
-import { extractElementPropsPathFromPropPath, extractParentElementPropsPath } from './ref-elements/path-utils'
 
 import type { ComponentType } from 'react'
 

package/src/module-loader.test.ts

@@ -2,9 +2,9 @@ import { createRequire } from 'node:module'
 import path from 'node:path'
 import { pathToFileURL } from 'node:url'
 import { describe, expect, it } from 'vitest'
+import { readRefElementComponentType } from './extensions/ref-elements/component-tag'
 import { loadCjsModule, loadModuleForExtraction } from './module-loader'
 import { readLoaderPortRefStateForTests, readRegisteredRefElementModuleUrlsForTests } from './react-runtime-loader'
-import { readRefElementComponentType } from './ref-elements/component-tag'
 
 const require = createRequire(import.meta.url)
 

package/src/module-loader.ts

@@ -4,8 +4,8 @@ import React from 'react'
 import type { ComponentType } from 'react'
 import ReactDOM from 'react-dom'
 
+import type { RefElementModule } from './extensions/ref-elements/types'
 import { ensureReactRuntimeLoader, loadTaggedCjsModule, prepareRefElementEsmImport } from './react-runtime-loader'
-import type { RefElementModule } from './ref-elements/types'
 
 /**
  * Structured failure from `loadModule` when both ESM import and CJS require fail.

package/src/react-runtime-loader.ts

@@ -2,12 +2,13 @@ import { createHash } from 'node:crypto'
 import { createRequire, register } from 'node:module'
 import { fileURLToPath, pathToFileURL } from 'node:url'
 import { MessageChannel, MessagePort } from 'node:worker_threads'
-import { buildRefElementTaggerSource, tagRefElementComponentForCleanup } from './ref-elements/component-tag'
-import type { RefElementModule } from './ref-elements/types'
+import { buildRefElementTaggerSource, tagRefElementComponentForCleanup } from './extensions/ref-elements/component-tag'
+import type { RefElementModule } from './extensions/ref-elements/types'
 
 const JSX_INTERCEPTOR_STATE_KEY = 'zero-config:jsx-interceptor'
 const JSX_ORIGINALS_KEY = 'zero-config:jsx-originals'
 const REACT_ORIGINALS_KEY = 'zero-config:react-originals'
+const REACT_ORIGINALS_SYMBOL = Symbol.for(REACT_ORIGINALS_KEY)
 const LOADER_REGISTERED_KEY = Symbol.for('zero-config:jsx-loader-registered')
 const RUNTIME_STATE_KEY = Symbol.for('zero-config:ref-element-runtime-state')
 const REF_ELEMENT_LOADER_REQUEST_TIMEOUT_MS = 1_000
@@ -67,6 +68,32 @@ interface RefElementRuntimeState {
   originalModuleLoad: InternalModuleApi['_load'] | null
 }
 
+/**
+ * Returns the URL that the ESM loader will assign to a context provider module
+ * when it is imported transitively from a ref-entry-signed statics bundle.
+ *
+ * When runtimeModules is non-empty, the statics bundle entry URL carries a
+ * `?zero-config-ref-entry=HASH` signature, and `shouldPropagateRefElementEntrySignature`
+ * propagates that hash to every non-node_modules file: import. `captureHookContext`
+ * must import the provider at the same signed URL so both paths hit the same Node.js
+ * ESM cache entry (and therefore the same `createContext` return value).
+ *
+ * When runtimeModules is empty, no hash is ever appended, so the plain URL is correct.
+ */
+export function computeSignedContextProviderUrl(resolvedModuleUrl: string, runtimeModules: RefElementModule[]): string {
+  if (runtimeModules.length === 0) {
+    return resolvedModuleUrl
+  }
+
+  const registeredModules = mergeRefElementModules(runtimeModules)
+  const signature = buildRefElementEntrySignature(registeredModules)
+  const signedUrl = new URL(resolvedModuleUrl)
+  if (!signedUrl.searchParams.has(REF_ELEMENT_ENTRY_SIGNATURE_QUERY_PARAM)) {
+    signedUrl.searchParams.set(REF_ELEMENT_ENTRY_SIGNATURE_QUERY_PARAM, signature)
+  }
+  return signedUrl.href
+}
+
 export async function prepareRefElementEsmImport(
   entryPath: string,
   runtimeModules: RefElementModule[] = [],
@@ -128,9 +155,9 @@ export function ensureReactRuntimeLoader(): void {
       jsxDEV: reactDevRuntime.jsxDEV,
       jsxs: reactRuntime.jsxs,
     }
-    ;(globalThis as Record<PropertyKey, unknown>)[Symbol.for(REACT_ORIGINALS_KEY)] = reactModule
+    ;(globalThis as Record<PropertyKey, unknown>)[REACT_ORIGINALS_SYMBOL] = reactModule
 
-    const reactModuleUrl = `data:text/javascript,${encodeURIComponent(buildReactModuleShimSource(Object.keys(reactModule)))}`
+    const reactModuleUrl = `data:text/javascript,${encodeURIComponent(buildReactModuleShimSource(reactModule))}`
     const interceptorModuleUrl = `data:text/javascript,${encodeURIComponent(buildJsxRuntimeInterceptorSource())}`
     const loaderSource = buildReactRuntimeLoaderSource({
       interceptorModuleUrl,
@@ -151,6 +178,14 @@ export function ensureReactRuntimeLoader(): void {
   }
 }
 
+export function getActiveReactModule(): unknown {
+  return (globalThis as Record<symbol, unknown>)[REACT_ORIGINALS_SYMBOL]
+}
+
+export function setActiveReactModule(reactModule: unknown): void {
+  ;(globalThis as Record<symbol, unknown>)[REACT_ORIGINALS_SYMBOL] = reactModule
+}
+
 export function readRegisteredRefElementModuleUrlsForTests(): string[] {
   return [...readRefElementRuntimeState().registeredEsmModulesByResolvedUrl.keys()]
 }
@@ -483,14 +518,20 @@ function isComponent(value: unknown): boolean {
   return typeof value === 'object' && value !== null && '$$typeof' in value
 }
 
-function buildReactModuleShimSource(reactExportNames: string[]): string {
-  const reactExportLines = reactExportNames
+function buildReactModuleShimSource(reactModule: Record<string, unknown>): string {
+  const reactExportLines = Object.keys(reactModule)
     .map((exportName) => {
       if (exportName === 'createElement') {
         return `export function createElement() {
   var state = globalThis[STATE_KEY];
   if (state && state.currentCreateElement) return state.currentCreateElement.apply(null, arguments);
-  return reactModule.createElement.apply(null, arguments);
+  return globalThis[REACT_KEY].createElement.apply(null, arguments);
+}`
+      }
+
+      if (exportName.startsWith('use') && typeof reactModule[exportName] === 'function') {
+        return `export function ${exportName}() {
+  return globalThis[REACT_KEY].${exportName}.apply(null, arguments);
 }`
       }
 
@@ -502,7 +543,7 @@ function buildReactModuleShimSource(reactExportNames: string[]): string {
 var REACT_KEY = Symbol.for(${JSON.stringify(REACT_ORIGINALS_KEY)});
 var STATE_KEY = Symbol.for(${JSON.stringify(JSX_INTERCEPTOR_STATE_KEY)});
 var reactModule = globalThis[REACT_KEY];
-export default reactModule;
+export default new Proxy({}, { get: function(_, prop) { return globalThis[REACT_KEY][prop]; } });
 ${reactExportLines}
 `
 }