2 files changed, 8 insertions, 19 deletions

diff --git a/recipes/element/package.json b/recipes/element/package.json
index 58ca48d..beba777 100644
--- a/recipes/element/package.json
+++ b/recipes/element/package.json
@@ -1,7 +1,7 @@
 {
   "id": "element",
   "name": "Element",
-  "version": "1.3.0",
+  "version": "1.3.1",
   "license": "MIT",
   "aliases": [
     "Riot.im",
diff --git a/recipes/element/webview.js b/recipes/element/webview.js
index ebe557f..3528418 100644
--- a/recipes/element/webview.js
+++ b/recipes/element/webview.js
@@ -1,25 +1,14 @@
 module.exports = Ferdium => {
   function getMessages() {
-    const badges = document.querySelectorAll('.mx_RoomList .mx_RoomSublist_tiles .mx_NotificationBadge');
-    // Number of messages from rooms which has "All Messages" notifications enabled.
-    // Always incremented for private rooms by default, incremented for group chats if
-    // "All Messages" is selected for them in notifications settings.
     let directCount = 0;
-    // Number of messages for rooms which has "Only Highlights" notifications level set.
-    // Appears in rooms list with dots on right.
-    let indirectCount = 0;
-    
-    for (const badge of badges) {
-      if (badge.classList.contains('mx_NotificationBadge_dot')) {
-        indirectCount++;
-      } else {
-        directCount += Ferdium.safeParseInt(badge.childNodes[0].textContent);
-      }
-    }
-
-    // set Ferdium badge
+    const spacesBar = document.querySelector('.mx_SpaceTreeLevel');
+    spacesBar.querySelectorAll('.mx_NotificationBadge_count').forEach((badge) => {
+      directCount += Ferdium.safeParseInt(badge.textContent);
+    });
+    const indirectCount = spacesBar.querySelectorAll('.mx_NotificationBadge_dot')
+      .length;
     Ferdium.setBadge(directCount, indirectCount);
   }
+
   Ferdium.loop(getMessages);
 };
-