fix(runtime-vapor): handle KeepAlive template ref cleanup for async branches

packages/runtime-vapor/__tests__/components/KeepAlive.spec.ts

@@ -4,6 +4,7 @@ import {
   nextTick,
   onActivated,
   onBeforeMount,
+  onBeforeUpdate,
   onDeactivated,
   onMounted,
   onUnmounted,
@@ -20,6 +21,7 @@ import {
   child,
   createComponent,
   createDynamicComponent,
+  createFor,
   createIf,
   createSlot,
   createTemplateRefSetter,
@@ -2352,6 +2354,89 @@ describe('VaporKeepAlive', () => {
     expect(instanceRef.value).not.toBe(null)
   })
 
+  test('should keep sibling ref_for entries when switching away from unresolved async KeepAlive branch in v-for', async () => {
+    let resolveAsync: (comp: VaporComponent) => void
+    const AsyncComp = defineVaporAsyncComponent(
+      () =>
+        new Promise(r => {
+          resolveAsync = r as any
+        }),
+    )
+
+    const CompB = defineVaporComponent({
+      name: 'CompB',
+      setup(_, { expose }) {
+        expose({ name: 'B' })
+        return template('<div>B</div>')()
+      },
+    })
+
+    const CompC = defineVaporComponent({
+      name: 'CompC',
+      setup(_, { expose }) {
+        expose({ name: 'C' })
+        return template('<div>C</div>')()
+      },
+    })
+
+    const listRef = ref<any[]>([])
+    const toggle = ref(true)
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        const items = ['async', 'stable']
+        return createFor(
+          () => items,
+          item => {
+            if (item.value === 'async') {
+              return createComponent(VaporKeepAlive, null, {
+                default: () =>
+                  createIf(
+                    () => toggle.value,
+                    () => {
+                      const comp = createComponent(AsyncComp)
+                      setRef(comp, listRef as any, true)
+                      return comp
+                    },
+                    () => {
+                      const comp = createComponent(CompB)
+                      setRef(comp, listRef as any, true)
+                      return comp
+                    },
+                  ),
+              })
+            }
+
+            const comp = createComponent(CompC)
+            setRef(comp, listRef as any, true)
+            return comp
+          },
+          item => item,
+        )
+      },
+    }).render()
+
+    await nextTick()
+    expect(listRef.value).toHaveLength(1)
+    expect(listRef.value[0]).toMatchObject({ name: 'C' })
+
+    toggle.value = false
+    await nextTick()
+
+    expect(listRef.value).toHaveLength(2)
+    expect(listRef.value.map(item => item.name).sort()).toEqual(['B', 'C'])
+
+    resolveAsync!(
+      defineVaporComponent({
+        name: 'ResolvedA',
+        setup() {
+          return template('<div>A</div>')()
+        },
+      }),
+    )
+  })
+
   test('should clear old ref when switching KeepAlive branches', async () => {
     const CompA = defineVaporComponent({
       name: 'CompA',
@@ -2409,4 +2494,221 @@ describe('VaporKeepAlive', () => {
     expect(refA.value).not.toBe(null)
     expect(refB.value).toBe(null)
   })
+
+  test('should not restore stale ref when current KeepAlive branch rerenders and then switches', async () => {
+    const refresh = ref(0)
+    const refA = ref<any>(null)
+    const refB = ref<any>(null)
+    const current = shallowRef<VaporComponent>()
+    let switched = false
+
+    const CompB = defineVaporComponent({
+      name: 'CompB',
+      setup(_, { expose }) {
+        expose({ name: 'B' })
+        return template('<div>B</div>')()
+      },
+    })
+
+    const CompA = defineVaporComponent({
+      name: 'CompA',
+      props: ['n'],
+      setup(props, { expose }) {
+        expose({ name: 'A' })
+
+        onBeforeUpdate(() => {
+          if (!switched && props.n === 1) {
+            switched = true
+            current.value = CompB
+          }
+        })
+
+        const n0 = template('<div> </div>')() as any
+        const x0 = child(n0) as any
+        renderEffect(() => setText(x0, `A${props.n}`))
+        return n0
+      },
+    })
+
+    current.value = CompA
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        return createComponent(VaporKeepAlive, null, {
+          default: () =>
+            createIf(
+              () => current.value === CompA,
+              () => {
+                const comp = createComponent(CompA, { n: () => refresh.value })
+                setRef(comp, refA)
+                return comp
+              },
+              () => {
+                const comp = createComponent(CompB)
+                setRef(comp, refB)
+                return comp
+              },
+            ),
+        })
+      },
+    }).render()
+
+    await nextTick()
+    expect(refA.value).toMatchObject({ name: 'A' })
+    expect(refB.value).toBe(null)
+
+    refresh.value = 1
+    await nextTick()
+
+    expect(refA.value).toBe(null)
+    expect(refB.value).toMatchObject({ name: 'B' })
+  })
+
+  test('should not restore stale ref when resolved async KeepAlive branch switches away in the same tick', async () => {
+    const timeout = (n: number = 0) => new Promise(r => setTimeout(r, n))
+
+    let resolveAsync: (comp: VaporComponent) => void
+    const refA = ref<any>(null)
+    const refB = ref<any>(null)
+    const current = shallowRef<VaporComponent>()
+
+    const CompB = defineVaporComponent({
+      name: 'CompB',
+      setup(_, { expose }) {
+        expose({ name: 'B' })
+        return template('<div>B</div>')()
+      },
+    })
+
+    const AsyncComp = defineVaporAsyncComponent(
+      () =>
+        new Promise(r => {
+          resolveAsync = r as any
+        }),
+    )
+
+    current.value = AsyncComp
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        return createComponent(VaporKeepAlive, null, {
+          default: () =>
+            createIf(
+              () => current.value === AsyncComp,
+              () => {
+                const comp = createComponent(AsyncComp)
+                setRef(comp, refA)
+                return comp
+              },
+              () => {
+                const comp = createComponent(CompB)
+                setRef(comp, refB)
+                return comp
+              },
+            ),
+        })
+      },
+    }).render()
+
+    await nextTick()
+    expect(refA.value).toBe(null)
+    expect(refB.value).toBe(null)
+
+    resolveAsync!(
+      defineVaporComponent({
+        name: 'ResolvedA',
+        setup(_, { expose }) {
+          expose({ name: 'A' })
+          onBeforeMount(() => {
+            current.value = CompB
+          })
+          return template('<div>A</div>')()
+        },
+      }),
+    )
+
+    await timeout()
+    await nextTick()
+
+    expect(refA.value).toBe(null)
+    expect(refB.value).toMatchObject({ name: 'B' })
+  })
+
+  test('should clear function ref when resolved async KeepAlive branch switches away in the same tick', async () => {
+    const timeout = (n: number = 0) => new Promise(r => setTimeout(r, n))
+
+    let resolveAsync: (comp: VaporComponent) => void
+    const fnA = vi.fn()
+    const fnB = vi.fn()
+    const current = shallowRef<VaporComponent>()
+
+    const CompB = defineVaporComponent({
+      name: 'CompB',
+      setup(_, { expose }) {
+        expose({ name: 'B' })
+        return template('<div>B</div>')()
+      },
+    })
+
+    const AsyncComp = defineVaporAsyncComponent(
+      () =>
+        new Promise(r => {
+          resolveAsync = r as any
+        }),
+    )
+
+    current.value = AsyncComp
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        return createComponent(VaporKeepAlive, null, {
+          default: () =>
+            createIf(
+              () => current.value === AsyncComp,
+              () => {
+                const comp = createComponent(AsyncComp)
+                setRef(comp, fnA as any)
+                return comp
+              },
+              () => {
+                const comp = createComponent(CompB)
+                setRef(comp, fnB as any)
+                return comp
+              },
+            ),
+        })
+      },
+    }).render()
+
+    await nextTick()
+    expect(fnA).toHaveBeenCalled()
+    expect(fnA.mock.calls[0][0]).toBe(null)
+    expect(fnB).not.toHaveBeenCalled()
+
+    resolveAsync!(
+      defineVaporComponent({
+        name: 'ResolvedA',
+        setup(_, { expose }) {
+          expose({ name: 'A' })
+          onBeforeMount(() => {
+            current.value = CompB
+          })
+          return template('<div>A</div>')()
+        },
+      }),
+    )
+
+    await timeout()
+    await nextTick()
+
+    const fnAArgs = fnA.mock.calls.map(args => args[0])
+    expect(fnAArgs.some(arg => arg && arg.name === 'A')).toBe(true)
+    expect(fnAArgs[fnAArgs.length - 1]).toBe(null)
+    expect(fnB.mock.calls[fnB.mock.calls.length - 1][0]).toMatchObject({
+      name: 'B',
+    })
+  })
 })

