diff --git a/lib/features/mailbox/presentation/extensions/list_mailbox_node_extension.dart b/lib/features/mailbox/presentation/extensions/list_mailbox_node_extension.dart
deleted file mode 100644
index d53ed8e8a..000000000
--- a/lib/features/mailbox/presentation/extensions/list_mailbox_node_extension.dart
+++ /dev/null
@@ -1,45 +0,0 @@
-
-import 'package:jmap_dart_client/jmap/mail/mailbox/mailbox.dart';
-import 'package:model/mailbox/expand_mode.dart';
-import 'package:model/mailbox/select_mode.dart';
-import 'package:tmail_ui_user/features/mailbox/presentation/model/mailbox_node.dart';
-
-extension ListMailboxNodeExtensions on List<MailboxNode> {
-  List<MailboxNode> updateNode(MailboxId currentMailboxId, MailboxNode newNode) {
-    return map((MailboxNode child) {
-      if (child.item.id == currentMailboxId) {
-        return newNode;
-      } else {
-        if (child.hasChildren()) {
-          return child.copyWith(mailboxNodes: child.updateNode(currentMailboxId, newNode));
-        }
-        return child;
-      }
-    }).toList();
-  }
-
-  List<MailboxNode> toggleSelectMailboxNode(MailboxNode selectedMailboxMode) {
-    return map((MailboxNode child) {
-      if (child.item.id == selectedMailboxMode.item.id) {
-        return child.toggleSelectMailboxNode();
-      } else {
-        if (child.hasChildren()) {
-          return child.copyWith(mailboxNodes: child.toggleSelectNode(selectedMailboxMode));
-        }
-        return child;
-      }
-    }).toList();
-  }
-
-  List<MailboxNode> toSelectMailboxNode({required SelectMode selectMode, ExpandMode? newExpandMode}) {
-    return map((MailboxNode child) {
-      if (child.hasChildren()) {
-        return child.copyWith(
-            mailboxNodes: child.toSelectedNode(selectMode: selectMode, newExpandMode: newExpandMode),
-            newSelectMode: selectMode,
-            newExpandMode: newExpandMode);
-      }
-      return child.toSelectedMailboxNode(selectMode: selectMode, newExpandMode: newExpandMode);
-    }).toList();
-  }
-}
\ No newline at end of file
diff --git a/lib/features/mailbox/presentation/model/mailbox_node.dart b/lib/features/mailbox/presentation/model/mailbox_node.dart
index 2cc30e541..8bbb8c882 100644
--- a/lib/features/mailbox/presentation/model/mailbox_node.dart
+++ b/lib/features/mailbox/presentation/model/mailbox_node.dart
@@ -4,7 +4,6 @@ import 'package:jmap_dart_client/jmap/mail/mailbox/mailbox.dart';
 import 'package:model/mailbox/expand_mode.dart';
 import 'package:model/mailbox/presentation_mailbox.dart';
 import 'package:model/mailbox/select_mode.dart';
