1 files changed, 14 insertions, 4 deletions

diff --git a/recipes/mattermost/webview.js b/recipes/mattermost/webview.js
index f348da4..c69c370 100644
--- a/recipes/mattermost/webview.js
+++ b/recipes/mattermost/webview.js
@@ -1,13 +1,23 @@
 "use strict";
 
 module.exports = Franz => {
+  const DIRECT_MESSAGES_INDIVIDUAL = '#sidebar-left .unread-title .DirectChannel__profile-picture';
+  const DIRECT_MESSAGES_GROUP = '#sidebar-left .unread-title .status--group';
+  const DIRECT_MESSAGES_LEGACY = '.sidebar--left .has-badge .badge';
+  const ALL_MESSAGES = '#sidebar-left .unread-title';
+  const ALL_MESSAGES_LEGACY = '#sidebar-left .unread-title';
+
   const getMessages = function getMessages() {
-    const directMessages = document.querySelectorAll('.sidebar--left .has-badge .badge').length;
-    const allMessages = document.querySelectorAll('.sidebar--left .has-badge').length - directMessages;
-    const channelMessages = document.querySelectorAll('.sidebar--left .unread-title').length - allMessages;
+    const directMessagesSelector = [DIRECT_MESSAGES_LEGACY, DIRECT_MESSAGES_INDIVIDUAL, DIRECT_MESSAGES_GROUP].join(', ');
+    const directMessages = document.querySelectorAll(directMessagesSelector).length;
+
+    const allMessagesSelector = [ALL_MESSAGES, ALL_MESSAGES_LEGACY].join(', ');
+    const allMessages = document.querySelectorAll(allMessagesSelector).length - directMessages;
+
     const teamDirectMessages = document.querySelectorAll('.team-wrapper .team-container .badge').length;
     const teamMessages = document.querySelectorAll('.team-wrapper .unread').length - teamDirectMessages;
-    Franz.setBadge(directMessages + teamDirectMessages, allMessages + channelMessages + teamMessages);
+
+    Franz.setBadge(directMessages + teamDirectMessages, allMessages + teamMessages);
   };
 
   Franz.loop(getMessages);