diff --git a/desktop/plugins/public/ui-debugger/components/Controls.tsx b/desktop/plugins/public/ui-debugger/components/Controls.tsx
new file mode 100644
index 000000000..6a5d5fdd9
--- /dev/null
+++ b/desktop/plugins/public/ui-debugger/components/Controls.tsx
@@ -0,0 +1,46 @@
+/**
+ * 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 React from 'react';
+import {plugin} from '../index';
+import {Button, Input, Tooltip} from 'antd';
+import {PauseCircleOutlined, PlayCircleOutlined} from '@ant-design/icons';
+import {usePlugin, useValue, Layout} from 'flipper-plugin';
+
+/**
+ * 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
+ */
+
+export const Controls: React.FC = () => {
+  const instance = usePlugin(plugin);
+  const searchTerm = useValue(instance.uiState.searchTerm);
+  const isPaused = useValue(instance.uiState.isPaused);
+  return (
+    <Layout.Horizontal pad="small" gap="small">
+      <Input
+        value={searchTerm}
+        onChange={(e) => instance.uiState.searchTerm.set(e.target.value)}
+      />
+      <Button
+        type="default"
+        shape="circle"
+        onClick={() => instance.setPlayPause(!instance.uiState.isPaused.get())}
+        icon={
+          <Tooltip
+            title={isPaused ? 'Resume live updates' : 'Pause incoming updates'}>
+            {isPaused ? <PlayCircleOutlined /> : <PauseCircleOutlined />}
+          </Tooltip>
+        }></Button>
+    </Layout.Horizontal>
+  );
+};
diff --git a/desktop/plugins/public/ui-debugger/components/main.tsx b/desktop/plugins/public/ui-debugger/components/main.tsx
index 7d4fb19a9..0d24585b3 100644
--- a/desktop/plugins/public/ui-debugger/components/main.tsx
+++ b/desktop/plugins/public/ui-debugger/components/main.tsx
@@ -17,8 +17,8 @@ import {Tree} from './Tree';
 import {Visualization2D} from './Visualization2D';
 import {useKeyboardModifiers} from '../hooks/useKeyboardModifiers';
 import {Inspector} from './sidebar/Inspector';
-import {Button, Input, Spin, Tooltip} from 'antd';
-import {PauseCircleOutlined, PlayCircleOutlined} from '@ant-design/icons';
+import {Spin} from 'antd';
+import {Controls} from './Controls';
 
 export function Component() {
   const instance = usePlugin(plugin);
@@ -31,10 +31,8 @@ export function Component() {
 
   useHotkeys('ctrl+i', () => setShowPerfStats((show) => !show));
 
-  const searchTerm = useValue(instance.uiState.searchTerm);
   const {ctrlPressed} = useKeyboardModifiers();
 
-  const isPaused = useValue(instance.uiState.isPaused);
   function renderSidebar(
     node: UINode | undefined,
     metadata: Map<MetadataId, Metadata>,
@@ -53,46 +51,31 @@ export function Component() {
 
   if (rootId) {
     return (
-      <Layout.Horizontal grow>
-        <Layout.Container grow pad="medium" gap="small">
-          <Layout.Horizontal padh="small" gap="small">
-            <Input
-              value={searchTerm}
-              onChange={(e) => instance.uiState.searchTerm.set(e.target.value)}
-            />
-            <Button
-              type="default"
-              shape="circle"
-              onClick={() =>
-                instance.setPlayPause(!instance.uiState.isPaused.get())
-              }
-              icon={
-                <Tooltip
-                  title={
-                    isPaused ? 'Resume live updates' : 'Pause incoming updates'
-                  }>
-                  {isPaused ? <PlayCircleOutlined /> : <PauseCircleOutlined />}
-                </Tooltip>
-              }></Button>
-          </Layout.Horizontal>
-          <Layout.ScrollContainer>
-            <Tree
-              selectedNode={selectedNode}
-              onSelectNode={setSelectedNode}
-              nodes={nodes}
-              rootId={rootId}
-            />
-          </Layout.ScrollContainer>
-        </Layout.Container>
-        <Visualization2D
-          rootId={rootId}
-          nodes={nodes}
-          selectedNode={selectedNode}
-          onSelectNode={setSelectedNode}
-          modifierPressed={ctrlPressed}
-        />
-        {selectedNode && renderSidebar(nodes.get(selectedNode), metadata)}
-      </Layout.Horizontal>
+      <Layout.Container grow padh="small" padv="medium">
+        <Controls />
+        <Layout.Horizontal grow pad="small" gap="small">
+          <Layout.Container grow gap="small">
+            <Layout.ScrollContainer>
+              <Tree
+                selectedNode={selectedNode}
+                onSelectNode={setSelectedNode}
+                nodes={nodes}
+                rootId={rootId}
+              />
+            </Layout.ScrollContainer>
+          </Layout.Container>
+
+          <Visualization2D
+            rootId={rootId}
+            nodes={nodes}
+            selectedNode={selectedNode}
+            onSelectNode={setSelectedNode}
+            modifierPressed={ctrlPressed}
+          />
+
+          {selectedNode && renderSidebar(nodes.get(selectedNode), metadata)}
+        </Layout.Horizontal>
+      </Layout.Container>
     );
   }
 
diff --git a/desktop/plugins/public/ui-debugger/index.tsx b/desktop/plugins/public/ui-debugger/index.tsx
index 9e1f09d87..441a3c900 100644
--- a/desktop/plugins/public/ui-debugger/index.tsx
+++ b/desktop/plugins/public/ui-debugger/index.tsx
@@ -32,6 +32,15 @@ type LiveClientState = {
   nodes: Map<Id, UINode>;
 };
 
+type UIState = {
+  isPaused: Atom<boolean>;
+  searchTerm: Atom<string>;
+  isContextMenuOpen: Atom<boolean>;
+  hoveredNodes: Atom<Id[]>;
+  focusedNode: Atom<Id | undefined>;
+  treeState: Atom<TreeState>;
+};
+
 export function plugin(client: PluginClient<Events>) {
   const rootId = createState<Id | undefined>(undefined);
   const metadata = createState<Map<MetadataId, Metadata>>(new Map());
@@ -63,7 +72,7 @@ export function plugin(client: PluginClient<Events>) {
   const nodes = createState<Map<Id, UINode>>(new Map());
   const snapshot = createState<SnapshotInfo | null>(null);
 
-  const uiState = {
+  const uiState: UIState = {
     //used to disabled hover effects which cause rerenders and mess up the existing context menu
     isContextMenuOpen: createState<boolean>(false),