fix(keep-alive): handle KeepAlive teardown for keyed live branches

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

@@ -1765,6 +1765,148 @@ describe('VaporKeepAlive', () => {
     expect(keptAliveScopes.size).toBe(0)
   })
 
+  test('should use live keyed branch when tearing down KeepAlive after same-tick switch', async () => {
+    const show = ref(true)
+    const toggle = ref(true)
+    let keepAlive: any
+    const deactivatedA = vi.fn()
+    const deactivatedB = vi.fn()
+    const unmountedA = vi.fn()
+
+    const Comp = defineVaporComponent({
+      name: 'Comp',
+      props: ['id'],
+      setup(props: any) {
+        const n0 = template('<div> </div>')() as any
+        const n1 = child(n0) as any
+        onBeforeMount(() => {
+          if (props.id === 'b') {
+            show.value = false
+          }
+        })
+        onDeactivated(() => {
+          if (props.id === 'a') {
+            deactivatedA()
+          } else {
+            deactivatedB()
+          }
+        })
+        onUnmounted(() => {
+          if (props.id === 'a') {
+            unmountedA()
+          }
+        })
+        renderEffect(() => setText(n1, props.id))
+        return n0
+      },
+    })
+
+    define({
+      setup() {
+        return createIf(
+          () => show.value,
+          () => {
+            keepAlive = createComponent(VaporKeepAlive, null, {
+              default: () =>
+                createIf(
+                  () => toggle.value,
+                  () => createComponent(Comp, { id: () => 'a' }),
+                  () => createComponent(Comp, { id: () => 'b' }),
+                  undefined,
+                  undefined,
+                  0,
+                ),
+            })
+            return keepAlive
+          },
+        )
+      },
+    }).render()
+
+    await nextTick()
+
+    toggle.value = false
+    await nextTick()
+
+    expect(show.value).toBe(false)
+    expect(deactivatedA).toHaveBeenCalledTimes(1)
+    expect(unmountedA).toHaveBeenCalledTimes(1)
+    expect(deactivatedB).toHaveBeenCalledTimes(1)
+    expect(keepAlive.ctx.getStorageContainer().innerHTML).toBe('')
+  })
+
+  test('should not retain cached keyed branch when current branch is unresolved async during KeepAlive teardown', async () => {
+    const show = ref(true)
+    const toggle = ref(true)
+    let keepAlive: any
+    const deactivatedA = vi.fn()
+    const unmountedA = vi.fn()
+
+    const AsyncComp = defineVaporAsyncComponent(
+      () =>
+        new Promise(() => {
+          // keep unresolved
+        }),
+    )
+
+    const Comp = defineVaporComponent({
+      name: 'Comp',
+      props: ['id'],
+      setup(props: any) {
+        const n0 = template('<div> </div>')() as any
+        const n1 = child(n0) as any
+        onDeactivated(() => {
+          if (props.id === 'a') {
+            deactivatedA()
+          }
+        })
+        onUnmounted(() => {
+          if (props.id === 'a') {
+            unmountedA()
+          }
+        })
+        renderEffect(() => setText(n1, props.id))
+        return n0
+      },
+    })
+
+    define({
+      setup() {
+        return createIf(
+          () => show.value,
+          () => {
+            keepAlive = createComponent(VaporKeepAlive, null, {
+              default: () =>
+                createIf(
+                  () => toggle.value,
+                  () => createComponent(Comp, { id: () => 'a' }),
+                  () => createComponent(AsyncComp),
+                  undefined,
+                  undefined,
+                  0,
+                ),
+            })
+            return keepAlive
+          },
+        )
+      },
+    }).render()
+
+    await nextTick()
+
+    toggle.value = false
+    await nextTick()
+
+    expect(deactivatedA).toHaveBeenCalledTimes(1)
+    expect(unmountedA).toHaveBeenCalledTimes(0)
+
+    show.value = false
+    await nextTick()
+
+    expect(unmountedA).toHaveBeenCalledTimes(1)
+    expect(keepAlive.ctx.getStorageContainer().innerHTML).toBe('')
+  })
+
   test('should recreate composite cache key after max prunes keyed branch entry', async () => {
     const Comp = defineVaporComponent({
       name: 'Comp',

packages/runtime-vapor/src/components/KeepAlive.ts

@@ -401,31 +401,53 @@ const VaporKeepAliveImpl = defineVaporComponent({
 
     onMounted(cacheBlock)
     onUpdated(cacheBlock)
+
+    const getCurrentBlockState = () => {
+      const block = keepAliveInstance.block!
+      const [currentBlock, interop] = getInnerBlock(block)
+      const branchKey =
+        isDynamicFragment(block) && block.keyed
+          ? block.current
+          : currentBranchKey
+
+      return {
+        currentBlock,
+        interop,
+        currentKey:
+          currentBlock && getCacheKey(currentBlock, interop, branchKey),
+      }
+    }
+
     onBeforeUnmount(() => {
-      cache.forEach((cached, key) => {
+      const { currentBlock, interop, currentKey } = getCurrentBlockState()
+      const deactivateCached = (
+        cached: VaporComponentInstance | VaporFragment,
+      ): void => {
+        resetCachedShapeFlag(cached)
         const instance = getInstanceFromCache(cached)
+        if (instance) {
+          const da = instance.da
+          da && queuePostFlushCb(da)
+        }
+      }
 
+      let matched = false
+      cache.forEach((cached, key) => {
         // current instance will be unmounted as part of keep-alive's unmount
-        if (current) {
-          const currentKey = getCacheKey(
-            current,
-            !isVaporComponent(current),
-            currentBranchKey,
-          )
-          if (currentKey === key) {
-            resetCachedShapeFlag(cached)
-            // call deactivated hook
-            if (instance) {
-              const da = instance.da
-              da && queuePostFlushCb(da)
-            }
-            return
-          }
+        if (currentKey === key) {
+          matched = true
+          deactivateCached(cached)
+          return
         }
 
         resetCachedShapeFlag(cached)
         remove(cached, storageContainer)
       })
+
+      if (!matched && currentBlock && isKeptAlive(currentBlock, interop)) {
+        deactivateCached(currentBlock)
+      }
+
       keptAliveScopes.forEach(scope => scope.stop())
       keptAliveScopes.clear()
     })
@@ -520,6 +542,19 @@ const resetCachedShapeFlag = (
   }
 }
 
+function isKeptAlive(
+  cached: VaporComponentInstance | VaporFragment,
+  interop: boolean,
+): boolean {
+  if (interop && isInteropEnabled && isInteropFragment(cached)) {
+    return !!(cached.vnode!.shapeFlag! & ShapeFlags.COMPONENT_SHOULD_KEEP_ALIVE)
+  }
+  return !!(
+    (cached as VaporComponentInstance).shapeFlag! &
+    ShapeFlags.COMPONENT_SHOULD_KEEP_ALIVE
+  )
+}
+
 type InnerBlockResult =
   | [VaporFragment, true]
   | [VaporComponentInstance, false]