packages/runtime-vapor/src/apiTemplateRef.ts

@@ -33,7 +33,12 @@ import {
   isFragment,
 } from './fragment'
 import { isInteropEnabled } from './vdomInteropState'
-import { refCleanups } from './refCleanup'
+import {
+  type RefCleanupState,
+  invalidatePendingRef,
+  refCleanups,
+  unsetRef,
+} from './refCleanup'
 
 export type NodeRef =
   | string
@@ -52,11 +57,12 @@ export type setRefFn = (
   refKey?: string,
 ) => NodeRef | undefined
 
-function ensureCleanup(el: RefEl): { fn: () => void } {
+function ensureCleanup(el: RefEl): RefCleanupState {
   let cleanupRef = refCleanups.get(el)
   if (!cleanupRef) {
     refCleanups.set(el, (cleanupRef = { fn: NOOP }))
     onScopeDispose(() => {
+      invalidatePendingRef(el)
       cleanupRef!.fn()
       refCleanups.delete(el)
     })
@@ -105,16 +111,6 @@ function setRef(
 ): NodeRef | undefined {
   if (!instance || instance.isUnmounted) return
 
-  // async component
-  if (isVaporComponent(el) && isAsyncWrapper(el)) {
-    // resolved: set ref to the inner component
-    if (el.type.__asyncResolved) {
-      el = (el.block as DynamicFragment).nodes as VaporComponentInstance
-    }
-    // unresolved: el stays as async wrapper, getRefValue returns null
-    // → ref will be cleared through normal setRef logic below
-  }
-
   const setupState: any = __DEV__ ? instance.setupState || {} : null
   const refValue = getRefValue(el)
 
@@ -151,6 +147,7 @@ function setRef(
 
   // dynamic ref changed. unset old ref
   if (oldRef != null && oldRef !== ref) {
+    invalidatePendingRef(el)
     if (isString(oldRef)) {
       refs[oldRef] = null
       if (__DEV__ && canSetSetupRef(oldRef)) {
@@ -172,8 +169,7 @@ function setRef(
       ])
     } else if (refFor) {
       // For dynamic ref-for branches, remove only this branch's previous value.
-      const cleanup = refCleanups.get(el)
-      if (cleanup) cleanup.fn()
+      unsetRef(el)
     }
   }
 
@@ -241,11 +237,12 @@ function setRef(
           warn('Invalid template ref type:', ref, `(${typeof ref})`)
         }
       }
-      queuePostFlushCb(doSet, -1)
-
-      ensureCleanup(el).fn = () => {
-        if (isArray(existing)) {
-          remove(existing, refValue)
+      const cleanup = ensureCleanup(el)
+      cleanup.fn = () => {
+        if (refFor) {
+          if (isArray(existing)) {
+            remove(existing, refValue)
+          }
         } else if (_isString) {
           refs[ref] = null
           if (__DEV__ && canSetSetupRef(ref)) {
@@ -256,6 +253,18 @@ function setRef(
           if (refKey) refs[refKey] = null
         }
       }
+
+      invalidatePendingRef(el)
+      if (refValue != null) {
+        const job: SchedulerJob = () => {
+          doSet()
+          if (cleanup.job === job) cleanup.job = undefined
+        }
+        cleanup.job = job
+        queuePostFlushCb(job, -1)
+      } else {
+        doSet()
+      }
     } else if (__DEV__) {
       warn('Invalid template ref type:', ref, `(${typeof ref})`)
     }
@@ -265,8 +274,11 @@ function setRef(
 
 const getRefValue = (el: RefEl) => {
   if (isVaporComponent(el)) {
-    // unresolved async wrapper: return null so ref gets cleared
-    if (isAsyncWrapper(el) && !el.type.__asyncResolved) return null
+    if (isAsyncWrapper(el)) {
+      // unresolved async wrapper: return null so ref gets cleared
+      if (!el.type.__asyncResolved) return null
+      return getRefValue((el.block as DynamicFragment).nodes as RefEl)
+    }
     return getExposed(el) || el
   } else if (isTeleportFragment(el)) {
     return null

packages/runtime-vapor/src/refCleanup.ts

@@ -1,10 +1,24 @@
 import type { Block } from './block'
+import { type SchedulerJob, SchedulerJobFlags } from '@vue/runtime-dom'
+
+export interface RefCleanupState {
+  fn: () => void
+  job?: SchedulerJob
+}
 
 /**
  * Stores ref cleanup functions keyed by the element/component they are set on.
  * Shared between apiTemplateRef.ts (writes) and KeepAlive deactivate (reads).
  */
-export const refCleanups: WeakMap<Block, { fn: () => void }> = new WeakMap()
+export const refCleanups: WeakMap<Block, RefCleanupState> = new WeakMap()
+
+export function invalidatePendingRef(el: Block): void {
+  const c = refCleanups.get(el)
+  if (c && c.job) {
+    c.job.flags = c.job.flags! | SchedulerJobFlags.DISPOSED
+    c.job = undefined
+  }
+}
 
 /**
  * Synchronously clear the ref for an element being deactivated by KeepAlive.
@@ -13,6 +27,7 @@ export const refCleanups: WeakMap<Block, { fn: () => void }> = new WeakMap()
  * this explicit sync cleanup path.
  */
 export function unsetRef(el: Block): void {
+  invalidatePendingRef(el)
   const c = refCleanups.get(el)
   if (c) c.fn()
 }