From 7b0e4d0edcc86de9523019aa17968af7512186e2 Mon Sep 17 00:00:00 2001
From: Aman Harwara <amanharwara@protonmail.com>
Date: Thu, 14 Sep 2023 15:25:24 +0530
Subject: [PATCH] chore: show current vault in template note if using exclusive
 mode

---
 .../src/Vaults/VaultDisplayService.ts         | 14 +++++
 .../Vaults/VaultDisplayServiceInterface.ts    |  4 +-
 .../src/javascripts/Hooks/useItemVaultInfo.ts | 53 +++++++++++++------
 3 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/packages/ui-services/src/Vaults/VaultDisplayService.ts b/packages/ui-services/src/Vaults/VaultDisplayService.ts
index 40208bae8..0fc25a48a 100644
--- a/packages/ui-services/src/Vaults/VaultDisplayService.ts
+++ b/packages/ui-services/src/Vaults/VaultDisplayService.ts
@@ -14,6 +14,7 @@ import {
   VaultLockServiceEvent,
 } from '@standardnotes/services'
 import {
+  DecryptedItemInterface,
   ExclusionaryOptions,
   VaultDisplayOptions,
   VaultDisplayOptionsPersistable,
@@ -103,6 +104,19 @@ export class VaultDisplayService
     return this.options.isInExclusiveDisplayMode()
   }
 
+  getItemVault(item: DecryptedItemInterface): VaultListingInterface | undefined {
+    if (this.application.items.isTemplateItem(item)) {
+      if (this.exclusivelyShownVault) {
+        return this.exclusivelyShownVault
+      }
+
+      return undefined
+    }
+
+    const vault = this.application.vaults.getItemVault(item)
+    return vault
+  }
+
   changeToMultipleVaultDisplayMode(): void {
     const vaults = this.application.vaults.getVaults()
     const lockedVaults = this.application.vaultLocks.getLockedvaults()
diff --git a/packages/ui-services/src/Vaults/VaultDisplayServiceInterface.ts b/packages/ui-services/src/Vaults/VaultDisplayServiceInterface.ts
index f7056569f..b69fbd498 100644
--- a/packages/ui-services/src/Vaults/VaultDisplayServiceInterface.ts
+++ b/packages/ui-services/src/Vaults/VaultDisplayServiceInterface.ts
@@ -1,9 +1,11 @@
-import { VaultDisplayOptions, VaultListingInterface } from '@standardnotes/models'
+import { DecryptedItemInterface, VaultDisplayOptions, VaultListingInterface } from '@standardnotes/models'
 import { AbstractUIServiceInterface } from '../Abstract/AbstractUIServiceInterface'
 
 export interface VaultDisplayServiceInterface extends AbstractUIServiceInterface {
   exclusivelyShownVault?: VaultListingInterface
 
+  getItemVault(item: DecryptedItemInterface): VaultListingInterface | undefined
+
   getOptions(): VaultDisplayOptions
 
   isVaultDisabledOrLocked(vault: VaultListingInterface): boolean
diff --git a/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts b/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts
index 833bd9a96..b002d02a5 100644
--- a/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts
+++ b/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts
@@ -1,7 +1,12 @@
 import { useApplication } from '@/Components/ApplicationProvider'
-import { DecryptedItemInterface, TrustedContactInterface, VaultListingInterface } from '@standardnotes/snjs'
-import useItem from './useItem'
-import { useRef } from 'react'
+import {
+  ContentType,
+  DecryptedItemInterface,
+  TrustedContactInterface,
+  VaultListingInterface,
+} from '@standardnotes/snjs'
+import { useCallback, useEffect, useLayoutEffect, useState } from 'react'
+import { useStateRef } from './useStateRef'
 
 type ItemVaultInfo = {
   vault?: VaultListingInterface
@@ -12,21 +17,37 @@ type ItemVaultInfo = {
 export const useItemVaultInfo = (item: DecryptedItemInterface): ItemVaultInfo => {
   const application = useApplication()
 
-  const info = useRef<ItemVaultInfo>({
-    vault: undefined,
-    lastEditedByContact: undefined,
-    sharedByContact: undefined,
-  })
+  const [vault, setVault] = useState<VaultListingInterface>()
+  const vaultRef = useStateRef(vault)
+  const [lastEditedByContact, setLastEditedByContact] = useState<TrustedContactInterface>()
+  const [sharedByContact, setSharedByContact] = useState<TrustedContactInterface>()
 
-  info.current.vault = useItem(application.vaults.getItemVault(item)?.uuid)
+  const updateInfo = useCallback(() => {
+    if (!application.featuresController.isEntitledToVaults()) {
+      return
+    }
 
-  const lastEditedBy = application.sharedVaults.getItemLastEditedBy(item)
-  info.current.lastEditedByContact = lastEditedBy || info.current.lastEditedByContact
-  info.current.sharedByContact = application.sharedVaults.getItemSharedBy(item)
+    setVault(application.vaultDisplayService.getItemVault(item))
+    setLastEditedByContact((lastEditedBy) => application.sharedVaults.getItemLastEditedBy(item) || lastEditedBy)
+    setSharedByContact(application.sharedVaults.getItemSharedBy(item))
+  }, [application.featuresController, application.sharedVaults, application.vaultDisplayService, item])
 
-  if (!application.featuresController.isEntitledToVaults()) {
-    return info.current
+  useLayoutEffect(() => {
+    updateInfo()
+  }, [updateInfo])
+
+  useEffect(() => {
+    return application.items.streamItems(ContentType.TYPES.VaultListing, ({ changed, inserted }) => {
+      const matchingItem = changed.concat(inserted).find((vault) => vault.uuid === vaultRef.current?.uuid)
+      if (matchingItem) {
+        setVault(matchingItem as VaultListingInterface)
+      }
+    })
+  }, [application.items, vaultRef])
+
+  return {
+    vault,
+    lastEditedByContact,
+    sharedByContact,
   }
-
-  return info.current
 }