Fix grid/graph view topological sort for group-level and cross-group dependencies

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

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 = {

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)

airflow-core/src/airflow/api_fastapi/core_api/services/ui/task_group.py

@@ -19,22 +19,26 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 from functools import cache
-from operator import methodcaller
+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
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any
+
+    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 +48,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, *, 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
@@ -65,10 +69,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 +101,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 +138,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 = {

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")
@@ -215,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.
 
@@ -231,11 +234,13 @@ def topological_sort(self) -> list[DAGNode]:
         nodes = list(children.values())
         n = len(nodes)
         id_to_idx = {nid: i for i, nid in enumerate(children)}
+        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
         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 +253,13 @@ 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
+        upstream_ids = child._topological_upstream_ids
         if not upstream_ids:
             return ()
         sib_deps: set[int] = set()
@@ -260,8 +269,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: