From ca73c8d4060de4c1d23d8fc2bbf743ef1a1e2a14 Mon Sep 17 00:00:00 2001
From: Himanshu Chawla
 <241715894+himanshuchawla0807-gif@users.noreply.github.com>
Date: Mon, 19 Oct 2026 14:49:19 +0530
Subject: [PATCH] Keep parent DAG folders collapsed after one click

---
 .../DagFolderTree/DagFolderTree.test.tsx      | 58 ++++++++++++++++++-
 .../DagsList/DagFolderTree/DagFolderTree.tsx  |  6 +-
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.test.tsx b/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.test.tsx
index 2d2cf2387bb2a..a6bd748270ec1 100644
--- a/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.test.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.test.tsx
@@ -17,7 +17,7 @@
  * under the License.
  */
 import "@testing-library/jest-dom/vitest";
-import { fireEvent, render, screen } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
 
 import type { DagFolderResponse } from "openapi/requests/types.gen";
@@ -48,6 +48,53 @@ const node = (label: string) =>
     (item) => item.querySelector('[data-part="branch-text"], [data-part="item-text"]')?.textContent === label,
   );
 
+it.each([
+  {
+    bundleName: "dags-folder",
+    folder: "team_a",
+    folders: SINGLE_BUNDLE,
+    parent: "team_a",
+    selectedBundle: undefined,
+  },
+  {
+    bundleName: "ml",
+    folder: undefined,
+    folders: MULTI_BUNDLE,
+    parent: "ml",
+    selectedBundle: "ml",
+  },
+])(
+  "collapses $parent in one click when a child is selected",
+  async ({ bundleName, folder, folders, parent, selectedBundle }) => {
+    const onSelectFolder = vi.fn();
+    const { rerender } = render(
+      <DagFolderTree
+        folders={folders}
+        onSelectFolder={onSelectFolder}
+        selectedBundle={selectedBundle}
+        selectedFolder="team_a/etl"
+      />,
+      { wrapper: BaseWrapper },
+    );
+
+    expect(node(parent)).toHaveAttribute("aria-expanded", "true");
+    fireEvent.click(screen.getByTitle(parent), { button: 0 });
+    await waitFor(() => expect(onSelectFolder).toHaveBeenLastCalledWith({ bundleName, folder }));
+
+    rerender(
+      <DagFolderTree
+        folders={folders}
+        onSelectFolder={onSelectFolder}
+        selectedBundle={bundleName}
+        selectedFolder={folder}
+      />,
+    );
+
+    await waitFor(() => expect(node(parent)).toHaveAttribute("aria-expanded", "false"));
+    expect(node(parent)).toHaveAttribute("aria-selected", "true");
+  },
+);
+
 describe("DagFolderTree (single bundle)", () => {
   it("renders the top-level folders and an 'All Dags' entry, without a bundle level", () => {
     render(
@@ -306,6 +353,15 @@ describe("DagFolderTree (multiple bundles)", () => {
     expect(selected.map((item) => item.getAttribute("data-value"))).toEqual(["folder:ml:team_a/etl"]);
   });
 
+  it("keeps a selected bundle collapsed when the folder parameter is empty", () => {
+    render(
+      <DagFolderTree folders={MULTI_BUNDLE} onSelectFolder={vi.fn()} selectedBundle="ml" selectedFolder="" />,
+      { wrapper: BaseWrapper },
+    );
+
+    expect(node("ml")).toHaveAttribute("aria-expanded", "false");
+  });
+
   it("marks a bundle as selected when only the bundle is selected", () => {
     render(
       <DagFolderTree
diff --git a/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.tsx b/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.tsx
index f3790657dba3a..5f537b998749c 100644
--- a/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/DagsList/DagFolderTree/DagFolderTree.tsx
@@ -94,7 +94,7 @@ const expandedForSelection = (
 ): Array<string> => {
   const values: Array<string> = [];
 
-  if (isMultiBundle && bundleName !== undefined) {
+  if (isMultiBundle && bundleName !== undefined && folder !== undefined && folder !== "") {
     values.push(bundleNodeValue(bundleName));
   }
 
@@ -102,7 +102,9 @@ const expandedForSelection = (
     const segments = folder.split("/");
 
     values.push(
-      ...segments.map((_, index) => folderNodeValue(bundleName, segments.slice(0, index + 1).join("/"))),
+      ...segments
+        .slice(0, -1)
+        .map((_, index) => folderNodeValue(bundleName, segments.slice(0, index + 1).join("/"))),
     );
   }