fix(loaders): restore context in sequential awaits

Close #2653
Close #2652

Author: posva

packages/router/src/experimental/data-loaders/defineColadaLoader.ts

@@ -506,14 +506,17 @@ export function defineColadaLoader<Data>(
 
     // load ensures there is a pending load
     const promise = entry
-      .pendingLoad!.then(() => {
+      .pendingLoad!.then(() =>
         // nested loaders might wait for all loaders to be ready before setting data
         // so we need to return the staged value if it exists as it will be the latest one
-        return entry.staged === STAGED_NO_VALUE ? ext.data.value : entry.staged
-      })
+        entry.staged === STAGED_NO_VALUE ? ext.data.value : entry.staged
+      )
       // we only want the error if we are nesting the loader
       // otherwise this will end up in "Unhandled promise rejection"
-      .catch(e => (parentEntry ? Promise.reject(e) : null))
+      .catch((e: unknown) => (parentEntry ? Promise.reject(e) : null))
+      // restore the caller's context so that sequential awaits
+      // (e.g. await useOne(); await useTwo()) preserve the parent
+      .finally(() => setCurrentContext(currentContext))
 
     // Restore the context to avoid sequential calls to be nested
     setCurrentContext(currentContext)

packages/router/src/experimental/data-loaders/defineLoader.ts

@@ -397,14 +397,17 @@ export function defineBasicLoader<Data>(
 
     // load ensures there is a pending load
     const promise = entry
-      .pendingLoad!.then(() => {
+      .pendingLoad!.then(() =>
         // nested loaders might wait for all loaders to be ready before setting data
         // so we need to return the staged value if it exists as it will be the latest one
-        return entry.staged === STAGED_NO_VALUE ? data.value : entry.staged
-      })
+        entry.staged === STAGED_NO_VALUE ? data.value : entry.staged
+      )
       // we only want the error if we are nesting the loader
       // otherwise this will end up in "Unhandled promise rejection"
       .catch((e: unknown) => (parentEntry ? Promise.reject(e) : null))
+      // restore the caller's context so that sequential awaits
+      // (e.g. await useOne(); await useTwo()) preserve the parent
+      .finally(() => setCurrentContext(currentContext))
 
     setCurrentContext(currentContext)
     return Object.assign(promise, useDataLoaderResult)

packages/router/src/tests/data-loaders/tester.ts

@@ -1302,6 +1302,65 @@ export function testDefineLoader<Context = void>(
     expect(data.value).toEqual('one,two')
   })
 
+  it('preserves context across sequential nested loader awaits', async () => {
+    const l1 = mockedLoader({ key: 'one' })
+    const l2 = mockedLoader({ key: 'two' })
+
+    const useLoaderThree = loaderFactory({
+      fn: async () => {
+        const one = await l1.loader()
+        const two = await l2.loader()
+        return `${one},${two}`
+      },
+      key: 'three',
+    })
+
+    const router = getRouter()
+    router.addRoute({
+      name: '_test',
+      path: '/fetch',
+      component: defineComponent({
+        setup() {
+          const { data, error, isLoading } = useLoaderThree()
+
+          return { data, error, isLoading }
+        },
+        template: `<p>{{ isLoading }}|{{ error ?? '<no>' }}|{{ data }}</p>`,
+      }),
+      meta: {
+        // All three exported as root loaders (like when all are exported from page component)
+        // The nested loaders are listed BEFORE the parent to trigger the bug:
+        // the guard calls load() for l1/l2 as roots with context [],
+        // so their .finally() restores context to [] instead of the parent's context
+        loaders: [l1.loader, l2.loader, useLoaderThree],
+        nested: { foo: 'bar' },
+      },
+    })
+    const wrapper = mount(RouterViewMock, {
+      global: {
+        plugins: [
+          [DataLoaderPlugin, { router }],
+          ...(plugins?.(customContext!) || []),
+          router,
+        ],
+      },
+    })
+
+    const p = router.push('/fetch')
+    await vi.advanceTimersByTimeAsync(0)
+    l1.resolve('1')
+    await vi.advanceTimersByTimeAsync(0)
+    // If context is lost after l1 resolves, l2.loader() inside useLoaderThree's fn
+    // would throw because router is undefined (no parent context)
+    l2.resolve('2')
+    await vi.advanceTimersByTimeAsync(0)
+    // Navigation should complete without errors
+    await expect(p).resolves.not.toThrow()
+    expect(l1.spy).toHaveBeenCalledTimes(1)
+    expect(l2.spy).toHaveBeenCalledTimes(1)
+    expect(wrapper.text()).toBe('false|<no>|1,2')
+  })
+
   it('fetches once with lazy across components', async () => {
     const router = getRouter()
     router.addRoute({