From db200661a25431cb999105ffb4d17a2617b84c68 Mon Sep 17 00:00:00 2001
From: Hemkumar Chheda <hemkumar.chheda@astronomer.io>
Date: Wed, 15 Jul 2026 21:54:50 +0530
Subject: [PATCH 1/8] Fix grid/graph view topological sort for group-level and
 cross-group deps

TaskGroup._project_child_deps only looked at a group's own upstream_task_ids,
which stays empty for a direct group-to-group dependency (list or individual
`>>`) and for a task-level dependency that crosses into another group's entry
task. Both cases sorted the group as if it had no upstream at all.

Now also pulls in the group's upstream_group_ids and its root tasks' upstream
task ids before projecting sibling dependencies. Applied to both the
serialization-layer sort and the mirrored design-time sort in task-sdk.

closes: #65291

Related: apache/airflow#67964 (closed for inactivity, written against the
topological_sort implementation before PR #67288/#67688 rewrote it) and
apache/airflow#65639 (draft, same issue, also predates the rewrite).
---
 .../serialization/definitions/taskgroup.py    |  28 ++++-
 .../tests/unit/utils/test_task_group.py       | 101 ++++++++++++++++--
 .../src/airflow/sdk/definitions/taskgroup.py  |  28 ++++-
 .../task_sdk/definitions/test_taskgroup.py    |   3 +-
 4 files changed, 138 insertions(+), 22 deletions(-)

diff --git a/airflow-core/src/airflow/serialization/definitions/taskgroup.py b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
index 65d59cb15f17e..55d1a229d33b4 100644
--- a/airflow-core/src/airflow/serialization/definitions/taskgroup.py
+++ b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
@@ -231,11 +231,12 @@ def topological_sort(self) -> list[DAGNode]:
         nodes = list(children.values())
         n = len(nodes)
         id_to_idx = {nid: i for i, nid in enumerate(children)}
+        group_dict = self.dag.task_group.get_task_group_dict()
 
         projected: list[tuple[int, ...]] = [()] * n
         nodes_with_back_edge = 0
         for i, child in enumerate(nodes):
-            deps = self._project_child_deps(i, child, id_to_idx)
+            deps = self._project_child_deps(i, child, id_to_idx, group_dict)
             if deps:
                 projected[i] = deps
                 if any(d > i for d in deps):
@@ -248,9 +249,24 @@ def topological_sort(self) -> list[DAGNode]:
         return self._sweep_projection(nodes, projected)
 
     def _project_child_deps(
-        self, child_idx: int, child: DAGNode, id_to_idx: dict[str, int]
+        self,
+        child_idx: int,
+        child: DAGNode,
+        id_to_idx: dict[str, int],
+        group_dict: dict[str | None, SerializedTaskGroup],
     ) -> tuple[int, ...]:
-        upstream_ids = child.upstream_task_ids
+        if isinstance(child, SerializedTaskGroup):
+            # A group's own upstream_task_ids only reflects direct group-to-task edges.
+            # Group-to-group edges (`group_a >> group_b`, list or individual) only populate
+            # upstream_group_ids, and task-level edges crossing into the group (a sibling
+            # task feeding one of this group's entry tasks) never touch the group at all —
+            # both need to be pulled in explicitly here.
+            upstream_ids: set[str] = set(child.upstream_task_ids)
+            upstream_ids.update(gid for gid in child.upstream_group_ids if gid is not None)
+            for root_task in child.get_roots():
+                upstream_ids.update(root_task.upstream_task_ids)
+        else:
+            upstream_ids = child.upstream_task_ids
         if not upstream_ids:
             return ()
         sib_deps: set[int] = set()
@@ -260,8 +276,10 @@ def _project_child_deps(
                 if j != child_idx:
                     sib_deps.add(j)
                 continue
-            edge = self.dag.get_task(edge_id)
-            tg = edge.task_group
+            tg = group_dict.get(edge_id)
+            if tg is None:
+                edge = self.dag.get_task(edge_id)
+                tg = edge.task_group
             while tg is not None:
                 anc_idx = id_to_idx.get(tg.node_id)
                 if anc_idx is not None:
diff --git a/airflow-core/tests/unit/utils/test_task_group.py b/airflow-core/tests/unit/utils/test_task_group.py
index c0cd58f30dd80..0b24ac377af92 100644
--- a/airflow-core/tests/unit/utils/test_task_group.py
+++ b/airflow-core/tests/unit/utils/test_task_group.py
@@ -502,17 +502,11 @@ def test_task_group_to_dict_and_dag_edges(dag_maker):
     nodes = task_group_to_dict(dag.task_group)
     edges = dag_edges(dag)
 
+    # group_d depends on group_c (`group_d << group_c`), so it must sort after group_c,
+    # not before task1 as it did prior to the #65291/#67964 topological-sort fix.
     expected_node_id = {
         "id": None,
         "children": [
-            {
-                "id": "group_d",
-                "children": [
-                    {"id": "group_d.task11"},
-                    {"id": "group_d.task12"},
-                    {"id": "group_d.upstream_join_id"},
-                ],
-            },
             {"id": "task1"},
             {
                 "id": "group_a",
@@ -541,6 +535,14 @@ def test_task_group_to_dict_and_dag_edges(dag_maker):
                     {"id": "group_c.downstream_join_id"},
                 ],
             },
+            {
+                "id": "group_d",
+                "children": [
+                    {"id": "group_d.task11"},
+                    {"id": "group_d.task12"},
+                    {"id": "group_d.upstream_join_id"},
+                ],
+            },
             {"id": "task10"},
             {"id": "task9"},
         ],
@@ -728,12 +730,18 @@ def section_2(value2):
     assert dag.task_dict["section_1.section_2.task_4"].downstream_task_ids == {"task_end"}
 
     # Node IDs test
+    # task_start feeds section_1.task_1 directly (a task-level dep crossing into the
+    # group), so section_1 must sort after task_start, not before it — see the
+    # #65291/#67964 topological-sort fix.
     node_ids = {
         "id": None,
         "children": [
+            {"id": "task_start"},
             {
                 "id": "section_1",
                 "children": [
+                    {"id": "section_1.task_1"},
+                    {"id": "section_1.task_2"},
                     {
                         "id": "section_1.section_2",
                         "children": [
@@ -741,12 +749,9 @@ def section_2(value2):
                             {"id": "section_1.section_2.task_4"},
                         ],
                     },
-                    {"id": "section_1.task_1"},
-                    {"id": "section_1.task_2"},
                 ],
             },
             {"id": "task_end"},
-            {"id": "task_start"},
         ],
     }
 
@@ -1187,6 +1192,80 @@ def test_topological_sort_serialized_layered():
                 )
 
 
