From 53785c7100ffd0e142f5c3fba0d78c4d064f1d93 Mon Sep 17 00:00:00 2001
From: Aman Harwara <amanharwara@protonmail.com>
Date: Tue, 12 Sep 2023 00:02:01 +0530
Subject: [PATCH] chore: correctly refresh vault info badge

---
 .../services/src/Domain/Vault/VaultService.ts     |  5 +++++
 .../ContentListView/ListItemVaultInfo.tsx         | 15 ++-------------
 .../web/src/javascripts/Hooks/useItemVaultInfo.ts | 10 +++++-----
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/packages/services/src/Domain/Vault/VaultService.ts b/packages/services/src/Domain/Vault/VaultService.ts
index 0e580bddc..d807aa2bd 100644
--- a/packages/services/src/Domain/Vault/VaultService.ts
+++ b/packages/services/src/Domain/Vault/VaultService.ts
@@ -243,7 +243,12 @@ export class VaultService
   }
 
   getItemVault(item: DecryptedItemInterface): VaultListingInterface | undefined {
+    if (this.items.isTemplateItem(item)) {
+      return undefined
+    }
+
     const latestItem = this.items.findItem(item.uuid)
+
     if (!latestItem) {
       throw new Error('Cannot find latest version of item to get vault for')
     }
diff --git a/packages/web/src/javascripts/Components/ContentListView/ListItemVaultInfo.tsx b/packages/web/src/javascripts/Components/ContentListView/ListItemVaultInfo.tsx
index f93a38378..948f7ac7e 100644
--- a/packages/web/src/javascripts/Components/ContentListView/ListItemVaultInfo.tsx
+++ b/packages/web/src/javascripts/Components/ContentListView/ListItemVaultInfo.tsx
@@ -1,8 +1,8 @@
 import { FunctionComponent } from 'react'
-import { useApplication } from '../ApplicationProvider'
 import { DecryptedItemInterface, classNames } from '@standardnotes/snjs'
 import VaultNameBadge from '../Vaults/VaultNameBadge'
 import SharedByContactBadge from '../Vaults/SharedByContactBadge'
+import { useItemVaultInfo } from '@/Hooks/useItemVaultInfo'
 
 type Props = {
   item: DecryptedItemInterface
@@ -10,23 +10,12 @@ type Props = {
 }
 
 const ListItemVaultInfo: FunctionComponent<Props> = ({ item, className }) => {
-  const application = useApplication()
+  const { vault, sharedByContact } = useItemVaultInfo(item)
 
-  if (!application.featuresController.isEntitledToVaults()) {
-    return null
-  }
-
-  if (application.items.isTemplateItem(item)) {
-    return null
-  }
-
-  const vault = application.vaults.getItemVault(item)
   if (!vault) {
     return null
   }
 
-  const sharedByContact = application.sharedVaults.getItemSharedBy(item)
-
   return (
     <div className={classNames('flex flex-wrap items-center gap-2', className)}>
       <VaultNameBadge vault={vault} />
diff --git a/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts b/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts
index de5ce0498..8ca99b401 100644
--- a/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts
+++ b/packages/web/src/javascripts/Hooks/useItemVaultInfo.ts
@@ -1,9 +1,11 @@
 import { useApplication } from '@/Components/ApplicationProvider'
 import { DecryptedItemInterface, TrustedContactInterface, VaultListingInterface } from '@standardnotes/snjs'
+import useItem from './useItem'
 
 type ItemVaultInfo = {
   vault?: VaultListingInterface
   lastEditedByContact?: TrustedContactInterface
+  sharedByContact?: TrustedContactInterface
 }
 
 export const useItemVaultInfo = (item: DecryptedItemInterface): ItemVaultInfo => {
@@ -14,16 +16,14 @@ export const useItemVaultInfo = (item: DecryptedItemInterface): ItemVaultInfo =>
     lastEditedByContact: undefined,
   }
 
+  info.vault = useItem(application.vaults.getItemVault(item)?.uuid)
+
   if (!application.featuresController.isEntitledToVaults()) {
     return info
   }
 
-  if (application.items.isTemplateItem(item)) {
-    return info
-  }
-
-  info.vault = application.vaults.getItemVault(item)
   info.lastEditedByContact = application.sharedVaults.getItemLastEditedBy(item)
+  info.sharedByContact = application.sharedVaults.getItemSharedBy(item)
 
   return info
 }