1 files changed, 2 insertions, 6 deletions

diff --git a/recipes/instagram-direct-messages/webview.js b/recipes/instagram-direct-messages/webview.js
index b59b9f2..290f596 100644
--- a/recipes/instagram-direct-messages/webview.js
+++ b/recipes/instagram-direct-messages/webview.js
@@ -14,12 +14,8 @@ setInterval(() => {
 
 module.exports = (Ferdium) => {
   const getMessages = () => {
-    let count = 0;
-    const elements = document.querySelectorAll('.rBNOH.soMvl');
-    if (elements.length > 0) {
-      count += 1;
-    }
-    Ferdium.setBadge(count);
+    const element = document.querySelector('a[href^="/direct/inbox"]');
+    Ferdium.setBadge(element ? Ferdium.safeParseInt(element.textContent.match(/\d+/)[0]) : 0);
   };
 
   Ferdium.loop(getMessages);