feat(runtime-core): return onUnmounted in onMounted

packages/runtime-core/__tests__/apiLifecycle.spec.ts

@@ -183,6 +183,36 @@ describe('api: lifecycle hooks', () => {
     expect(fn).toHaveBeenCalledTimes(1)
   })
 
+  it('onMounted return onUnmounted', async () => {
+    const toggle = ref(true)
+    const root = nodeOps.createElement('div')
+    const fn1 = vi.fn(() => fn2)
+    const fn2 = vi.fn()
+
+    const Comp = {
+      setup() {
+        return () => (toggle.value ? h(Child) : null)
+      }
+    }
+
+    const Child = {
+      setup() {
+        onMounted(fn1)
+        return () => h('div')
+      }
+    }
+
+    render(h(Comp), root)
+
+    expect(fn1).toHaveBeenCalledTimes(1)
+    expect(fn2).toHaveBeenCalledTimes(0)
+
+    toggle.value = false
+
+    await nextTick()
+    expect(fn2).toHaveBeenCalledTimes(1)
+  })
+
   it('onBeforeUnmount in onMounted', async () => {
     const toggle = ref(true)
     const root = nodeOps.createElement('div')

packages/runtime-core/src/apiLifecycle.ts

@@ -8,7 +8,7 @@ import {
 import { ComponentPublicInstance } from './componentPublicInstance'
 import { callWithAsyncErrorHandling, ErrorTypeStrings } from './errorHandling'
 import { warn } from './warning'
-import { toHandlerKey } from '@vue/shared'
+import { isPromise, isFunction, toHandlerKey } from '@vue/shared'
 import { DebuggerEvent, pauseTracking, resetTracking } from '@vue/reactivity'
 import { LifecycleHooks } from './enums'
 
@@ -68,7 +68,24 @@ export const createHook =
   (hook: T, target: ComponentInternalInstance | null = currentInstance) =>
     // post-create lifecycle registrations are noops during SSR (except for serverPrefetch)
     (!isInSSRComponentSetup || lifecycle === LifecycleHooks.SERVER_PREFETCH) &&
-    injectHook(lifecycle, (...args: unknown[]) => hook(...args), target)
+    injectHook(
+      lifecycle,
+      (...args: unknown[]) => {
+        const res = hook(...args)
+        if (lifecycle === LifecycleHooks.MOUNTED) {
+          // In the mounted hook, if a function is returned, it will be called in the unmounted hook.
+          if (isPromise(res)) {
+            res
+              .then(_res => isFunction(_res) && onUnmounted(_res, target))
+              .catch(() => {})
+          } else if (isFunction(res)) {
+            onUnmounted(res)
+          }
+        }
+        return res
+      },
+      target
+    )
 
 export const onBeforeMount = createHook(LifecycleHooks.BEFORE_MOUNT)
 export const onMounted = createHook(LifecycleHooks.MOUNTED)