diff --git a/app/packages/core/src/plugins/SchemaIO/components/TreeSelectionView.tsx b/app/packages/core/src/plugins/SchemaIO/components/TreeSelectionView.tsx
index 1a7dac3662..74a187d764 100644
--- a/app/packages/core/src/plugins/SchemaIO/components/TreeSelectionView.tsx
+++ b/app/packages/core/src/plugins/SchemaIO/components/TreeSelectionView.tsx
@@ -3,6 +3,7 @@ import ChevronRightIcon from "@mui/icons-material/ChevronRight";
 import ExpandMoreIcon from "@mui/icons-material/ExpandMore";
 import { Box, Checkbox, FormControlLabel, IconButton } from "@mui/material";
 import React, { useEffect } from "react";
+import { getComponentProps } from "../utils";
 import { ViewPropsType } from "../utils/types";
 
 interface CheckedState {
@@ -73,80 +74,6 @@ export default function TreeSelectionView(props: ViewPropsType) {
 
   const unboundState = useUnboundState(checkedState);
 
-  useEffect(() => {
-    // this only runs when data and checkboxstate are different
-    // meaning the user selected samples from the grid
-    // we will handle the state change of checkedState here
-    // ! do not use onChange here. A change triggered by python should not be sent back to python.
-    if (!data) return;
-
-    // extract the selected sample ids from the unboundState
-    const selectedIdsFromUnboundState = Object.keys(unboundState).filter(
-      (key) => {
-        const isSample =
-          !structure.some(([parentId]) => parentId === key) &&
-          key !== "selectAll";
-        return isSample && unboundState[key].checked; // Only checked samples
-      }
-    );
-
-    const dataSet: Set<string> = new Set(data);
-    const unboundSet: Set<string> = new Set(selectedIdsFromUnboundState);
-    const hasDifference =
-      dataSet.size !== unboundSet.size ||
-      [...dataSet].some((id) => !unboundSet.has(id));
-
-    if (hasDifference) {
-      setCheckedState((prevState) => {
-        const updatedState = { ...prevState };
-
-        // Update the checked state of individual samples based on `dataSet`
-        Object.keys(updatedState).forEach((key) => {
-          if (
-            key !== "selectAll" &&
-            !structure.some(([parentId]) => parentId === key)
-          ) {
-            updatedState[key].checked = dataSet.has(key);
-          }
-        });
-
-        // Update group (parent) states based on their children's state
-        structure.forEach(([parentId, children]) => {
-          const allChildrenChecked = children.every((childId) =>
-            typeof childId === "string"
-              ? updatedState[childId].checked
-              : updatedState[childId[0]].checked
-          );
-          const someChildrenChecked = children.some((childId) =>
-            typeof childId === "string"
-              ? updatedState[childId].checked
-              : updatedState[childId[0]].checked
-          );
-
-          updatedState[parentId] = {
-            checked: allChildrenChecked,
-            indeterminate: someChildrenChecked && !allChildrenChecked,
-          };
-        });
-
-        // Update `selectAll` based on the sample checkboxes
-        const allSamplesChecked = Object.keys(updatedState).every((key) => {
-          const isSample =
-            !structure.some(([parentId]) => parentId === key) &&
-            key !== "selectAll";
-          return !isSample || updatedState[key].checked;
-        });
-
-        updatedState["selectAll"] = {
-          checked: allSamplesChecked,
-          indeterminate: !allSamplesChecked && dataSet.size > 0,
-        };
-
-        return updatedState;
-      });
-    }
-  }, [data, unboundState]);
-
   // Initialize collapsed state for all parents
   const initialCollapsedState: CollapsedState = React.useMemo(() => {
     const state: CollapsedState = {};
@@ -160,6 +87,19 @@ export default function TreeSelectionView(props: ViewPropsType) {
     initialCollapsedState
   );
 
+  const [allCollapsed, setAllCollapsed] = React.useState(false);
+
+  const handleExpandCollapseAll = () => {
+    setCollapsedState((prevState) => {
+      const newState = { ...prevState };
+      Object.keys(newState).forEach((key) => {
+        newState[key] = allCollapsed;
+      });
+      return newState;
+    });
+    setAllCollapsed(!allCollapsed); // Toggle the expand/collapse state
+  };
+
   const handleCheckboxChange = (id: string, isChecked: boolean) => {
     setCheckedState((prevState) => {
       const updatedState = {
@@ -269,6 +209,7 @@ export default function TreeSelectionView(props: ViewPropsType) {
     return idx === -1 ? 0 : idx + 1;
   };
 
+  // On init, all samples are selected by default
   useEffect(() => {
     const sampleIds = view?.data.flatMap(([parentId, children]) => {
       return children.map((childId) =>
@@ -279,6 +220,79 @@ export default function TreeSelectionView(props: ViewPropsType) {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);
 
+  // this only runs when data and checkboxstate are different
+  // meaning the user selected samples from the grid
+  // we will handle the state change of checkedState here
+  // ! do not use onChange here. A change triggered by python should not be sent back to python.
+  useEffect(() => {
+    if (!data) return;
+    // extract the selected sample ids from the unboundState
+    const selectedIdsFromUnboundState = Object.keys(unboundState).filter(
+      (key) => {
+        const isSample =
+          !structure.some(([parentId]) => parentId === key) &&
+          key !== "selectAll";
+        return isSample && unboundState[key].checked; // Only checked samples
+      }
+    );
+
+    const dataSet: Set<string> = new Set(data);
+    const unboundSet: Set<string> = new Set(selectedIdsFromUnboundState);
+    const hasDifference =
+      dataSet.size !== unboundSet.size ||
+      [...dataSet].some((id) => !unboundSet.has(id));
+
+    if (hasDifference) {
+      setCheckedState((prevState) => {
+        const updatedState = { ...prevState };
+
+        // Update the checked state of individual samples based on `dataSet`
+        Object.keys(updatedState).forEach((key) => {
+          if (
+            key !== "selectAll" &&
+            !structure.some(([parentId]) => parentId === key)
+          ) {
+            updatedState[key].checked = dataSet.has(key);
+          }
+        });
+
+        // Update group (parent) states based on their children's state
+        structure.forEach(([parentId, children]) => {
+          const allChildrenChecked = children.every((childId) =>
+            typeof childId === "string"
+              ? updatedState[childId].checked
+              : updatedState[childId[0]].checked
+          );
+          const someChildrenChecked = children.some((childId) =>
+            typeof childId === "string"
+              ? updatedState[childId].checked
+              : updatedState[childId[0]].checked
+          );
+
+          updatedState[parentId] = {
+            checked: allChildrenChecked,
+            indeterminate: someChildrenChecked && !allChildrenChecked,
+          };
+        });
+
+        // Update `selectAll` based on the sample checkboxes
+        const allSamplesChecked = Object.keys(updatedState).every((key) => {
+          const isSample =
+            !structure.some(([parentId]) => parentId === key) &&
+            key !== "selectAll";
+          return !isSample || updatedState[key].checked;
+        });
+
+        updatedState["selectAll"] = {
+          checked: allSamplesChecked,
+          indeterminate: !allSamplesChecked && dataSet.size > 0,
+        };
+
+        return updatedState;
+      });
+    }
+  }, [data, unboundState]);
+
   // CheckboxView: Represents a single checkbox (either parent or child)
   function CheckboxView({
     id,
@@ -313,7 +327,10 @@ export default function TreeSelectionView(props: ViewPropsType) {
   }: TreeNodeProps) {
     const isCollapsed = collapsedState[nodeId] || false;
     const count = childrenIds.length;
-    const title = `Group ${getGroupIdx(nodeId, structure)}  • ${count} Samples`;
+    const title = `Group ${getGroupIdx(
+      nodeId,
+      structure
+    )}  • ${count} Samples • ${nodeId}`;
     return (
       <Box sx={{ display: "flex", flexDirection: "column", ml: 3 }}>
         <Box sx={{ display: "flex", alignItems: "center" }}>
@@ -363,17 +380,20 @@ export default function TreeSelectionView(props: ViewPropsType) {
   }
 
   return (
-    <div>
-      {/* Select All Checkbox */}
-      <CheckboxView
-        id="selectAll"
-        label="Select All"
-        isChecked={checkedState.selectAll.checked}
-        isIndeterminate={checkedState.selectAll.indeterminate || false}
-        onChange={handleCheckboxChange}
-      />
+    <Box {...getComponentProps(props, "container")}>
+      <Box display="flex" alignItems="center">
+        <CheckboxView
+          id="selectAll"
+          label="Select All"
+          isChecked={checkedState.selectAll.checked}
+          isIndeterminate={checkedState.selectAll.indeterminate || false}
+          onChange={handleCheckboxChange}
+        />
+        <IconButton onClick={handleExpandCollapseAll} size="small">
+          {allCollapsed ? <ChevronRightIcon /> : <ExpandMoreIcon />}
+        </IconButton>
+      </Box>
 
-      {/* Render Tree Structure */}
       {structure.map(([parentId, children]) => (
         <TreeNode
           key={parentId}
@@ -385,6 +405,6 @@ export default function TreeSelectionView(props: ViewPropsType) {
           onToggleCollapse={handleToggleCollapse}
         />
       ))}
-    </div>
+    </Box>
   );
 }