diff --git a/desktop/plugins/public/ui-debugger/components/Tree.tsx b/desktop/plugins/public/ui-debugger/components/Tree.tsx
index 0aed68602..009ff10dc 100644
--- a/desktop/plugins/public/ui-debugger/components/Tree.tsx
+++ b/desktop/plugins/public/ui-debugger/components/Tree.tsx
@@ -41,7 +41,7 @@ export function Tree(props: {
   onSelectNode: (id: Id) => void;
 }) {
   const instance = usePlugin(plugin);
-  const expandedItems = useValue(instance.uiState.treeState).expandedNodes;
+  const expandedItems = useValue(instance.uiState.expandedNodes);
   const focused = useValue(instance.uiState.focusedNode);
 
   const items = useMemo(
@@ -87,7 +87,7 @@ export function Tree(props: {
           viewState={{
             tree: {
               focusedItem: head(hoveredNodes),
-              expandedItems,
+              expandedItems: [...expandedItems],
               selectedItems: props.selectedNode ? [props.selectedNode] : [],
             },
           }}
@@ -95,15 +95,13 @@ export function Tree(props: {
             instance.uiState.hoveredNodes.set([item.index]);
           }}
           onExpandItem={(item) => {
-            instance.uiState.treeState.update((draft) => {
-              draft.expandedNodes.push(item.index);
+            instance.uiState.expandedNodes.update((draft) => {
+              draft.add(item.index);
             });
           }}
           onCollapseItem={(item) =>
-            instance.uiState.treeState.update((draft) => {
-              draft.expandedNodes = draft.expandedNodes.filter(
-                (expandedItemIndex) => expandedItemIndex !== item.index,
-              );
+            instance.uiState.expandedNodes.update((draft) => {
+              draft.delete(item.index);
             })
           }
           renderItem={renderItem}
diff --git a/desktop/plugins/public/ui-debugger/components/Tree2.tsx b/desktop/plugins/public/ui-debugger/components/Tree2.tsx
new file mode 100644
index 000000000..3f97873bc
--- /dev/null
+++ b/desktop/plugins/public/ui-debugger/components/Tree2.tsx
@@ -0,0 +1,152 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ */
+import {Id, UINode} from '../types';
+import React from 'react';
+import {
+  HighlightManager,
+  HighlightProvider,
+  styled,
+  theme,
+  useHighlighter,
+  usePlugin,
+  useValue,
+} from 'flipper-plugin';
+import {plugin} from '../index';
+
+export function Tree2({
+  nodes,
+  rootId,
+  selectedNode,
+  onSelectNode,
+}: {
+  nodes: Map<Id, UINode>;
+  rootId: Id;
+  selectedNode?: Id;
+  onSelectNode: (node?: Id) => void;
+}) {
+  const instance = usePlugin(plugin);
+  const expandedNodes = useValue(instance.uiState.expandedNodes);
+  const searchTerm = useValue(instance.uiState.searchTerm);
+
+  const items = toTreeList(nodes, rootId, expandedNodes);
+
+  return (
+    <HighlightProvider
+      text={searchTerm}
+      highlightColor={theme.searchHighlightBackground.yellow}>
+      <div>
+        {items.map((treeNode) => (
+          <TreeItemContainer
+            key={treeNode.id}
+            treeNode={treeNode}
+            selectedNode={selectedNode}
+            onSelectNode={onSelectNode}
+          />
+        ))}
+      </div>
+    </HighlightProvider>
+  );
+}
+
+export type TreeNode = UINode & {
+  depth: number;
+};
+
+function TreeItemContainer({
+  treeNode,
+  selectedNode,
+  hoveredNode,
+  onSelectNode,
+}: {
+  treeNode: TreeNode;
+  selectedNode?: Id;
+  hoveredNode?: Id;
+  onSelectNode: (node?: Id) => void;
+}) {
+  return (
+    <TreeItem
+      isSelected={treeNode.id === selectedNode}
+      isHovered={treeNode.id === hoveredNode}
+      onClick={() => {
+        onSelectNode(treeNode.id);
+      }}
+      item={treeNode}>
+      {/*{arrow}*/}
+      {defaultIcon(treeNode)}
+      <HighlightedText text={treeNode.name} />
+    </TreeItem>
+  );
+}
+
+const TreeItem = styled.li<{
+  item: TreeNode;
+  isHovered: boolean;
+  isSelected: boolean;
+}>(({item, isHovered, isSelected}) => ({
+  display: 'flex',
+  alignItems: 'center',
+  height: '26px',
+  paddingLeft: `${(item.depth + 1) * renderDepthOffset}px`,
+  borderWidth: '1px',
+  borderRadius: '3px',
+  borderColor: isHovered ? theme.selectionBackgroundColor : 'transparent',
+  borderStyle: 'solid',
+  backgroundColor: isSelected ? theme.selectionBackgroundColor : theme.white,
+}));
+
+function HighlightedText(props: {text: string}) {
+  const highlightManager: HighlightManager = useHighlighter();
+  return <span>{highlightManager.render(props.text)}</span>;
+}
+
+function defaultIcon(node: UINode) {
+  if (node.tags.includes('Litho')) {
+    return <DecorationImage src="icons/litho-logo.png" />;
+  }
+}
+
+const DecorationImage = styled.img({
+  height: 12,
+  marginRight: 5,
+  width: 12,
+});
+
+const renderDepthOffset = 4;
+
+function toTreeList(
+  nodes: Map<Id, UINode>,
+  rootId: Id,
+  expanded: Set<Id>,
+): TreeNode[] {
+  const stack = [[nodes.get(rootId), 0]] as [UINode, number][];
+
+  const res = [] as TreeNode[];
+
+  while (stack.length > 0) {
+    const [cur, depth] = stack.pop()!!;
+
+    res.push({
+      ...cur,
+      depth,
+    });
+
+    if (expanded.has(cur.id)) {
+      for (const childId of cur.children) {
+        const child = nodes.get(childId);
+        if (child != null) {
+          stack.push([child, depth + 1]);
+        } else {
+          console.log('null', childId);
+        }
+      }
+    }
+  }
+
+  return res;
+}
diff --git a/desktop/plugins/public/ui-debugger/components/main.tsx b/desktop/plugins/public/ui-debugger/components/main.tsx
index 6166fe97d..05c8512e6 100644
--- a/desktop/plugins/public/ui-debugger/components/main.tsx
+++ b/desktop/plugins/public/ui-debugger/components/main.tsx
@@ -20,6 +20,7 @@ import {Inspector} from './sidebar/Inspector';
 import {Controls} from './Controls';
 import {Input, Spin} from 'antd';
 import FeedbackRequest from './fb-stubs/feedback';
+import {Tree2} from './Tree2';
 
 export function Component() {
   const instance = usePlugin(plugin);
@@ -44,7 +45,7 @@ export function Component() {
         <Layout.Horizontal grow pad="small" gap="small">
           <Layout.Container grow gap="small">
             <Layout.ScrollContainer>
-              <Tree
+              <Tree2
                 selectedNode={selectedNode}
                 onSelectNode={setSelectedNode}
                 nodes={nodes}
diff --git a/desktop/plugins/public/ui-debugger/index.tsx b/desktop/plugins/public/ui-debugger/index.tsx
index 441a3c900..1f4a4f2d7 100644
--- a/desktop/plugins/public/ui-debugger/index.tsx
+++ b/desktop/plugins/public/ui-debugger/index.tsx
@@ -25,6 +25,7 @@ import {
   UINode,
 } from './types';
 import './node_modules/react-complex-tree/lib/style.css';
+import {Draft} from 'immer';
 
 type SnapshotInfo = {nodeId: Id; base64Image: Snapshot};
 type LiveClientState = {
@@ -38,7 +39,7 @@ type UIState = {
   isContextMenuOpen: Atom<boolean>;
   hoveredNodes: Atom<Id[]>;
   focusedNode: Atom<Id | undefined>;
-  treeState: Atom<TreeState>;
+  expandedNodes: Atom<Set<Id>>;
 };
 
 export function plugin(client: PluginClient<Events>) {
@@ -84,7 +85,7 @@ export function plugin(client: PluginClient<Events>) {
 
     searchTerm: createState<string>(''),
     focusedNode: createState<Id | undefined>(undefined),
-    treeState: createState<TreeState>({expandedNodes: []}),
+    expandedNodes: createState<Set<Id>>(new Set()),
   };
 
   client.onMessage('coordinateUpdate', (event) => {
@@ -110,7 +111,7 @@ export function plugin(client: PluginClient<Events>) {
     if (!isPaused) {
       //When going back to play mode then set the atoms to the live state to rerender the latest
       //Also need to fixed expanded state for any change in active child state
-      uiState.treeState.update((draft) => {
+      uiState.expandedNodes.update((draft) => {
         liveClientData.nodes.forEach((node) => {
           collapseinActiveChildren(node, draft);
         });
@@ -144,10 +145,10 @@ export function plugin(client: PluginClient<Events>) {
       setParentPointers(rootId.get()!!, undefined, draft.nodes);
     });
 
-    uiState.treeState.update((draft) => {
+    uiState.expandedNodes.update((draft) => {
       for (const node of event.nodes) {
         if (!seenNodes.has(node.id)) {
-          draft.expandedNodes.push(node.id);
+          draft.add(node.id);
         }
         seenNodes.add(node.id);
 
@@ -193,17 +194,7 @@ function setParentPointers(
   });
 }
 
-function checkFocusedNodeStillActive(
-  uiState: {
-    isPaused: Atom<boolean>;
-    searchTerm: Atom<string>;
-    isContextMenuOpen: Atom<boolean>;
-    hoveredNodes: Atom<Id[]>;
-    focusedNode: Atom<Id | undefined>;
-    treeState: Atom<TreeState>;
-  },
-  nodes: Map<Id, UINode>,
-) {
+function checkFocusedNodeStillActive(uiState: UIState, nodes: Map<Id, UINode>) {
   const focusedNodeId = uiState.focusedNode.get();
   const focusedNode = focusedNodeId && nodes.get(focusedNodeId);
   if (focusedNode && !isFocusedNodeAncestryAllActive(focusedNode, nodes)) {
@@ -239,16 +230,14 @@ function isFocusedNodeAncestryAllActive(
   return false;
 }
 
-function collapseinActiveChildren(node: UINode, draft: TreeState) {
+function collapseinActiveChildren(node: UINode, expandedNodes: Draft<Set<Id>>) {
   if (node.activeChild) {
-    const inactiveChildren = node.children.filter(
-      (child) => child !== node.activeChild,
-    );
-
-    draft.expandedNodes = draft.expandedNodes.filter(
-      (nodeId) => !inactiveChildren.includes(nodeId),
-    );
-    draft.expandedNodes.push(node.activeChild);
+    expandedNodes.add(node.activeChild);
+    for (const child of node.children) {
+      if (child !== node.activeChild) {
+        expandedNodes.delete(child);
+      }
+    }
   }
 }