-import 'package:tmail_ui_user/features/mailbox/presentation/model/mailbox_tree.dart';
 
 class MailboxNode with EquatableMixin{
   static PresentationMailbox _root = PresentationMailbox(MailboxId(Id('root')));
@@ -36,31 +35,6 @@ class MailboxNode with EquatableMixin{
     childrenItems?.add(node);
   }
 
-  String getPathMailboxNode(MailboxTree mailboxTree, List<PresentationMailbox> defaultMailboxList) {
-    String path = '${item.name?.name}';
-
-    var parentId = item.parentId;
-
-    while(parentId != null) {
-      var parent = mailboxTree.findNode(parentId);
-      if (parent == null) {
-        try {
-          final parentItem = defaultMailboxList.firstWhere((mailbox) => mailbox.id == parentId);
-          parent = MailboxNode(parentItem);
-        } catch(e) {}
-      }
-
-      if (parent != null) {
-        path = '${parent.item.name?.name}/$path';
-        parentId = parent.item.parentId;
-      } else {
-        break;
-      }
-    }
-
-    return path;
-  }
-
   List<MailboxNode>? updateNode(MailboxId mailboxId, MailboxNode newNode, {MailboxNode? parent}) {
     List<MailboxNode>? _children = parent == null ? this.childrenItems : parent.childrenItems;
     return _children?.map((MailboxNode child) {
@@ -69,7 +43,7 @@ class MailboxNode with EquatableMixin{
       } else {
         if (child.hasChildren()) {
           return child.copyWith(
-            mailboxNodes: updateNode(
+            children: updateNode(
               mailboxId,
               newNode,
               parent: child,
@@ -81,6 +55,24 @@ class MailboxNode with EquatableMixin{
     }).toList();
   }
 
+  List<MailboxNode> descendantsAsList(){
+    List<MailboxNode> listOfNodes = <MailboxNode>[];
+    _appendDescendants(this, listOfNodes);
+    return listOfNodes;
+  }
+
+  void _appendDescendants(MailboxNode? node, List<MailboxNode> listOfNodes) {
+    if (node != null) {
+      listOfNodes.add(node);
+      List<MailboxNode>? childrenItems = node.childrenItems;
+      if (childrenItems != null) {
+        childrenItems.forEach((child) {
+          _appendDescendants(child, listOfNodes);
+        });
+      }
+    }
+  }
+
   List<MailboxNode>? toggleSelectNode(MailboxNode selectedMailboxMode, {MailboxNode? parent}) {
     List<MailboxNode>? _children = parent == null ? this.childrenItems : parent.childrenItems;
     return _children?.map((MailboxNode child) {
@@ -88,7 +80,7 @@ class MailboxNode with EquatableMixin{
         return child.toggleSelectMailboxNode();
       } else {
         if (child.hasChildren()) {
-          return child.copyWith(mailboxNodes: toggleSelectNode(selectedMailboxMode, parent: child));
+          return child.copyWith(children: toggleSelectNode(selectedMailboxMode, parent: child));
         }
         return child;
       }
@@ -100,7 +92,7 @@ class MailboxNode with EquatableMixin{
     return _children?.map((MailboxNode child) {
       if (child.hasChildren()) {
         return child.copyWith(
-            mailboxNodes: toSelectedNode(selectMode: selectMode, newExpandMode: newExpandMode, parent: child),
+            children: toSelectedNode(selectMode: selectMode, newExpandMode: newExpandMode, parent: child),
             newSelectMode: selectMode,
             newExpandMode: newExpandMode);
       }
@@ -109,18 +101,18 @@ class MailboxNode with EquatableMixin{
   }
 
   @override
-  List<Object?> get props => [item, childrenItems];
+  List<Object?> get props => [item, childrenItems, expandMode, selectMode];
 }
 
 extension MailboxNodeExtension on MailboxNode {
   MailboxNode copyWith({
     SelectMode? newSelectMode,
     ExpandMode? newExpandMode,
-    List<MailboxNode>? mailboxNodes
+    List<MailboxNode>? children
   }) {
     return MailboxNode(
       item,
-      childrenItems: mailboxNodes ?? childrenItems,
+      childrenItems: children ?? childrenItems,
       expandMode: newExpandMode ?? expandMode,
       selectMode: newSelectMode ?? selectMode,
     );
diff --git a/lib/features/mailbox/presentation/model/mailbox_tree.dart b/lib/features/mailbox/presentation/model/mailbox_tree.dart
index 116cc719d..c61c7577f 100644
--- a/lib/features/mailbox/presentation/model/mailbox_tree.dart
+++ b/lib/features/mailbox/presentation/model/mailbox_tree.dart
@@ -3,20 +3,24 @@ import 'dart:collection';
 
 import 'package:equatable/equatable.dart';
 import 'package:jmap_dart_client/jmap/mail/mailbox/mailbox.dart';
+import 'package:model/mailbox/expand_mode.dart';
+import 'package:model/mailbox/select_mode.dart';
 
 import 'mailbox_node.dart';
 
+typedef NodeQuery = bool Function(MailboxNode node);
+
 class MailboxTree with EquatableMixin {
   MailboxNode root;
   MailboxTree(this.root);
 
-  MailboxNode? findNode(MailboxId? mailboxId) {
+  MailboxNode? findNode(NodeQuery nodeQuery) {
     var result;
     final queue = ListQueue<MailboxNode>();
     queue.addLast(root);
-    while (queue.isNotEmpty && mailboxId != null) {
+    while (queue.isNotEmpty) {
       final currentNode = queue.removeFirst();
-      if (mailboxId == currentNode.item.id) {
+      if (nodeQuery(currentNode)) {
         result = currentNode;
         break;
       }
@@ -27,22 +31,72 @@ class MailboxTree with EquatableMixin {
     return result;
   }
 
-  List<MailboxNode> getAllNodes(MailboxNode node){
-    List<MailboxNode> listOfNodes = <MailboxNode>[];
-    _addAllNodes(node, listOfNodes);
-    return listOfNodes;
+  List<MailboxNode> findNodes(NodeQuery nodeQuery) {
+    final listResult = List<MailboxNode>.empty(growable: true);
+    final queue = ListQueue<MailboxNode>();
+    queue.addLast(root);
+    while (queue.isNotEmpty) {
+      final currentNode = queue.removeFirst();
+      if (nodeQuery(currentNode)) {
+        listResult.add(currentNode);
+      }
+      currentNode.childrenItems?.forEach((child) {
+        queue.addLast(child);
+      });
+    }
+    return listResult;
   }
 
-  void _addAllNodes(MailboxNode? node, List<MailboxNode> listOfNodes) {
-    if (node != null) {
-      listOfNodes.add(node);
-      List<MailboxNode>? childrenItems = node.childrenItems;
-      if (childrenItems != null) {
-        childrenItems.forEach((child) {
-          _addAllNodes(child, listOfNodes);
-        });
-      }
+  MailboxNode? updateExpandedNode(MailboxNode selectedNode, ExpandMode newExpandMode) {
+    var matchedNode = findNode((node) => node.item.id == selectedNode.item.id);
+    matchedNode?.expandMode = newExpandMode;
+    return matchedNode;
+  }
+
+  MailboxNode? updateSelectedNode(MailboxNode selectedNode, SelectMode newSelectMode) {
+    var matchedNode = findNode((node) => node.item.id == selectedNode.item.id);
+    matchedNode?.selectMode = newSelectMode;
+    return matchedNode;
+  }
+
+  void updateNodesUIMode(SelectMode? selectMode, ExpandMode? expandMode) {
+    if (selectMode == null && expandMode == null) {
+      return;
     }
+    final queue = ListQueue<MailboxNode>();
+    queue.addLast(root);
+    while (queue.isNotEmpty) {
+      final currentNode = queue.removeFirst();
+      if (expandMode != null) {
+        currentNode.expandMode = expandMode;
+      }
+      if (selectMode != null) {
+        currentNode.selectMode = selectMode;
+      }
+      currentNode.childrenItems?.forEach((child) {
+        queue.addLast(child);
+      });
+    }
+  }
+
+  String? getNodePath(MailboxId mailboxId) {
+    final matchedNode = findNode((node) => node.item.id == mailboxId);
+    if (matchedNode == null) {
+      return null;
+    }
+    String path = '${matchedNode.item.name?.name}';
+
+    var parentId = matchedNode.item.parentId;
+
+    while(parentId != null) {
+      var parentNode = findNode((node) => node.item.id == parentId);
+      if (parentNode == null) {
+        break;
+      }
+      path = '${parentNode.item.name?.name}/$path';
+      parentId = parentNode.item.parentId;
+    }
+    return path;
   }
 
   @override
diff --git a/lib/features/mailbox/presentation/model/mailbox_tree_builder.dart b/lib/features/mailbox/presentation/model/mailbox_tree_builder.dart
index e483b0947..09e2601f0 100644
--- a/lib/features/mailbox/presentation/model/mailbox_tree_builder.dart
+++ b/lib/features/mailbox/presentation/model/mailbox_tree_builder.dart
@@ -1,6 +1,7 @@
 import 'dart:collection';
-import 'package:collection/collection.dart';
 
+import 'package:collection/collection.dart';
+import 'package:dartz/dartz.dart';
 import 'package:jmap_dart_client/jmap/mail/mailbox/mailbox.dart';
 import 'package:model/model.dart';
 
@@ -37,4 +38,40 @@ class TreeBuilder {
     });
     return tree;
   }
+
+  Future<Tuple2<MailboxTree, MailboxTree>> generateMailboxTreeInUI(List<PresentationMailbox> allMailboxes) async {
+    final Map<MailboxId, MailboxNode> mailboxDictionary = HashMap();
+
+    final defaultTree = MailboxTree(MailboxNode.root());
+    final folderTree = MailboxTree(MailboxNode.root());
+
+    allMailboxes.forEach((mailbox) {
+      mailboxDictionary[mailbox.id] = MailboxNode(mailbox);
+    });
+
+    allMailboxes.forEach((mailbox) {
+      final parentId = mailbox.parentId;
+      final parentNode = mailboxDictionary[parentId];
+      final node = mailboxDictionary[mailbox.id];
+      if (node != null) {
+        if (parentNode != null) {
+          parentNode.addChildNode(node);
+          parentNode.childrenItems?.sortByCompare<MailboxName?>(
+            (node) => node.item.name,
+            (name, other) => name?.compareAlphabetically(other) ?? -1
+          );
+        } else {
+          var tree = mailbox.hasRole() ? defaultTree : folderTree;
+
+          tree.root.addChildNode(node);
+          tree.root.childrenItems?.sortByCompare<MailboxName?>(
+            (node) => node.item.name,
+            (name, other) => name?.compareAlphabetically(other) ?? -1
+          );
+        }
+      }
+    });
+
+    return Tuple2(defaultTree, folderTree);
+  }
 }
\ No newline at end of file