refactor(runtime-vapor): simplify KeepAlive cache key resolution

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

@@ -1705,10 +1705,20 @@ describe('VaporKeepAlive', () => {
   })
 
   test('should stop branch scope when cache entry is pruned (keyed branches)', async () => {
+    const mountedA = vi.fn()
+    const mountedB = vi.fn()
+
     const Comp = defineVaporComponent({
       name: 'Comp',
       props: ['id'],
       setup(props: any) {
+        onMounted(() => {
+          if (props.id === 'a') {
+            mountedA()
+          } else {
+            mountedB()
+          }
+        })
         return template('<div></div>')()
       },
     })
@@ -1750,14 +1760,16 @@ describe('VaporKeepAlive', () => {
     await nextTick()
 
     // both branches should be independently cached
-    // keyed branches use branchKey to distinguish same-type components
     expect(cache.size).toBe(2)
-    expect(keptAliveScopes.size).toBe(2)
+    expect(mountedA).toHaveBeenCalledTimes(1)
+    expect(mountedB).toHaveBeenCalledTimes(1)
 
     // switch back to branch A
     toggle.value = true
     await nextTick()
     expect(cache.size).toBe(2)
+    expect(mountedA).toHaveBeenCalledTimes(1)
+    expect(mountedB).toHaveBeenCalledTimes(1)
 
     // prune by excluding Comp — all entries should be cleaned
     exclude.value = 'Comp'
@@ -1908,7 +1920,7 @@ describe('VaporKeepAlive', () => {
     expect(keepAlive.ctx.getStorageContainer().innerHTML).toBe('')
   })
 
-  test('should recreate composite cache key after max prunes keyed branch entry', async () => {
+  test('should recreate cached entry while preserving branch cache key after max prunes keyed branch entry', async () => {
     const Comp = defineVaporComponent({
       name: 'Comp',
       props: ['id'],
@@ -1945,6 +1957,7 @@ describe('VaporKeepAlive', () => {
     await nextTick()
     expect(cache.size).toBe(1)
     const keyA1 = Array.from(cache.keys())[0]
+    const cachedA1 = cache.get(keyA1)
 
     toggle.value = false
     await nextTick()
@@ -1956,12 +1969,13 @@ describe('VaporKeepAlive', () => {
     await nextTick()
     expect(cache.size).toBe(1)
     const keyA2 = Array.from(cache.keys())[0]
+    const cachedA2 = cache.get(keyA2)
 
-    expect((keyA1 as any).branchKey).toBe((keyA2 as any).branchKey)
-    expect(keyA2).not.toBe(keyA1)
+    expect(keyA2).toBe(keyA1)
+    expect(cachedA2).not.toBe(cachedA1)
   })
 
-  test('should recreate composite cache key after KeepAlive hmr rerender', async () => {
+  test('should recreate cached entry while preserving branch cache key after KeepAlive hmr rerender', async () => {
     const Comp = defineVaporComponent({
       name: 'Comp',
       props: ['id'],
@@ -1993,17 +2007,19 @@ describe('VaporKeepAlive', () => {
     await nextTick()
     expect(cache.size).toBe(1)
     const keyA1 = Array.from(cache.keys())[0]
+    const cachedA1 = cache.get(keyA1)
 
     keepAliveInstance.hmrRerender!()
     await nextTick()
 
     expect(cache.size).toBe(1)
     const keyA2 = Array.from(cache.keys())[0]
-    expect((keyA1 as any).branchKey).toBe((keyA2 as any).branchKey)
-    expect(keyA2).not.toBe(keyA1)
+    const cachedA2 = cache.get(keyA2)
+    expect(keyA2).toBe(keyA1)
+    expect(cachedA2).not.toBe(cachedA1)
   })
 
-  test('should not retain composite key entries for uncached keyed branches', async () => {
+  test('should not create cache entries for uncached keyed branches', async () => {
     const Comp = defineVaporComponent({
       name: 'Comp',
       setup() {
@@ -2013,76 +2029,42 @@ describe('VaporKeepAlive', () => {
 
     const include = ref('OtherComp')
     const routeKey = ref('a')
+    const { instance } = define({
+      setup() {
+        return createComponent(
+          VaporKeepAlive,
+          { include: () => include.value },
+          {
+            default: () =>
+              createKeyedFragment(
+                () => routeKey.value,
+                () => createComponent(Comp),
+              ),
+          },
+        )
+      },
+    }).render()
 
-    const rawSet = Map.prototype.set
-    const rawDelete = Map.prototype.delete
-    const compositeMaps = new WeakSet<Map<any, any>>()
-    let compositeSetCount = 0
-    let compositeDeleteCount = 0
-
-    const setSpy = vi.spyOn(Map.prototype, 'set').mockImplementation(function (
-      this: Map<any, any>,
-      key: any,
-      value: any,
-    ) {
-      if (
-        value &&
-        typeof value === 'object' &&
-        'branchKey' in value &&
-        'type' in value &&
-        value.type === Comp
-      ) {
-        compositeMaps.add(this)
-        compositeSetCount++
-      }
-      return rawSet.call(this, key, value)
-    })
-
-    const deleteSpy = vi
-      .spyOn(Map.prototype, 'delete')
-      .mockImplementation(function (this: Map<any, any>, key: any) {
-        if (compositeMaps.has(this)) {
-          compositeDeleteCount++
-        }
-        return rawDelete.call(this, key)
-      })
-
-    try {
-      const { instance } = define({
-        setup() {
-          return createComponent(
-            VaporKeepAlive,
-            { include: () => include.value },
-            {
-              default: () =>
-                createKeyedFragment(
-                  () => routeKey.value,
-                  () => createComponent(Comp),
-                ),
-            },
-          )
-        },
-      }).render()
-
-      const keepAliveInstance = instance!.block as any
-      const cache = keepAliveInstance.__v_cache as Map<any, any>
-
-      await nextTick()
-      expect(cache.size).toBe(0)
+    const keepAliveInstance = instance!.block as any
+    const cache = keepAliveInstance.__v_cache as Map<any, any>
+    const keptAliveScopes = keepAliveInstance.__v_keptAliveScopes as Map<
+      any,
+      any
+    >
 
-      routeKey.value = 'b'
-      await nextTick()
-      expect(cache.size).toBe(0)
+    await nextTick()
+    expect(cache.size).toBe(0)
+    expect(keptAliveScopes.size).toBe(0)
 
-      routeKey.value = 'c'
-      await nextTick()
-      expect(cache.size).toBe(0)
+    routeKey.value = 'b'
+    await nextTick()
+    expect(cache.size).toBe(0)
+    expect(keptAliveScopes.size).toBe(0)
 
-      expect(compositeSetCount - compositeDeleteCount).toBeLessThanOrEqual(1)
-    } finally {
-      setSpy.mockRestore()
-      deleteSpy.mockRestore()
-    }
+    routeKey.value = 'c'
+    await nextTick()
+    expect(cache.size).toBe(0)
+    expect(keptAliveScopes.size).toBe(0)
   })
 
   test('handle error in async onActivated', async () => {

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

@@ -30,13 +30,7 @@ import {
   type DefineVaporComponent,
   defineVaporComponent,
 } from '../apiDefineComponent'
-import {
-  ShapeFlags,
-  invokeArrayFns,
-  isArray,
-  isFunction,
-  isObject,
-} from '@vue/shared'
+import { ShapeFlags, invokeArrayFns, isArray } from '@vue/shared'
 import { createElement } from '../dom/node'
 import { unsetRef } from '../refCleanup'
 import { type VaporFragment, isDynamicFragment, isFragment } from '../fragment'
@@ -46,9 +40,8 @@ import { isInteropEnabled } from '../vdomInteropState'
 export interface VaporKeepAliveContext {
   processShapeFlag(block: Block): CacheKey | false
   cacheBlock(): void
-  cacheScope(cacheKey: CacheKey, branchKey: any, scope: EffectScope): void
-  getScope(branchKey: any): EffectScope | undefined
-  setCurrentBranchKey(key: any): any
+  cacheScope(cacheKey: CacheKey, scopeLookupKey: any, scope: EffectScope): void
+  getScope(key: any): EffectScope | undefined
 }
 
 export let currentKeepAliveCtx: VaporKeepAliveContext | null = null
@@ -63,6 +56,17 @@ export function setCurrentKeepAliveCtx(
   }
 }
 
+let currentCacheKey: any | undefined
+export function withCurrentCacheKey<T>(key: any, fn: () => T): T {
+  const prev = currentCacheKey
+  currentCacheKey = key
+  try {
+    return fn()
+  } finally {
+    currentCacheKey = prev
+  }
+}
+
 export interface KeepAliveInstance extends VaporComponentInstance {
   ctx: {
     activate: (
@@ -79,84 +83,9 @@ export interface KeepAliveInstance extends VaporComponentInstance {
   }
 }
 
-const COMPOSITE_KEY: unique symbol = Symbol('keepAliveCompositeKey')
-
-type BaseCacheKey = VaporComponent | VNode['type']
-type CacheKey = BaseCacheKey | CompositeKey
+type CacheKey = any
 type Cache = Map<CacheKey, VaporComponentInstance | VaporFragment>
 type Keys = Set<CacheKey>
-type CompositeKey = {
-  type: BaseCacheKey
-  branchKey: any
-  [COMPOSITE_KEY]: true
-}
-
-// Returns a stable composite key object for a given (type, branchKey) pair.
-// Caches are passed as parameters (per KeepAlive instance) to avoid
-// module-level Map leaks for primitive type keys (e.g. string tag names
-// from VDOM interop).
-function getOrCreateCompositeKey(
-  type: BaseCacheKey,
-  branchKey: any,
-  compositeKeyCache: WeakMap<object, Map<any, CompositeKey>>,
-  compositeKeyCachePrimitive: Map<any, Map<any, CompositeKey>>,
-): CacheKey {
-  const isObjectType = isObject(type) || isFunction(type)
-  const perType = isObjectType
-    ? compositeKeyCache.get(type) || new Map<any, CompositeKey>()
-    : compositeKeyCachePrimitive.get(type) || new Map<any, CompositeKey>()
-  if (isObjectType) {
-    if (!compositeKeyCache.has(type)) compositeKeyCache.set(type, perType)
-  } else if (!compositeKeyCachePrimitive.has(type)) {
-    compositeKeyCachePrimitive.set(type, perType)
-  }
-
-  let composite = perType.get(branchKey)
-  if (!composite) {
-    composite = { type, branchKey, [COMPOSITE_KEY]: true }
-    perType.set(branchKey, composite)
-  }
-  return composite
-}
-
-function getCompositeKey(
-  type: BaseCacheKey,
-  branchKey: any,
-  compositeKeyCache: WeakMap<object, Map<any, CompositeKey>>,
-  compositeKeyCachePrimitive: Map<any, Map<any, CompositeKey>>,
-): CacheKey | undefined {
-  const isObjectType = isObject(type) || isFunction(type)
-  const perType = isObjectType
-    ? compositeKeyCache.get(type)
-    : compositeKeyCachePrimitive.get(type)
-
-  return perType && perType.get(branchKey)
-}
-
-function deleteCompositeKey(
-  key: CacheKey,
-  compositeKeyCache: WeakMap<object, Map<any, CompositeKey>>,
-  compositeKeyCachePrimitive: Map<any, Map<any, CompositeKey>>,
-): void {
-  if (!isObject(key) || !(COMPOSITE_KEY in key)) return
-
-  const { type, branchKey } = key as CompositeKey
-  const isObjectType = isObject(type) || isFunction(type)
-  const perType = isObjectType
-    ? compositeKeyCache.get(type)
-    : compositeKeyCachePrimitive.get(type)
-
-  if (!perType || perType.get(branchKey) !== key) return
-
-  perType.delete(branchKey)
-  if (perType.size) return
-
-  if (isObjectType) {
-    compositeKeyCache.delete(type)
-  } else {
-    compositeKeyCachePrimitive.delete(type)
-  }
-}
 
 const VaporKeepAliveImpl = defineVaporComponent({
   name: 'VaporKeepAlive',
@@ -185,80 +114,28 @@ const VaporKeepAliveImpl = defineVaporComponent({
     const keys: Keys = new Set()
     const storageContainer = createElement('div')
     const keptAliveScopes = new Map<any, EffectScope>()
-    // Per-instance composite key caches for generating stable cache keys.
-    // Using WeakMap for object types (auto GC) and Map for primitive types
-    // (e.g. string tag names from VDOM interop). Both are per-instance so
-    // they are cleaned up when the KeepAlive instance is destroyed.
-    const compositeKeyCache = new WeakMap<object, Map<any, CompositeKey>>()
-    const compositeKeyCachePrimitive = new Map<any, Map<any, CompositeKey>>()
-
-    const resolveKey = (
-      type: BaseCacheKey,
-      key?: any,
-      branchKey?: any,
-    ): CacheKey => {
-      if (key != null) {
-        return getOrCreateCompositeKey(
-          type,
-          key,
-          compositeKeyCache,
-          compositeKeyCachePrimitive,
-        )
-      }
-      if (branchKey !== undefined) {
-        return getOrCreateCompositeKey(
-          type,
-          branchKey,
-          compositeKeyCache,
-          compositeKeyCachePrimitive,
-        )
-      }
-      return type as CacheKey
-    }
-
-    const resolveKeyForLookup = (
-      type: BaseCacheKey,
-      key?: any,
-      branchKey?: any,
-    ): CacheKey | undefined => {
-      if (key != null) {
-        return getCompositeKey(
-          type,
-          key,
-          compositeKeyCache,
-          compositeKeyCachePrimitive,
-        )
-      }
-      if (branchKey !== undefined) {
-        return getCompositeKey(
-          type,
-          branchKey,
-          compositeKeyCache,
-          compositeKeyCachePrimitive,
-        )
-      }
-      return type as CacheKey
-    }
 
-    const getCacheKey = (
+    const resolveCacheKeyFromBlock = (
       block: VaporComponentInstance | VaporFragment,
       interop: boolean,
-      branchKey?: any,
+      branchKey = currentCacheKey,
     ): CacheKey => {
       if (interop && isInteropEnabled) {
         const frag = block as VaporFragment
-        return resolveKey(
-          frag.vnode!.type,
-          frag.$key !== undefined ? frag.$key : frag.vnode!.key,
-          branchKey,
+        return (
+          (frag.$key !== undefined
+            ? frag.$key
+            : (frag.vnode!.key ?? branchKey)) ?? frag.vnode!.type
         )
       }
-      const instance = block as VaporComponentInstance
-      return resolveKey(instance.type, instance.key, branchKey)
+
+      return (
+        (block as VaporComponentInstance).key ??
+        branchKey ??
+        (block as VaporComponentInstance).type
+      )
     }
-    // Track active keyed DynamicFragment branch key so KeepAlive can combine
-    // branch key + component type into a stable isolated cache key.
-    let currentBranchKey: any | undefined
+
     let current: VaporComponentInstance | VaporFragment | undefined
 
     if (__DEV__ || __FEATURE_PROD_DEVTOOLS__) {
@@ -272,10 +149,7 @@ const VaporKeepAliveImpl = defineVaporComponent({
       const rerender = keepAliveInstance.hmrRerender
       keepAliveInstance.hmrRerender = () => {
         keepAliveInstance.exposed = null
-        cache.forEach((cached, key) => {
-          resetCachedShapeFlag(cached)
-          deleteCompositeKey(key, compositeKeyCache, compositeKeyCachePrimitive)
-        })
+        cache.forEach(cached => resetCachedShapeFlag(cached))
         cache.clear()
         keys.clear()
         keptAliveScopes.forEach(scope => scope.stop())
@@ -290,15 +164,9 @@ const VaporKeepAliveImpl = defineVaporComponent({
       getStorageContainer: () => storageContainer,
       getCachedComponent: (comp, key) => {
         if (isInteropEnabled && isVNode(comp)) {
-          const cacheKey = resolveKeyForLookup(
-            comp.type,
-            comp.key,
-            currentBranchKey,
-          )
-          return cacheKey === undefined ? undefined : cache.get(cacheKey)
+          return cache.get(comp.key ?? currentCacheKey ?? comp.type)
         }
-        const cacheKey = resolveKeyForLookup(comp, key, currentBranchKey)
-        return cacheKey === undefined ? undefined : cache.get(cacheKey)
+        return cache.get(key ?? currentCacheKey ?? comp)
       },
       activate: (instance, parentNode, anchor) => {
         current = instance
@@ -347,16 +215,13 @@ const VaporKeepAliveImpl = defineVaporComponent({
         ) {
           return
         }
-        // For keyed DynamicFragment, read branch key from the fragment
-        // since currentBranchKey is already restored at lifecycle hook time
-        if (block.keyed) {
-          currentBranchKey = block.current
-        }
       }
       const [innerBlock, interop] = getInnerBlock(block)
       if (!innerBlock || !shouldCache(innerBlock, props, interop)) return
+      const branchKey =
+        isDynamicFragment(block) && block.keyed ? block.current : undefined
       innerCacheBlock(
-        getCacheKey(innerBlock, interop, currentBranchKey),
+        resolveCacheKeyFromBlock(innerBlock, interop, branchKey),
         innerBlock,
       )
     }
@@ -366,14 +231,14 @@ const VaporKeepAliveImpl = defineVaporComponent({
       if (!innerBlock || !shouldCache(innerBlock!, props, interop)) return false
 
       if (interop && isInteropEnabled) {
-        const cacheKey = getCacheKey(innerBlock, true, currentBranchKey)
+        const cacheKey = resolveCacheKeyFromBlock(innerBlock, true)
         if (cache.has(cacheKey)) {
           innerBlock.vnode!.shapeFlag! |= ShapeFlags.COMPONENT_KEPT_ALIVE
         }
         innerBlock.vnode!.shapeFlag! |= ShapeFlags.COMPONENT_SHOULD_KEEP_ALIVE
         return cacheKey
       } else {
-        const cacheKey = getCacheKey(innerBlock, false, currentBranchKey)
+        const cacheKey = resolveCacheKeyFromBlock(innerBlock, false)
         if (cache.has(cacheKey)) {
           ;(innerBlock as VaporComponentInstance)!.shapeFlag! |=
             ShapeFlags.COMPONENT_KEPT_ALIVE
@@ -400,7 +265,7 @@ const VaporKeepAliveImpl = defineVaporComponent({
     }
 
     // delete scope from keptAliveScopes by one key,
-    // also removes the paired entry (cache key ↔ branch key)
+    // also removes the paired entry (cache key ↔ scope lookup key)
     const deleteScope = (key: any): EffectScope | undefined => {
       const scope = keptAliveScopes.get(key)
       if (scope) {
@@ -427,7 +292,6 @@ const VaporKeepAliveImpl = defineVaporComponent({
       }
       cache.delete(key)
       keys.delete(key)
-      deleteCompositeKey(key, compositeKeyCache, compositeKeyCachePrimitive)
       const scope = deleteScope(key)
       if (scope) scope.stop()
     }
@@ -452,13 +316,14 @@ const VaporKeepAliveImpl = defineVaporComponent({
       const branchKey =
         isDynamicFragment(block) && block.keyed
           ? block.current
-          : currentBranchKey
+          : currentCacheKey
 
       return {
         currentBlock,
         interop,
         currentKey:
-          currentBlock && getCacheKey(currentBlock, interop, branchKey),
+          currentBlock &&
+          resolveCacheKeyFromBlock(currentBlock, interop, branchKey),
       }
     }
 
@@ -499,23 +364,27 @@ const VaporKeepAliveImpl = defineVaporComponent({
     const keepAliveCtx: VaporKeepAliveContext = {
       processShapeFlag,
       cacheBlock,
-      cacheScope(cacheKey, branchKey, scope) {
-        // store under both keys so the scope can be looked up by:
-        // - cache key: for cleanup in pruneCacheEntry
-        // - branch key: for reuse in getScope (before the cache key is known)
+      cacheScope(cacheKey, scopeLookupKey, scope) {
+        // remove stale scope
+        const prevScope = keptAliveScopes.get(cacheKey)
+        if (prevScope && prevScope !== scope) {
+          const staleScope = deleteScope(cacheKey)
+          if (staleScope) {
+            staleScope.stop()
+          }
+        }
+
+        // cacheKey is used for cleanup in pruneCacheEntry.
+        // scopeLookupKey is still needed for getScope() before a new block
+        // exists, but keyed branches may resolve to the same effective cacheKey.
         keptAliveScopes.set(cacheKey, scope)
-        keptAliveScopes.set(branchKey, scope)
-      },
-      getScope(branchKey) {
-        return deleteScope(branchKey)
-      },
-      setCurrentBranchKey(key) {
-        try {
-          return currentBranchKey
-        } finally {
-          currentBranchKey = key
+        if (scopeLookupKey !== cacheKey) {
+          keptAliveScopes.set(scopeLookupKey, scope)
         }
       },
+      getScope(key) {
+        return deleteScope(key)
+      },
     }
 
     const prevCtx = setCurrentKeepAliveCtx(keepAliveCtx)

packages/runtime-vapor/src/fragment.ts

@@ -41,6 +41,7 @@ import {
   type VaporKeepAliveContext,
   currentKeepAliveCtx,
   setCurrentKeepAliveCtx,
+  withCurrentCacheKey,
 } from './components/KeepAlive'
 
 export class VaporFragment<
@@ -136,14 +137,14 @@ export class DynamicFragment extends VaporFragment {
     // currently leaving: defer mounting the next branch until
     // the leave finishes.
     if (transition && transition.state.isLeaving) {
+      // Track the latest target key immediately so repeated updates during
+      // leave keep overwriting the pending branch instead of reviving stale
+      // keys when the deferred render finally runs.
       this.current = key
       this.pending = { render, key }
       return
     }
 
-    const prevKey = this.current
-    this.current = key
-
     const instance = currentInstance
     const prevSub = setActiveSub()
     const parent = isHydrating ? null : this.anchor.parentNode
@@ -154,18 +155,17 @@ export class DynamicFragment extends VaporFragment {
 
       // if keepAliveCtx exists and processShapeFlag returns a cache key,
       // cache the scope and retain it.
-      // For keyed fragments, temporarily set branchKey to prevKey so
-      // getCacheKey generates the correct composite key for the OLD block.
-      const needBranchKey = keepAliveCtx && this.keyed
-      const prevBranchKey = needBranchKey
-        ? keepAliveCtx.setCurrentBranchKey(prevKey)
-        : undefined
-      const cacheKey = keepAliveCtx && keepAliveCtx.processShapeFlag(this.nodes)
+      const cacheKey = keepAliveCtx
+        ? this.keyed
+          ? withCurrentCacheKey(this.current, () =>
+              keepAliveCtx.processShapeFlag(this.nodes),
+            )
+          : keepAliveCtx.processShapeFlag(this.nodes)
+        : false
       if (cacheKey) {
-        keepAliveCtx!.cacheScope(cacheKey, prevKey, this.scope)
+        keepAliveCtx!.cacheScope(cacheKey, this.current, this.scope)
         retainScope = true
       }
-      if (needBranchKey) keepAliveCtx.setCurrentBranchKey(prevBranchKey)
 
       if (!retainScope) {
         this.scope.stop()
@@ -190,10 +190,9 @@ export class DynamicFragment extends VaporFragment {
             const pending = this.pending
             if (pending) {
               this.pending = undefined
-              this.current = pending.key
-              this.renderBranch(pending.render, transition, parent)
+              this.renderBranch(pending.render, transition, parent, pending.key)
             } else {
-              this.renderBranch(render, transition, parent)
+              this.renderBranch(render, transition, parent, key)
             }
           } finally {
             setCurrentInstance(...prevInstance)
@@ -209,7 +208,7 @@ export class DynamicFragment extends VaporFragment {
       }
     }
 
-    this.renderBranch(render, transition, parent)
+    this.renderBranch(render, transition, parent, key)
     setActiveSub(prevSub)
 
     if (isHydrating) this.hydrate(render == null)
@@ -219,7 +218,9 @@ export class DynamicFragment extends VaporFragment {
     render: BlockFn | undefined,
     transition: VaporTransitionHooks | undefined,
     parent: ParentNode | null,
+    key: any,
   ): void {
+    this.current = key
     if (render) {
       const keepAliveCtx = this.keepAliveCtx
       // try to reuse the kept-alive scope
@@ -230,28 +231,33 @@ export class DynamicFragment extends VaporFragment {
         this.scope = new EffectScope()
       }
 
-      const needBranchKey = keepAliveCtx && this.keyed
-      const prevBranchKey = needBranchKey
-        ? keepAliveCtx.setCurrentBranchKey(this.current)
-        : undefined
-      try {
-        this.nodes = this.runWithRenderCtx(() => this.scope!.run(render) || [])
-      } finally {
-        // set key on blocks
-        if (this.keyed) setKey(this.nodes, this.current)
-
-        if (transition) {
-          this.$transition = applyTransitionHooks(this.nodes, transition)
-        }
+      const renderBranch = () => {
+        try {
+          this.nodes = this.runWithRenderCtx(
+            () => this.scope!.run(render) || [],
+          )
+        } finally {
+          // set key on blocks
+          if (this.keyed) setKey(this.nodes, this.current)
+
+          if (transition) {
+            this.$transition = applyTransitionHooks(this.nodes, transition)
+          }
 
-        // call processShapeFlag to mark shapeFlag before mounting.
-        // Must be called before restoring branchKey so getCacheKey
-        // generates the correct composite key.
-        if (keepAliveCtx) {
-          keepAliveCtx.processShapeFlag(this.nodes)
+          // call processShapeFlag to mark shapeFlag before mounting.
+          // This must run before leaving the keyed cache-key context so
+          // creating components inside the branch can still resolve the
+          // same cache key during initial mount.
+          if (keepAliveCtx) {
+            keepAliveCtx.processShapeFlag(this.nodes)
+          }
         }
+      }
 
-        if (needBranchKey) keepAliveCtx.setCurrentBranchKey(prevBranchKey)
+      if (keepAliveCtx && this.keyed) {
+        withCurrentCacheKey(key, renderBranch)
+      } else {
+        renderBranch()
       }
 
       if (parent) {