+def test_topological_group_dep_list_syntax():
+    """List-based deps (`[b0, b1] >> a`) must produce the same topological order as individual deps.
+
+    Regression test for apache/airflow#65291: declaring a group dependency via a list
+    (`groups >> a`) only populated `upstream_group_ids`, which `_project_child_deps` never
+    consulted, so `a` sorted as if it had no upstream at all.
+    """
+    with DAG("test_dag_list_dep", schedule=None, start_date=DEFAULT_DATE) as dag:
+        with TaskGroup("a") as tg_a:
+            EmptyOperator(task_id="task")
+
+        groups = []
+        for x in range(3):
+            with TaskGroup(f"b_{x}") as tg_b:
+                EmptyOperator(task_id="task")
+            groups.append(tg_b)
+
+        groups >> tg_a  # list-based dep — previously produced the wrong order
+
+    order = [node.node_id for node in dag.task_group.topological_sort()]
+    a_idx = order.index("a")
+    assert all(order.index(f"b_{x}") < a_idx for x in range(3)), (
+        f"Expected all b_x before a in topological order, got: {order!r}"
+    )
+
+
+def test_topological_sort_serialized_list_dep_between_groups():
+    """Same as test_topological_group_dep_list_syntax, exercised on the serialized variant."""
+    with DAG("test_dag_list_dep_serialized", schedule=None, start_date=DEFAULT_DATE) as dag:
+        with TaskGroup("a"):
+            EmptyOperator(task_id="task")
+
+        groups = []
+        for x in range(3):
+            with TaskGroup(f"b_{x}") as tg_b:
+                EmptyOperator(task_id="task")
+            groups.append(tg_b)
+
+        groups >> dag.task_group.children["a"]
+
+    serialized = create_scheduler_dag(dag)
+    order = [node.node_id for node in serialized.task_group.topological_sort()]
+    a_idx = order.index("a")
+    assert all(order.index(f"b_{x}") < a_idx for x in range(3)), (
+        f"Expected all b_x before a in topological order, got: {order!r}"
+    )
+
+
+def test_topological_sort_serialized_task_level_cross_group_dep():
+    """Task-level deps between groups are respected for ordering after serialization.
+
+    Regression test for apache/airflow#67964: a task-level dependency that crosses into
+    another group's entry task (bypassing any group-to-group edge) must still order the
+    downstream group after the upstream one.
+    """
+    with DAG("test_cross_group_task_dep", schedule=None, start_date=DEFAULT_DATE) as dag:
+        with TaskGroup("stage_b"):
+            b_start = EmptyOperator(task_id="b_start")
+            b_end = EmptyOperator(task_id="b_end")
+            b_start >> b_end
+
+        with TaskGroup("stage_a"):
+            a_start = EmptyOperator(task_id="a_start")
+            a_end = EmptyOperator(task_id="a_end")
+            a_start >> a_end
+
+        b_end >> a_start
+
+    serialized = create_scheduler_dag(dag)
+    order = [node.node_id for node in serialized.task_group.topological_sort()]
+
+    assert order.index("stage_b") < order.index("stage_a")
+
+
 def test_topological_sort_serialized_padded_reverse_chain_uses_pass_numbering(monkeypatch):
     dag = _make_padded_reverse_chain(chain_length=80, independent_count=80)
     serialized = create_scheduler_dag(dag)
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index c1be99c5766bd..a1794d3cf0ee0 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -569,11 +569,12 @@ def topological_sort(self) -> list[DAGNode]:
         nodes = list(children.values())
         n = len(nodes)
         id_to_idx = {nid: i for i, nid in enumerate(children)}
+        group_dict = self.dag.task_group.get_task_group_dict()
 
         projected: list[tuple[int, ...]] = [()] * n
         nodes_with_back_edge = 0
         for i, child in enumerate(nodes):
-            deps = self._project_child_deps(i, child, id_to_idx)
+            deps = self._project_child_deps(i, child, id_to_idx, group_dict)
             if deps:
                 projected[i] = deps
                 if any(d > i for d in deps):
@@ -586,9 +587,24 @@ def topological_sort(self) -> list[DAGNode]:
         return self._sweep_projection(nodes, projected)
 
     def _project_child_deps(
-        self, child_idx: int, child: DAGNode, id_to_idx: dict[str, int]
+        self,
+        child_idx: int,
+        child: DAGNode,
+        id_to_idx: dict[str, int],
+        group_dict: dict[str, TaskGroup],
     ) -> tuple[int, ...]:
-        upstream_ids = child.upstream_task_ids
+        if isinstance(child, TaskGroup):
+            # A group's own upstream_task_ids only reflects direct group-to-task edges.
+            # Group-to-group edges (`group_a >> group_b`, list or individual) only populate
+            # upstream_group_ids, and task-level edges crossing into the group (a sibling
+            # task feeding one of this group's entry tasks) never touch the group at all —
+            # both need to be pulled in explicitly here.
+            upstream_ids: set[str] = set(child.upstream_task_ids)
+            upstream_ids.update(gid for gid in child.upstream_group_ids if gid is not None)
+            for root_task in child.get_roots():
+                upstream_ids.update(root_task.upstream_task_ids)
+        else:
+            upstream_ids = child.upstream_task_ids
         if not upstream_ids:
             return ()
         sib_deps: set[int] = set()
