fix(react/runtime): fix runOnBackground() failure in certain cases

.changeset/kind-rocks-spend.md

@@ -0,0 +1,5 @@
+---
+"@lynx-js/react": patch
+---
+
+fix: main thread functions created during the initial render cannot correctly call `runOnBackground()` after hydration

packages/react/worklet-runtime/__test__/runOnBackground.test.js

@@ -27,29 +27,52 @@ describe('runOnBackground', () => {
     };
     const worklet = {
       _wkltId: 'ctx1',
-      _jsFn: { '_jsFn1': { '_jsFnId': 1 }, '_jsFn2': { '_jsFnId': 2 }, '_jsFn3': { '_jsFnId': 3 } },
+      _jsFn: {
+        '_jsFn1': { '_jsFnId': 1 },
+        '_jsFn2': { '_jsFnId': 2 },
+        '_jsFn3': { '_jsFnId': 3 },
+      },
       _execId: 8,
     };
-    // If the functions are not used in the first screen, they will not be hydrated
+    // If the functions are not used in the first screen, they should not be called
     globalThis.lynxWorkletImpl._hydrateCtx(worklet, firstScreenWorklet);
-    expect(globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.delayedBackgroundFunctionArray.length).toBe(0);
+    expect(
+      globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl
+        .delayedBackgroundFunctionArray.length,
+    ).toBe(0);
+    // functions in `firstScreenWorklet` should be hydrated
+    expect(firstScreenWorklet._jsFn._jsFn1._isFirstScreen).toBe(false);
+    expect(firstScreenWorklet._jsFn._jsFn1._jsFnId).toBe(1);
+    expect(firstScreenWorklet._jsFn._jsFn1._execId).toBe(8);
+    expect(firstScreenWorklet._jsFn._jsFn2._isFirstScreen).toBe(false);
+    expect(firstScreenWorklet._jsFn._jsFn2._jsFnId).toBe(2);
+    expect(firstScreenWorklet._jsFn._jsFn2._execId).toBe(8);
 
     // If the functions are used in the first screen, they will be hydrated
     const task = vi.fn();
     globalThis.registerWorklet('main-thread', 'ctx1', function() {
       const { _jsFn1 } = this._jsFn;
-      globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.delayRunOnBackground(_jsFn1, task);
+      globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.delayRunOnBackground(
+        _jsFn1,
+        task,
+      );
     });
     globalThis.runWorklet(firstScreenWorklet, []);
-    expect(globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.delayedBackgroundFunctionArray).toMatchInlineSnapshot(`
+    expect(
+      globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl
+        .delayedBackgroundFunctionArray,
+    ).toMatchInlineSnapshot(`
       [
         {
           "task": [MockFunction spy],
         },
       ]
     `);
     globalThis.lynxWorkletImpl._hydrateCtx(worklet, firstScreenWorklet);
-    expect(globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.delayedBackgroundFunctionArray).toMatchInlineSnapshot(`
+    expect(
+      globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl
+        .delayedBackgroundFunctionArray,
+    ).toMatchInlineSnapshot(`
       [
         {
           "jsFnHandle": {
@@ -61,7 +84,8 @@ describe('runOnBackground', () => {
       ]
     `);
 
-    globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.runDelayedBackgroundFunctions();
+    globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl
+      .runDelayedBackgroundFunctions();
     expect(task.mock.calls).toMatchInlineSnapshot(`
       [
         [
@@ -70,6 +94,9 @@ describe('runOnBackground', () => {
         ],
       ]
     `);
-    expect(globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl.delayedBackgroundFunctionArray.length).toBe(0);
+    expect(
+      globalThis.lynxWorkletImpl._runOnBackgroundDelayImpl
+        .delayedBackgroundFunctionArray.length,
+    ).toBe(0);
   });
 });

packages/react/worklet-runtime/src/hydrate.ts

@@ -20,8 +20,15 @@ export function hydrateCtx(ctx: Worklet, firstScreenCtx: Worklet): void {
   });
 }
 
-function hydrateCtxImpl(ctx: ClosureValueType, firstScreenCtx: ClosureValueType, execId: number): void {
-  if (!ctx || typeof ctx !== 'object' || !firstScreenCtx || typeof firstScreenCtx !== 'object') return;
+function hydrateCtxImpl(
+  ctx: ClosureValueType,
+  firstScreenCtx: ClosureValueType,
+  execId: number,
+): void {
+  if (
+    !ctx || typeof ctx !== 'object' || !firstScreenCtx
+    || typeof firstScreenCtx !== 'object'
+  ) return;
 
   const ctxObj = ctx as Record<string, ClosureValueType>;
   const firstScreenCtxObj = firstScreenCtx as Record<string, ClosureValueType>;
@@ -33,7 +40,10 @@ function hydrateCtxImpl(ctx: ClosureValueType, firstScreenCtx: ClosureValueType,
   // eslint-disable-next-line @typescript-eslint/no-for-in-array
   for (const key in ctx) {
     if (key === '_wvid') {
-      hydrateMainThreadRef(ctxObj[key] as WorkletRefId, firstScreenCtxObj as unknown as WorkletRefImpl<unknown>);
+      hydrateMainThreadRef(
+        ctxObj[key] as WorkletRefId,
+        firstScreenCtxObj as unknown as WorkletRefImpl<unknown>,
+      );
     } else if (key === '_jsFn') {
       hydrateDelayRunOnBackgroundTasks(
         ctxObj[key] as Record<string, JsFnHandle>,
@@ -57,7 +67,10 @@ function hydrateCtxImpl(ctx: ClosureValueType, firstScreenCtx: ClosureValueType,
  * @param refId The ID of the WorkletRef to hydrate.
  * @param value The new value for the WorkletRef.
  */
-function hydrateMainThreadRef(refId: WorkletRefId, value: WorkletRefImpl<unknown> | { current: unknown }) {
+function hydrateMainThreadRef(
+  refId: WorkletRefId,
+  value: WorkletRefImpl<unknown> | { current: unknown },
+) {
   if ('_initValue' in value) {
     // The ref has not been accessed yet.
     return;
@@ -81,10 +94,16 @@ function hydrateDelayRunOnBackgroundTasks(
     const fnObj = fnObjs[fnName]!;
     const firstScreenFnObj: JsFnHandle | undefined = firstScreenFnObjs[fnName];
     if (!firstScreenFnObj?._delayIndices) {
+      if (firstScreenFnObj) {
+        firstScreenFnObj._isFirstScreen = false;
+        firstScreenFnObj._execId = execId;
+        Object.assign(firstScreenFnObj, fnObj);
+      }
       continue;
     }
     for (const index of firstScreenFnObj._delayIndices) {
-      const details = lynxWorkletImpl!._runOnBackgroundDelayImpl.delayedBackgroundFunctionArray[index]!;
+      const details = lynxWorkletImpl!._runOnBackgroundDelayImpl
+        .delayedBackgroundFunctionArray[index]!;
       fnObj._execId = execId;
       details.jsFnHandle = fnObj;
     }