Fix clearing with upstream and downstream selecting unrelated tasks

Clearing a task in a specific Dag run with both upstream and downstream
enabled resolved the downstream relatives from the selected tasks plus
their upstream tasks, so every other descendant of those upstream tasks
was selected too. In a Dag whose branches share a root task this
cleared the whole run, and selecting one map index of a mapped task
cleared all of its map indexes.

Each direction has to be resolved from the requested tasks alone, as
clearing without a specific Dag run already does.

Author: sjyangkevin

airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Collection, Sequence
 from typing import Annotated, Literal, cast
 from uuid import UUID
 
@@ -1090,19 +1090,19 @@ def post_clear_task_instances(
         # Unmapped tasks are expressed in their task_ids (without map_indexes)
         normal_task_ids = {t for t in task_markers_to_clear if not isinstance(t, tuple)}
 
-        def _collect_relatives(run_id: str, direction: Literal["upstream", "downstream"]) -> None:
+        def _find_relatives(
+            run_id: str, direction: Literal["upstream", "downstream"]
+        ) -> Collection[str | tuple[str, int]]:
             from airflow.models.taskinstance import find_relevant_relatives
 
-            relevant_relatives = find_relevant_relatives(
+            return find_relevant_relatives(
                 normal_task_ids,
                 mapped_tasks_tuples,
                 dag=dag,
                 run_id=run_id,
                 direction=direction,
                 session=session,
             )
-            normal_task_ids.update(t for t in relevant_relatives if not isinstance(t, tuple))
-            mapped_tasks_tuples.update(t for t in relevant_relatives if isinstance(t, tuple))
 
         # We can't easily calculate upstream/downstream map indexes when not
         # working for a specific dag run. It's possible by looking at the runs
@@ -1119,10 +1119,15 @@ def _collect_relatives(run_id: str, direction: Literal["upstream", "downstream"]
                 )
                 normal_task_ids.update(partial_dag.task_dict)
         else:
+            # Both directions must start from the requested tasks only. Merging upstream results
+            # before the downstream pass would also select every descendant of those upstream tasks.
+            relatives: set[str | tuple[str, int]] = set()
             if upstream:
-                _collect_relatives(dag_run_id, "upstream")
+                relatives.update(_find_relatives(dag_run_id, "upstream"))
             if downstream:
-                _collect_relatives(dag_run_id, "downstream")
+                relatives.update(_find_relatives(dag_run_id, "downstream"))
+            normal_task_ids.update(t for t in relatives if not isinstance(t, tuple))
+            mapped_tasks_tuples.update(t for t in relatives if isinstance(t, tuple))
 
         task_markers_to_clear = [
             *normal_task_ids,

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py

@@ -5443,6 +5443,45 @@ def test_clear_by_task_group_id_handles_group_larger_than_page_size(
         assert response.status_code == 200
         assert response.json()["total_entries"] == group_size
 
+    @pytest.mark.parametrize(
+        ("mapped", "task_id", "expected"),
+        [
+            pytest.param(False, "a", [("a", -1), ("b", -1), ("root", -1)], id="unmapped"),
+            pytest.param(True, ["a", 0], [("a", 0), ("b", -1), ("root", -1)], id="mapped"),
+        ],
+    )
+    def test_clear_with_upstream_and_downstream_excludes_unrelated_tasks(
+        self, test_client, dag_maker, session, mapped, task_id, expected
+    ):
+        """Upstream + downstream must not include other descendants of the upstream tasks."""
+        dag_id = "clear_upstream_and_downstream"
+        with dag_maker(session=session, dag_id=dag_id, start_date=DEFAULT_DATETIME_1, serialized=True):
+            root = BaseOperator(task_id="root")
+            if mapped:
+                a = MockOperator.partial(task_id="a").expand(arg2=["x", "y", "z"])
+            else:
+                a = BaseOperator(task_id="a")
+            b = BaseOperator(task_id="b")
+            other = BaseOperator(task_id="other")
+            root >> a >> b
+            root >> other
+        dr = dag_maker.create_dagrun(run_id="run1", logical_date=DEFAULT_DATETIME_1)
+
+        response = test_client.post(
+            f"/dags/{dag_id}/clearTaskInstances",
+            json={
+                "dry_run": True,
+                "only_failed": False,
+                "dag_run_id": dr.run_id,
+                "task_ids": [task_id],
+                "include_upstream": True,
+                "include_downstream": True,
+            },
+        )
+        assert response.status_code == 200
+        response_tis = response.json()["task_instances"]
+        assert sorted((ti["task_id"], ti["map_index"]) for ti in response_tis) == expected
+
 
 class TestGetTaskInstanceTries(TestTaskInstanceEndpoint):
     def test_history_without_dag_version_serializes_null(self, test_client, session):