@@ -598,8 +614,10 @@ def _project_child_deps(
                 if j != child_idx:
                     sib_deps.add(j)
                 continue
-            edge = self.dag.get_task(edge_id)
-            tg = edge.task_group
+            tg = group_dict.get(edge_id)
+            if tg is None:
+                edge = self.dag.get_task(edge_id)
+                tg = edge.task_group
             while tg is not None:
                 anc_idx = id_to_idx.get(tg.node_id)
                 if anc_idx is not None:
diff --git a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
index 485346d6d28a8..98114598db9ee 100644
--- a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
+++ b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
@@ -1094,7 +1094,8 @@ def test_topological_sort_reverse_declared_order_matches_sweep():
     group = dag.task_group
     nodes = list(group.children.values())
     id_to_idx = {nid: i for i, nid in enumerate(group.children)}
-    projected = [group._project_child_deps(i, child, id_to_idx) for i, child in enumerate(nodes)]
+    group_dict = group.dag.task_group.get_task_group_dict()
+    projected = [group._project_child_deps(i, child, id_to_idx, group_dict) for i, child in enumerate(nodes)]
 
     sweep_order = [node.node_id for node in group._sweep_projection(nodes, projected)]
     pass_number_order = [node.node_id for node in group._sort_via_pass_numbering(nodes, projected)]

From ee77c2d6f2c7a7811495f386a536348710761dba Mon Sep 17 00:00:00 2001
From: Hemkumar Chheda <hemkumar.chheda@astronomer.io>
Date: Sat, 18 Jul 2026 00:52:14 +0530
Subject: [PATCH 2/8] Address review feedback: cache get_task_group_dict,
 describe test intent not issue numbers

viiccwen pointed out that fetching the group map inside topological_sort() rebuilds
the whole DAG's group tree on every nested group's own call, turning a render with
G groups into an O(G^2) cost. get_task_group_dict() is now memoized per DAG instance
(kept behind a small private helper since methodtools.lru_cache has no type stubs and
would otherwise widen the public method's return type to Any for every caller).

Also reworded test comments/docstrings that cited issue numbers to describe what's
actually being verified instead.
---
 .../serialization/definitions/taskgroup.py    |  8 +++-
 .../tests/unit/utils/test_task_group.py       | 38 ++++++++++++++-----
 .../src/airflow/sdk/definitions/taskgroup.py  | 10 ++++-
 .../task_sdk/definitions/test_taskgroup.py    | 19 ++++++++++
 4 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/airflow-core/src/airflow/serialization/definitions/taskgroup.py b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
index 55d1a229d33b4..9562f47d0f77d 100644
--- a/airflow-core/src/airflow/serialization/definitions/taskgroup.py
+++ b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
@@ -173,8 +173,14 @@ def recurse_for_first_non_teardown(task):
                     yield from recurse_for_first_non_teardown(task)
 
     def get_task_group_dict(self) -> dict[str | None, SerializedTaskGroup]:
-        """Create a flat dict of group_id: TaskGroup."""
+        """Create a flat dict of group_id: TaskGroup. Cached per instance/DAG."""
+        return self._get_task_group_dict_cached()
 
+    # methodtools.lru_cache has no type stubs, so it widens this method's return type to
+    # Any for every caller; kept private behind the explicitly-typed wrapper above so mypy
+    # still trusts get_task_group_dict()'s declared return type.
+    @methodtools.lru_cache(maxsize=None)
+    def _get_task_group_dict_cached(self) -> dict[str | None, SerializedTaskGroup]:
         def build_map(node: DAGNode) -> Generator[tuple[str | None, SerializedTaskGroup]]:
             if not isinstance(node, SerializedTaskGroup):
                 return
diff --git a/airflow-core/tests/unit/utils/test_task_group.py b/airflow-core/tests/unit/utils/test_task_group.py
index 0b24ac377af92..8d815ffae0830 100644
--- a/airflow-core/tests/unit/utils/test_task_group.py
+++ b/airflow-core/tests/unit/utils/test_task_group.py
@@ -502,8 +502,8 @@ def test_task_group_to_dict_and_dag_edges(dag_maker):
     nodes = task_group_to_dict(dag.task_group)
     edges = dag_edges(dag)
 
-    # group_d depends on group_c (`group_d << group_c`), so it must sort after group_c,
-    # not before task1 as it did prior to the #65291/#67964 topological-sort fix.
+    # group_d depends on group_c (`group_d << group_c`), so it must sort after group_c
+    # rather than before task1, which has no dependency on it at all.
     expected_node_id = {
         "id": None,
         "children": [
@@ -731,8 +731,7 @@ def section_2(value2):
 
     # Node IDs test
     # task_start feeds section_1.task_1 directly (a task-level dep crossing into the
-    # group), so section_1 must sort after task_start, not before it — see the
-    # #65291/#67964 topological-sort fix.
+    # group), so section_1 must sort after task_start, not before it.
     node_ids = {
         "id": None,
         "children": [
@@ -1195,9 +1194,9 @@ def test_topological_sort_serialized_layered():
 def test_topological_group_dep_list_syntax():
     """List-based deps (`[b0, b1] >> a`) must produce the same topological order as individual deps.
 
-    Regression test for apache/airflow#65291: declaring a group dependency via a list
-    (`groups >> a`) only populated `upstream_group_ids`, which `_project_child_deps` never
-    consulted, so `a` sorted as if it had no upstream at all.
+    Declaring a group dependency via a list (`groups >> a`) only populates
+    `upstream_group_ids`, not `upstream_task_ids`, so `a` must not sort as if it had no
+    upstream at all.
     """
     with DAG("test_dag_list_dep", schedule=None, start_date=DEFAULT_DATE) as dag:
         with TaskGroup("a") as tg_a:
@@ -1243,9 +1242,8 @@ def test_topological_sort_serialized_list_dep_between_groups():
 def test_topological_sort_serialized_task_level_cross_group_dep():
     """Task-level deps between groups are respected for ordering after serialization.
 
-    Regression test for apache/airflow#67964: a task-level dependency that crosses into
-    another group's entry task (bypassing any group-to-group edge) must still order the
-    downstream group after the upstream one.
+    A task-level dependency that crosses into another group's entry task (bypassing any
+    group-to-group edge) must still order the downstream group after the upstream one.
     """
     with DAG("test_cross_group_task_dep", schedule=None, start_date=DEFAULT_DATE) as dag:
         with TaskGroup("stage_b"):
@@ -1266,6 +1264,26 @@ def test_topological_sort_serialized_task_level_cross_group_dep():
     assert order.index("stage_b") < order.index("stage_a")
 
 
+def test_topological_sort_serialized_reuses_cached_group_dict():
+    with DAG("test_group_dict_cache", schedule=None, start_date=DEFAULT_DATE) as dag:
+        with TaskGroup("a"):
+            EmptyOperator(task_id="task")
+        with TaskGroup("b"):
+            EmptyOperator(task_id="task")
+
+    serialized = create_scheduler_dag(dag)
+    root = serialized.task_group
+    assert root.get_task_group_dict() is root.get_task_group_dict()
+    assert root._get_task_group_dict_cached.cache_info().misses == 1
+
+    for group in root.children.values():
+        if hasattr(group, "topological_sort"):
+            group.topological_sort()
+    cache_info = root._get_task_group_dict_cached.cache_info()
+    assert cache_info.misses == 1
+    assert cache_info.hits >= len(root.children)
+
+
 def test_topological_sort_serialized_padded_reverse_chain_uses_pass_numbering(monkeypatch):
     dag = _make_padded_reverse_chain(chain_length=80, independent_count=80)
     serialized = create_scheduler_dag(dag)
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index a1794d3cf0ee0..4efec64d283c4 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -27,6 +27,7 @@
 from typing import TYPE_CHECKING, Any
 
 import attrs
+import methodtools
 
 from airflow.sdk import TriggerRule
 from airflow.sdk.definitions._internal.node import DAGNode, validate_group_key
@@ -499,7 +500,14 @@ def downstream_join_id(self) -> str:
         return f"{self.group_id}.downstream_join_id"
 
     def get_task_group_dict(self) -> dict[str, TaskGroup]:
-        """Return a flat dictionary of group_id: TaskGroup."""
+        """Return a flat dictionary of group_id: TaskGroup. Cached per instance/DAG."""
+        return self._get_task_group_dict_cached()
+
+    # methodtools.lru_cache has no type stubs, so it widens this method's return type to
+    # Any for every caller; kept private behind the explicitly-typed wrapper above so mypy
+    # still trusts get_task_group_dict()'s declared return type.
+    @methodtools.lru_cache(maxsize=None)
+    def _get_task_group_dict_cached(self) -> dict[str, TaskGroup]:
         task_group_map = {}
 
         def build_map(task_group):
diff --git a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
index 98114598db9ee..116397313854e 100644
--- a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
+++ b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
@@ -1103,6 +1103,25 @@ def test_topological_sort_reverse_declared_order_matches_sweep():
     assert pass_number_order == sweep_order
 
 
+def test_topological_sort_reuses_cached_group_dict():
+    with DAG("test_group_dict_cache", schedule=None, start_date=DEFAULT_DATE) as test_dag:
+        with TaskGroup("a"):
+            EmptyOperator(task_id="task")
+        with TaskGroup("b"):
+            EmptyOperator(task_id="task")
+
+    root = test_dag.task_group
+    assert root.get_task_group_dict() is root.get_task_group_dict()
+    assert root._get_task_group_dict_cached.cache_info().misses == 1
+
+    for group in root.children.values():
+        if isinstance(group, TaskGroup):
+            group.topological_sort()
+    cache_info = root._get_task_group_dict_cached.cache_info()
+    assert cache_info.misses == 1
+    assert cache_info.hits >= len(root.children)
+
+
 def test_topological_sort_padded_reverse_chain_uses_pass_numbering(monkeypatch):
     dag = _make_padded_reverse_chain(chain_length=80, independent_count=80)
     called = {"value": False}

From 0ebe633a9fd99296b0303e8db718a81140123f5c Mon Sep 17 00:00:00 2001
From: TP <uranusjr@apache.org>
Date: Tue, 28 Jul 2026 08:57:20 +0800
Subject: [PATCH 3/8] Hoist common logic into shared lib

---
 .../serialization/definitions/taskgroup.py    | 16 ++------
 .../src/airflow_shared/dagnode/node.py        | 38 +++++++++++++++++++
 .../src/airflow/sdk/definitions/taskgroup.py  | 16 ++------
 3 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/airflow-core/src/airflow/serialization/definitions/taskgroup.py b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
index 9562f47d0f77d..500e488efd823 100644
--- a/airflow-core/src/airflow/serialization/definitions/taskgroup.py
+++ b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
@@ -27,6 +27,7 @@
 import attrs
 import methodtools
 
+from airflow._shared.dagnode.node import TaskGroupMixin
 from airflow.serialization.definitions.node import DAGNode
 
 if TYPE_CHECKING:
@@ -38,7 +39,7 @@
 
 
 @attrs.define(eq=False, hash=False, kw_only=True)
-class SerializedTaskGroup(DAGNode):
+class SerializedTaskGroup(TaskGroupMixin, DAGNode):
     """Serialized representation of a TaskGroup used in protected processes."""
 
     _group_id: str | None = attrs.field(alias="group_id")
@@ -261,18 +262,7 @@ def _project_child_deps(
         id_to_idx: dict[str, int],
         group_dict: dict[str | None, SerializedTaskGroup],
     ) -> tuple[int, ...]:
-        if isinstance(child, SerializedTaskGroup):
-            # A group's own upstream_task_ids only reflects direct group-to-task edges.
-            # Group-to-group edges (`group_a >> group_b`, list or individual) only populate
-            # upstream_group_ids, and task-level edges crossing into the group (a sibling
-            # task feeding one of this group's entry tasks) never touch the group at all —
-            # both need to be pulled in explicitly here.
-            upstream_ids: set[str] = set(child.upstream_task_ids)
-            upstream_ids.update(gid for gid in child.upstream_group_ids if gid is not None)
-            for root_task in child.get_roots():
-                upstream_ids.update(root_task.upstream_task_ids)
-        else:
-            upstream_ids = child.upstream_task_ids
+        upstream_ids = child._topological_upstream_ids
         if not upstream_ids:
             return ()
         sib_deps: set[int] = set()
diff --git a/shared/dagnode/src/airflow_shared/dagnode/node.py b/shared/dagnode/src/airflow_shared/dagnode/node.py
index 7d52ff1ea1f4d..c1c9200d95515 100644
--- a/shared/dagnode/src/airflow_shared/dagnode/node.py
+++ b/shared/dagnode/src/airflow_shared/dagnode/node.py
@@ -140,6 +140,16 @@ def downstream_list(self) -> Iterable[Task]:
             raise RuntimeError(f"Operator {self} has not been assigned to a Dag yet")
         return [self.dag.get_task(tid) for tid in self.downstream_task_ids]
 
+    @property
+    def _topological_upstream_ids(self) -> Iterable[str]:
+        """
+        Node ids this node must be ordered after within its parent group.
+
+        A plain task depends only on its direct upstream tasks. Task groups override
+        this to also cover group-to-group and cross-group edges.
+        """
+        return self.upstream_task_ids
+
     def has_dag(self) -> bool:
         return self.dag is not None
 
@@ -263,3 +273,31 @@ def get_upstreams_only_setups(self) -> Iterable[Task]:
         for task in self.get_upstreams_only_setups_and_teardowns():
             if task.is_setup:
                 yield task
+
+
+class TaskGroupMixin:
+    """Mixin to host common logic between authored and serialized task group classes."""
+
+    upstream_task_ids: set[str]
+    upstream_group_ids: set[str | None]
+
+    def get_roots(self) -> Iterable[GenericDAGNode]:
+        raise NotImplementedError()
+
+    @property
+    def _topological_upstream_ids(self) -> Iterable[str]:
+        """
+        Node ids this node must be ordered after within its parent group.
+
+        A group's upstream_task_ids only reflects direct task-to-group edges
+        (e.g. ``task >> this_group``). This explicitly pulls in two more cases:
+
+        * Group-to-group edges (e.g. ``another_group >> this_group``).
+        * Task-level edges crossing into the group.
+        """
+        return self.upstream_task_ids.union(
+            (gid for gid in self.upstream_group_ids if gid is not None),
+            (t for root in self.get_roots() for t in root.upstream_task_ids),
+        )
+
+    # TODO: Move more duplicated logic between Core and SDK task group types.
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index 4efec64d283c4..c458c5102833d 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -30,6 +30,7 @@
 import methodtools
 
 from airflow.sdk import TriggerRule
+from airflow.sdk._shared.dagnode.node import TaskGroupMixin
 from airflow.sdk.definitions._internal.node import DAGNode, validate_group_key
 from airflow.sdk.exceptions import (
     AirflowDagCycleException,
@@ -93,7 +94,7 @@ def _convert_doc_md(doc_md: str | None) -> str | None:
 
 
 @attrs.define(repr=False)
-class TaskGroup(DAGNode):
+class TaskGroup(TaskGroupMixin, DAGNode):
     """
     A collection of tasks.
 
@@ -601,18 +602,7 @@ def _project_child_deps(
         id_to_idx: dict[str, int],
         group_dict: dict[str, TaskGroup],
     ) -> tuple[int, ...]:
-        if isinstance(child, TaskGroup):
-            # A group's own upstream_task_ids only reflects direct group-to-task edges.
-            # Group-to-group edges (`group_a >> group_b`, list or individual) only populate
-            # upstream_group_ids, and task-level edges crossing into the group (a sibling
-            # task feeding one of this group's entry tasks) never touch the group at all —
-            # both need to be pulled in explicitly here.
-            upstream_ids: set[str] = set(child.upstream_task_ids)
-            upstream_ids.update(gid for gid in child.upstream_group_ids if gid is not None)
-            for root_task in child.get_roots():
-                upstream_ids.update(root_task.upstream_task_ids)
-        else:
-            upstream_ids = child.upstream_task_ids
+        upstream_ids = child._topological_upstream_ids
         if not upstream_ids:
             return ()
         sib_deps: set[int] = set()

From b6158d4b5c1c225b9f6590e704cc46130de27c43 Mon Sep 17 00:00:00 2001
From: TP <uranusjr@apache.org>
Date: Tue, 28 Jul 2026 09:18:26 +0800
Subject: [PATCH 4/8] Remove caching on get_task_group_dict

---
 .../src/airflow/serialization/definitions/taskgroup.py |  8 +-------
 task-sdk/src/airflow/sdk/definitions/taskgroup.py      | 10 +---------
 2 files changed, 2 insertions(+), 16 deletions(-)

diff --git a/airflow-core/src/airflow/serialization/definitions/taskgroup.py b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
index 500e488efd823..b1510156b00ca 100644
--- a/airflow-core/src/airflow/serialization/definitions/taskgroup.py
+++ b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
@@ -174,14 +174,8 @@ def recurse_for_first_non_teardown(task):
                     yield from recurse_for_first_non_teardown(task)
 
     def get_task_group_dict(self) -> dict[str | None, SerializedTaskGroup]:
-        """Create a flat dict of group_id: TaskGroup. Cached per instance/DAG."""
-        return self._get_task_group_dict_cached()
+        """Create a flat dict of group_id: TaskGroup."""
 
-    # methodtools.lru_cache has no type stubs, so it widens this method's return type to
-    # Any for every caller; kept private behind the explicitly-typed wrapper above so mypy
-    # still trusts get_task_group_dict()'s declared return type.
-    @methodtools.lru_cache(maxsize=None)
-    def _get_task_group_dict_cached(self) -> dict[str | None, SerializedTaskGroup]:
         def build_map(node: DAGNode) -> Generator[tuple[str | None, SerializedTaskGroup]]:
             if not isinstance(node, SerializedTaskGroup):
                 return
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index c458c5102833d..3df18457f2e14 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -27,7 +27,6 @@
 from typing import TYPE_CHECKING, Any
 
 import attrs
-import methodtools
 
 from airflow.sdk import TriggerRule
 from airflow.sdk._shared.dagnode.node import TaskGroupMixin
@@ -501,14 +500,7 @@ def downstream_join_id(self) -> str:
         return f"{self.group_id}.downstream_join_id"
 
     def get_task_group_dict(self) -> dict[str, TaskGroup]:
-        """Return a flat dictionary of group_id: TaskGroup. Cached per instance/DAG."""
-        return self._get_task_group_dict_cached()
-
-    # methodtools.lru_cache has no type stubs, so it widens this method's return type to
-    # Any for every caller; kept private behind the explicitly-typed wrapper above so mypy
-    # still trusts get_task_group_dict()'s declared return type.
-    @methodtools.lru_cache(maxsize=None)
-    def _get_task_group_dict_cached(self) -> dict[str, TaskGroup]:
+        """Return a flat dictionary of group_id: TaskGroup."""
         task_group_map = {}
 
         def build_map(task_group):

From dedf54583e4d84649995f96466a8f718e8106ef4 Mon Sep 17 00:00:00 2001
From: LIU ZHE YOU <zhu424.dev@gmail.com>
Date: Tue, 28 Jul 2026 02:43:24 +0000
Subject: [PATCH 5/8] Remove stale get_task_group_dict cache tests

The cache these tests asserted was removed in the previous commit, so the
identity check and the _get_task_group_dict_cached.cache_info() assertions
no longer apply.
---
 .../tests/unit/utils/test_task_group.py       | 20 -------------------
 .../task_sdk/definitions/test_taskgroup.py    | 19 ------------------
 2 files changed, 39 deletions(-)

diff --git a/airflow-core/tests/unit/utils/test_task_group.py b/airflow-core/tests/unit/utils/test_task_group.py
index 8d815ffae0830..3776a0c9d203d 100644
--- a/airflow-core/tests/unit/utils/test_task_group.py
+++ b/airflow-core/tests/unit/utils/test_task_group.py
@@ -1264,26 +1264,6 @@ def test_topological_sort_serialized_task_level_cross_group_dep():
     assert order.index("stage_b") < order.index("stage_a")
 
 
-def test_topological_sort_serialized_reuses_cached_group_dict():
-    with DAG("test_group_dict_cache", schedule=None, start_date=DEFAULT_DATE) as dag:
-        with TaskGroup("a"):
-            EmptyOperator(task_id="task")
-        with TaskGroup("b"):
-            EmptyOperator(task_id="task")
-
-    serialized = create_scheduler_dag(dag)
-    root = serialized.task_group
-    assert root.get_task_group_dict() is root.get_task_group_dict()
-    assert root._get_task_group_dict_cached.cache_info().misses == 1
-
-    for group in root.children.values():
-        if hasattr(group, "topological_sort"):
-            group.topological_sort()
-    cache_info = root._get_task_group_dict_cached.cache_info()
-    assert cache_info.misses == 1
-    assert cache_info.hits >= len(root.children)
-
-
 def test_topological_sort_serialized_padded_reverse_chain_uses_pass_numbering(monkeypatch):
     dag = _make_padded_reverse_chain(chain_length=80, independent_count=80)
     serialized = create_scheduler_dag(dag)
diff --git a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
index 116397313854e..98114598db9ee 100644
--- a/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
+++ b/task-sdk/tests/task_sdk/definitions/test_taskgroup.py
@@ -1103,25 +1103,6 @@ def test_topological_sort_reverse_declared_order_matches_sweep():
     assert pass_number_order == sweep_order
 
 
-def test_topological_sort_reuses_cached_group_dict():
-    with DAG("test_group_dict_cache", schedule=None, start_date=DEFAULT_DATE) as test_dag:
-        with TaskGroup("a"):
-            EmptyOperator(task_id="task")
-        with TaskGroup("b"):
-            EmptyOperator(task_id="task")
-
-    root = test_dag.task_group
-    assert root.get_task_group_dict() is root.get_task_group_dict()
-    assert root._get_task_group_dict_cached.cache_info().misses == 1
-
-    for group in root.children.values():
-        if isinstance(group, TaskGroup):
-            group.topological_sort()
-    cache_info = root._get_task_group_dict_cached.cache_info()
-    assert cache_info.misses == 1
-    assert cache_info.hits >= len(root.children)
-
-
 def test_topological_sort_padded_reverse_chain_uses_pass_numbering(monkeypatch):
     dag = _make_padded_reverse_chain(chain_length=80, independent_count=80)
     called = {"value": False}

From e764e752c435356d879cc12a053f4182099055b0 Mon Sep 17 00:00:00 2001
From: TP <uranusjr@apache.org>
Date: Tue, 28 Jul 2026 15:21:55 +0800
Subject: [PATCH 6/8] Add call-level task group memo to reduce calc

---
 .../api_fastapi/core_api/routes/ui/grid.py    | 19 +++++--
 .../core_api/routes/ui/structure.py           |  6 ++-
 .../api_fastapi/core_api/services/ui/grid.py  |  7 ++-
 .../core_api/services/ui/task_group.py        | 50 ++++++++++++++-----
 .../serialization/definitions/taskgroup.py    |  7 ++-
 .../src/airflow/sdk/definitions/taskgroup.py  |  5 +-
 6 files changed, 71 insertions(+), 23 deletions(-)

diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
index 9cd5b6e60d9f1..2425f67bcc329 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/grid.py
@@ -200,15 +200,22 @@ def get_dag_structure(
     run_ids = list(session.scalars(dag_runs_select_filter))
 
     task_group_sort = get_task_group_children_getter()
+    latest_group_dict = latest_dag.task_group.get_task_group_dict()
     if not run_ids:
-        nodes = [task_group_to_dict_grid(x) for x in task_group_sort(latest_dag.task_group)]
+        nodes = [
+            task_group_to_dict_grid(x, group_dict=latest_group_dict)
+            for x in task_group_sort(latest_dag.task_group, latest_group_dict)
+        ]
         return [GridNodeResponse(**n) for n in nodes]
 
     # Process and merge the latest serdag first
     merged_nodes: list[dict[str, Any]] = []
-    nodes = [task_group_to_dict_grid(x) for x in task_group_sort(latest_dag.task_group)]
+    nodes = [
+        task_group_to_dict_grid(x, group_dict=latest_group_dict)
+        for x in task_group_sort(latest_dag.task_group, latest_group_dict)
+    ]
     _merge_node_dicts(merged_nodes, nodes)
-    del latest_dag
+    del latest_dag, latest_group_dict
 
     # Process serdags one by one and merge immediately to reduce memory usage.
     # Use yield_per() for streaming results and expunge each serdag after processing
@@ -243,7 +250,11 @@ def get_dag_structure(
                 depth=depth,
             )
         # Merge immediately instead of collecting all Dags in memory
-        nodes = [task_group_to_dict_grid(x) for x in task_group_sort(filtered_dag.task_group)]
+        filtered_group_dict = filtered_dag.task_group.get_task_group_dict()
+        nodes = [
+            task_group_to_dict_grid(x, group_dict=filtered_group_dict)
+            for x in task_group_sort(filtered_dag.task_group, filtered_group_dict)
+        ]
         _merge_node_dicts(merged_nodes, nodes)
 
         session.expunge(serdag)  # to allow garbage collection
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
index 597f44db424bb..a5cb9d1d8d188 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py
@@ -96,7 +96,11 @@ def structure_data(
             depth=depth,
         )
 
-    nodes = [task_group_to_dict(child) for child in dag.task_group.topological_sort()]
+    group_dict = dag.task_group.get_task_group_dict()
+    nodes = [
+        task_group_to_dict(child, group_dict=group_dict)
+        for child in dag.task_group.topological_sort(group_dict=group_dict)
+    ]
     edges = dag_edges(dag)
 
     data = {
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
index a40a6f726e213..12336b7386a44 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
@@ -145,6 +145,7 @@ def _find_aggregates(
     node: SerializedTaskGroup | SerializedBaseOperator | TaskMap,
     parent_node: SerializedTaskGroup | SerializedBaseOperator | TaskMap | None,
     ti_details: Mapping[str, GridNodeAgg],
+    group_dict: dict[str | None, SerializedTaskGroup] | None = None,
 ) -> Iterable[tuple[dict[str, Any], GridNodeAgg]]:
     """Recursively fill the Task Group Map."""
     node_id = node.node_id
@@ -171,10 +172,12 @@ def _find_aggregates(
 
         return
     if isinstance(node, SerializedTaskGroup):
+        if group_dict is None:
+            group_dict = node.dag.task_group.get_task_group_dict()
         children_summary = GridNodeAgg()
-        for child in get_task_group_children_getter()(node):
+        for child in get_task_group_children_getter()(node, group_dict):
             for child_node, child_summary in _find_aggregates(
-                node=child, parent_node=node, ti_details=ti_details
+                node=child, parent_node=node, ti_details=ti_details, group_dict=group_dict
             ):
                 if child_node["parent_id"] == node_id:
                     children_summary.merge(child_summary)
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py
index 139e0844091fe..1d7d9b7ccaf0f 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py
@@ -21,20 +21,20 @@
 
 from collections.abc import Callable
 from functools import cache
-from operator import methodcaller
+from typing import Any
 
 from airflow.configuration import conf
 from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
 from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator, is_mapped
+from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
 
 
 @cache
 def get_task_group_children_getter() -> Callable:
     """Get the Task Group Children Getter for the Dag."""
-    sort_order = conf.get("api", "grid_view_sorting_order")
-    if sort_order == "topological":
-        return methodcaller("topological_sort")
-    return methodcaller("hierarchical_alphabetical_sort")
+    if conf.get("api", "grid_view_sorting_order") == "topological":
+        return lambda task_group, group_dict=None: task_group.topological_sort(group_dict=group_dict)
+    return lambda task_group, group_dict=None: task_group.hierarchical_alphabetical_sort()
 
 
 def _ui_colors(node) -> dict[str, str]:
@@ -44,7 +44,7 @@ def _ui_colors(node) -> dict[str, str]:
     }
 
 
-def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False):
+def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False, group_dict=None):
     """Create a nested dict representation of this TaskGroup and its children used to construct the Graph."""
     if isinstance(task := task_item_or_group, (SerializedBaseOperator, SerializedMappedOperator)):
         # we explicitly want the short task ID here, not the full doted notation if in a group
@@ -65,10 +65,16 @@ def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False):
         return node_operator
 
     task_group = task_item_or_group
+    if group_dict is None:
+        group_dict = task_group.dag.task_group.get_task_group_dict()
     mapped = is_mapped(task_group)
     children = [
-        task_group_to_dict(child, parent_group_is_mapped=parent_group_is_mapped or mapped)
-        for child in get_task_group_children_getter()(task_group)
+        task_group_to_dict(
+            child,
+            parent_group_is_mapped=parent_group_is_mapped or mapped,
+            group_dict=group_dict,
+        )
+        for child in get_task_group_children_getter()(task_group, group_dict)
     ]
 
     if task_group.upstream_group_ids or task_group.upstream_task_ids:
@@ -91,8 +97,22 @@ def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False):
     return node
 
 
-def task_group_to_dict_grid(task_item_or_group, parent_group_is_mapped=False):
-    """Create a nested dict representation of this TaskGroup and its children used to construct the Grid."""
+def task_group_to_dict_grid(
+    task_item_or_group,
+    *,
+    group_dict: dict[str | None, SerializedTaskGroup] | None = None,
+    parent_group_is_mapped: bool = False,
+) -> dict[str, Any]:
+    """
+    Create a nested dict representation of this TaskGroup and its children used to construct the Grid.
+
+    :param group_dict: A ``{group_id: group}`` map used to resolve cross-group
+        dependencies. Built once at the top of a render and threaded through the
+        recursion so nested groups reuse it.
+    :param parent_group_is_mapped: Whether an ancestor task group is mapped, propagated to children.
+    """
+    node: dict[str, Any]
+
     if isinstance(task := task_item_or_group, (SerializedMappedOperator, SerializedBaseOperator)):
         mapped = None
         if parent_group_is_mapped or is_mapped(task):
@@ -114,11 +134,17 @@ def task_group_to_dict_grid(task_item_or_group, parent_group_is_mapped=False):
         return node
 
     task_group = task_item_or_group
+    if group_dict is None:
+        group_dict = task_group.dag.task_group.get_task_group_dict()
     task_group_sort = get_task_group_children_getter()
     mapped = is_mapped(task_group)
     children = [
-        task_group_to_dict_grid(x, parent_group_is_mapped=parent_group_is_mapped or mapped)
-        for x in task_group_sort(task_group)
+        task_group_to_dict_grid(
+            child,
+            group_dict=group_dict,
+            parent_group_is_mapped=parent_group_is_mapped or mapped,
+        )
+        for child in task_group_sort(task_group, group_dict)
     ]
 
     node = {
diff --git a/airflow-core/src/airflow/serialization/definitions/taskgroup.py b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
index b1510156b00ca..0e0ae06572dc8 100644
--- a/airflow-core/src/airflow/serialization/definitions/taskgroup.py
+++ b/airflow-core/src/airflow/serialization/definitions/taskgroup.py
@@ -216,7 +216,9 @@ def iter_mapped_task_groups(self) -> Iterator[SerializedMappedTaskGroup]:
                 yield group
             group = group.parent_group
 
-    def topological_sort(self) -> list[DAGNode]:
+    def topological_sort(
+        self, *, group_dict: dict[str | None, SerializedTaskGroup] | None = None
+    ) -> list[DAGNode]:
         """
         Sort children topologically — a task always comes after its upstream dependencies.
 
@@ -232,7 +234,8 @@ def topological_sort(self) -> list[DAGNode]:
         nodes = list(children.values())
         n = len(nodes)
         id_to_idx = {nid: i for i, nid in enumerate(children)}
-        group_dict = self.dag.task_group.get_task_group_dict()
+        if group_dict is None:
+            group_dict = self.dag.task_group.get_task_group_dict()
 
         projected: list[tuple[int, ...]] = [()] * n
         nodes_with_back_edge = 0
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index 3df18457f2e14..89bc37d7127bd 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -549,7 +549,7 @@ def hierarchical_alphabetical_sort(self):
             key=lambda node: (not isinstance(node, TaskGroup), node.node_id),
         )
 
-    def topological_sort(self) -> list[DAGNode]:
+    def topological_sort(self, *, group_dict: dict[str, TaskGroup] | None = None) -> list[DAGNode]:
         """
         Sort children topologically — a task always comes after its upstream dependencies.
 
@@ -570,7 +570,8 @@ def topological_sort(self) -> list[DAGNode]:
         nodes = list(children.values())
         n = len(nodes)
         id_to_idx = {nid: i for i, nid in enumerate(children)}
-        group_dict = self.dag.task_group.get_task_group_dict()
+        if group_dict is None:
+            group_dict = self.dag.task_group.get_task_group_dict()
 
         projected: list[tuple[int, ...]] = [()] * n
         nodes_with_back_edge = 0

From d0e6e2caa940ac70f1fc35a868e53d9c14447716 Mon Sep 17 00:00:00 2001
From: TP <uranusjr@apache.org>
Date: Tue, 28 Jul 2026 15:35:45 +0800
Subject: [PATCH 7/8] Tidy Typy

---
 .../api_fastapi/core_api/services/ui/task_group.py   | 12 ++++++++----
 shared/dagnode/src/airflow_shared/dagnode/node.py    |  4 ++--
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py
index 1d7d9b7ccaf0f..b84a1feef9380 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py
@@ -19,14 +19,18 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 from functools import cache
-from typing import Any
+from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
 from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
 from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator, is_mapped
-from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any
+
+    from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
 
 
 @cache
@@ -44,7 +48,7 @@ def _ui_colors(node) -> dict[str, str]:
     }
 
 
-def task_group_to_dict(task_item_or_group, parent_group_is_mapped=False, group_dict=None):
+def task_group_to_dict(task_item_or_group, *, group_dict=None, parent_group_is_mapped=False):
     """Create a nested dict representation of this TaskGroup and its children used to construct the Graph."""
     if isinstance(task := task_item_or_group, (SerializedBaseOperator, SerializedMappedOperator)):
         # we explicitly want the short task ID here, not the full doted notation if in a group
diff --git a/shared/dagnode/src/airflow_shared/dagnode/node.py b/shared/dagnode/src/airflow_shared/dagnode/node.py
index c1c9200d95515..651547a183bbe 100644
--- a/shared/dagnode/src/airflow_shared/dagnode/node.py
+++ b/shared/dagnode/src/airflow_shared/dagnode/node.py
@@ -141,7 +141,7 @@ def downstream_list(self) -> Iterable[Task]:
         return [self.dag.get_task(tid) for tid in self.downstream_task_ids]
 
     @property
-    def _topological_upstream_ids(self) -> Iterable[str]:
+    def _topological_upstream_ids(self) -> Collection[str]:
         """
         Node ids this node must be ordered after within its parent group.
 
@@ -285,7 +285,7 @@ def get_roots(self) -> Iterable[GenericDAGNode]:
         raise NotImplementedError()
 
     @property
-    def _topological_upstream_ids(self) -> Iterable[str]:
+    def _topological_upstream_ids(self) -> Collection[str]:
         """
         Node ids this node must be ordered after within its parent group.
 

From acf074aaf38cabdad4ab9c9d85f8b2c55e0e2803 Mon Sep 17 00:00:00 2001
From: TP <uranusjr@apache.org>
Date: Tue, 28 Jul 2026 15:37:29 +0800
Subject: [PATCH 8/8] Add test for task group memoing

---
 .../tests/unit/utils/test_task_group.py       | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/airflow-core/tests/unit/utils/test_task_group.py b/airflow-core/tests/unit/utils/test_task_group.py
index 3776a0c9d203d..04b233a37b3d5 100644
--- a/airflow-core/tests/unit/utils/test_task_group.py
+++ b/airflow-core/tests/unit/utils/test_task_group.py
@@ -33,6 +33,7 @@
     task_group as task_group_decorator,
     teardown,
 )
+from airflow.serialization.definitions.taskgroup import SerializedTaskGroup
 from airflow.utils.dag_edges import dag_edges
 
 from tests_common.test_utils.dag import create_scheduler_dag
@@ -252,6 +253,34 @@ def test_task_group_to_dict_alternative_syntax():
     assert task_group_to_dict(serialized_dag.task_group) == EXPECTED_JSON
 
 
+def test_task_group_to_dict_builds_group_dict_once(monkeypatch):
+    """Rendering the whole tree threads one group_dict; it is not rebuilt per nested group."""
+    with DAG("test_group_dict_once", schedule=None, start_date=DEFAULT_DATE) as dag:
+        with TaskGroup("outer"):
+            with TaskGroup("inner"):
+                EmptyOperator(task_id="a")
+                EmptyOperator(task_id="b")
+        with TaskGroup("sibling"):
+            EmptyOperator(task_id="c")
+
+    serialized = create_scheduler_dag(dag)
+
+    calls = 0
+    original = SerializedTaskGroup.get_task_group_dict
+
+    def counting(self):
+        nonlocal calls
+        calls += 1
+        return original(self)
+
+    monkeypatch.setattr(SerializedTaskGroup, "get_task_group_dict", counting)
+
+    # A full render must build the group map exactly once, not once per group.
+    task_group_to_dict(serialized.task_group)
+
+    assert calls == 1
+
+
 def test_task_group_to_dict_grid_includes_task_group_doc_md(dag_maker):
     logical_date = pendulum.parse("20200101")
     with dag_maker("test_task_group_to_dict_doc_md", schedule=None, start_date=logical_date) as dag: