Add additional logging and hack to demonstrate the problem

libs/common/src/platform/state/implementations/state-base.ts

@@ -1,10 +1,10 @@
 import {
-  Observable,
-  ReplaySubject,
   defer,
   filter,
   firstValueFrom,
   merge,
+  Observable,
+  ReplaySubject,
   share,
   switchMap,
   tap,
@@ -20,7 +20,7 @@ import {
   ObservableStorageService,
 } from "../../abstractions/storage.service";
 import { DebugOptions } from "../key-definition";
-import { StateUpdateOptions, populateOptionsWithDefault } from "../state-update-options";
+import { populateOptionsWithDefault, StateUpdateOptions } from "../state-update-options";
 
 import { getStoredValue } from "./util";
 
@@ -54,8 +54,20 @@ export abstract class StateBase<T, KeyDef extends KeyDefinitionRequirements<T>>
     );
 
     let state$ = merge(
-      defer(() => getStoredValue(key, storageService, keyDefinition.deserializer)),
-      storageUpdate$,
+      defer(() => getStoredValue(key, storageService, keyDefinition.deserializer)).pipe(
+        tap(() => {
+          if (keyDefinition.debug.enableRetrievalLogging) {
+            this.logService.info(`Emission from defer for '${key}'`);
+          }
+        }),
+      ),
+      storageUpdate$.pipe(
+        tap(() => {
+          if (keyDefinition.debug.enableRetrievalLogging) {
+            this.logService.info(`Emission from update for '${key}'`);
+          }
+        }),
+      ),
     );
 
     if (keyDefinition.debug.enableRetrievalLogging) {
@@ -75,7 +87,21 @@ export abstract class StateBase<T, KeyDef extends KeyDefinitionRequirements<T>>
       state$ = state$.pipe(
         share({
           connector: () => new ReplaySubject(1),
-          resetOnRefCountZero: () => timer(keyDefinition.cleanupDelayMs),
+          resetOnRefCountZero: () =>
+            timer(keyDefinition.cleanupDelayMs).pipe(
+              tap(() => {
+                if (keyDefinition.debug.enableRetrievalLogging) {
+                  this.logService.info(
+                    `Cleaning up '${key}' from memory cache after ${keyDefinition.cleanupDelayMs}ms`,
+                  );
+                }
+              }),
+            ),
+        }),
+        tap(() => {
+          if (keyDefinition.debug.enableRetrievalLogging) {
+            this.logService.info(`Retrieving '${key}' from state replay subject`);
+          }
         }),
       );
     }

libs/common/src/vault/services/folder/folder.service.ts

@@ -1,4 +1,4 @@
-import { Observable, firstValueFrom, map, shareReplay, switchMap } from "rxjs";
+import { firstValueFrom, map, Observable, shareReplay, switchMap, tap, timer } from "rxjs";
 
 import { EncryptService } from "@bitwarden/common/platform/abstractions/encrypt.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
@@ -18,6 +18,8 @@ import { Cipher } from "../../models/domain/cipher";
 import { FolderWithIdRequest } from "../../models/request/folder-with-id.request";
 import { FOLDER_DECRYPTED_FOLDERS, FOLDER_ENCRYPTED_FOLDERS } from "../key-state/folder.state";
 
+let folderViewCount = 0;
+
 export class FolderService implements InternalFolderServiceAbstraction {
   constructor(
     private keyService: KeyService,
@@ -40,9 +42,24 @@ export class FolderService implements InternalFolderServiceAbstraction {
   }
 
   folderViews$(userId: UserId): Observable<FolderView[]> {
+    const folderViewId = folderViewCount++;
+    console.log(`FolderView$ ${folderViewId} created`);
     return this.encryptedFoldersState(userId).state$.pipe(
       switchMap((folderData) => {
-        return this.decryptFolders(userId, folderData);
+        // RACE CONDITION HACK, DO NOT MERGE:
+        // Add artificial delay to let encryptedFoldersState finish updating, the switchMap will cancel the previous observable
+        // encryptedFoldersState has a shareReplay subject internally, so the old state is still available immediately after update.
+        // By adding the timer(0), we are effectively skipping the "stale" emission.
+
+        // A slightly better fix would be to cache these `folderView$` observables and return the same one for the same user.
+        // That way we can put a shareReplay on this observable which prevents future subscribers from starting a
+        // new subscription to encryptedFoldersState and getting stale value.
+        return timer(0).pipe(
+          tap(() => {
+            console.log(`FolderView$ ${folderViewId} emitting`);
+          }),
+          switchMap(() => this.decryptFolders(userId, folderData)),
+        );
       }),
     );
   }
@@ -94,7 +111,9 @@ export class FolderService implements InternalFolderServiceAbstraction {
 
   async upsert(folderData: FolderData | FolderData[], userId: UserId): Promise<void> {
     await this.clearDecryptedFolderState(userId);
+    console.log("starting encryptedFoldersState update");
     await this.encryptedFoldersState(userId).update((folders) => {
+      console.log("encryptedFoldersState update callback");
       if (folders == null) {
         folders = {};
       }
@@ -110,6 +129,7 @@ export class FolderService implements InternalFolderServiceAbstraction {
 
       return folders;
     });
+    console.log("ending encryptedFoldersState update");
   }
 
   async replace(folders: { [id: string]: FolderData }, userId: UserId): Promise<void> {

libs/common/src/vault/services/key-state/folder.state.ts

@@ -10,14 +10,22 @@ export const FOLDER_ENCRYPTED_FOLDERS = UserKeyDefinition.record<FolderData>(
   {
     deserializer: (obj: Jsonify<FolderData>) => FolderData.fromJSON(obj),
     clearOn: ["logout"],
+    debug: {
+      enableUpdateLogging: true,
+      enableRetrievalLogging: true,
+    },
   },
 );
 
 export const FOLDER_DECRYPTED_FOLDERS = new UserKeyDefinition<FolderView[]>(
   FOLDER_MEMORY,
-  "folder",
+  "folderViews",
   {
     deserializer: (obj: Jsonify<FolderView[]>) => obj?.map((f) => FolderView.fromJSON(f)) ?? [],
     clearOn: ["logout", "lock"],
+    debug: {
+      enableUpdateLogging: true,
+      enableRetrievalLogging: true,
+    },
   },
 );