fix(runtime-vapor): update setCurrentBranchKey to return previous key and handle context correctly

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

@@ -43,7 +43,7 @@ export interface KeepAliveContext {
   cacheBlock(): void
   cacheScope(key: any, scope: EffectScope): void
   getScope(key: any): EffectScope | undefined
-  setCurrentBranchKey(key: any): void
+  setCurrentBranchKey(key: any): any
 }
 
 export let currentKeepAliveCtx: KeepAliveContext | null = null
@@ -314,7 +314,11 @@ const KeepAliveImpl: ObjectVaporComponent = defineVaporComponent({
         }
       },
       setCurrentBranchKey(key) {
-        currentBranchKey = key
+        try {
+          return currentBranchKey
+        } finally {
+          currentBranchKey = key
+        }
       },
     }
 
@@ -386,10 +390,7 @@ function resolveKey(
   branchKey?: any,
 ): CacheKey {
   if (key != null) {
-    if (branchKey !== undefined && key === branchKey) {
-      return getCompositeKey(type, branchKey)
-    }
-    return key as CacheKey
+    return getCompositeKey(type, key)
   }
   if (branchKey !== undefined) {
     return getCompositeKey(type, branchKey)

packages/runtime-vapor/src/fragment.ts

@@ -199,14 +199,18 @@ export class DynamicFragment extends VaporFragment {
       const prevOwner = setCurrentSlotOwner(this.slotOwner)
       // set currentKeepAliveCtx so nested DynamicFragments and components can capture it
       const prevCtx = setCurrentKeepAliveCtx(keepAliveCtx)
+      let prevBranchKey: any
       if (keepAliveCtx && this.keyed) {
-        keepAliveCtx.setCurrentBranchKey(this.current)
+        prevBranchKey = keepAliveCtx.setCurrentBranchKey(this.current)
       }
       // switch current instance to parent instance during update
       // ensure that the parent instance is correct for nested components
       const prev = parent && instance ? setCurrentInstance(instance) : undefined
       this.nodes = this.scope.run(render) || []
       if (prev !== undefined) setCurrentInstance(...prev)
+      if (keepAliveCtx && this.keyed) {
+        keepAliveCtx.setCurrentBranchKey(prevBranchKey)
+      }
       setCurrentKeepAliveCtx(prevCtx)
       setCurrentSlotOwner(prevOwner)