Call listeners for running task instance when a Dag Run state is manually set

airflow-core/src/airflow/api/common/mark_tasks.py

@@ -223,7 +223,7 @@ def _set_dag_run_terminal_state(
     ti_state: TaskInstanceState,
     commit: bool,
     session: SASession,
-) -> list[TaskInstance]:
+) -> tuple[list[TaskInstance], list[TaskInstance]]:
     """
     Set the dag run's state to the given terminal state.
 
@@ -237,11 +237,15 @@ def _set_dag_run_terminal_state(
     :param ti_state: state to set on running task instances
     :param commit: commit Dag and tasks to be altered to the database
     :param session: database session
-    :return: If commit is true, list of tasks that have been updated,
-             otherwise list of tasks that will be updated
+    :return: ``(all_updated_tis, killed_tis)`` where ``all_updated_tis`` is the
+        combined list of pending (now SKIPPED) and running (now ``ti_state``) TIs,
+        and ``killed_tis`` contains only the non-teardown TIs that were in an active
+        running state and were forcefully terminated (teardown TIs are intentionally
+        left running so they can finish their own cleanup, and must not receive a
+        terminal listener event here).
     """
     if not dag:
-        return []
+        return [], []
     if not run_id:
         raise ValueError(f"Invalid dag_run_id: {run_id}")
 
@@ -264,9 +268,10 @@ def _set_dag_run_terminal_state(
             )
         ).all()
     )
-
     # Do not kill teardown tasks
     task_ids_of_running_tis = {ti.task_id for ti in running_tis if not dag.task_dict[ti.task_id].is_teardown}
+    # Keep task instances that were killed with no cleanup capability (all running minus teardown ones).
+    killed_tis = [ti for ti in running_tis if ti.task_id in task_ids_of_running_tis]
 
     def _set_running_task(task: Operator) -> Operator:
         task.dag = dag
@@ -302,13 +307,14 @@ def _set_running_task(task: Operator) -> Operator:
         if not any(dag.task_dict[ti.task_id].is_teardown for ti in (running_tis + pending_tis)):
             _set_dag_run_state(dag.dag_id, run_id, run_state, session)
 
-    return pending_normal_tis + set_state(
+    all_updated = pending_normal_tis + set_state(
         tasks=running_tasks,
         run_id=run_id,
         state=ti_state,
         commit=commit,
         session=session,
     )
+    return all_updated, killed_tis
 
 
 @provide_session
@@ -318,7 +324,7 @@ def set_dag_run_state_to_success(
     run_id: str | None = None,
     commit: bool = False,
     session: SASession = NEW_SESSION,
-) -> list[TaskInstance]:
+) -> tuple[list[TaskInstance], list[TaskInstance]]:
     """
     Set the dag run's state to success.
 
@@ -328,8 +334,15 @@ def set_dag_run_state_to_success(
     :param run_id: the run_id to start looking from
     :param commit: commit Dag and tasks to be altered to the database
     :param session: database session
-    :return: If commit is true, list of tasks that have been updated,
-             otherwise list of tasks that will be updated
+    :return: A tuple ``(all_updated_tis, killed_tis)``.
+        ``all_updated_tis`` is the combined list of task instances that were updated:
+        previously-running ones (now SUCCESS) and non-finished pending ones (now SKIPPED).
+        If ``commit`` is ``False``, this is the list of task instances *that would be* updated.
+        ``killed_tis`` contains only the non-teardown task instances that were in an active
+        running state (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) and were
+        forcefully terminated — teardown tasks are intentionally excluded because they are
+        left running to finish their own cleanup. Use ``killed_tis`` for firing terminal
+        listener hooks.
     :raises: ValueError if dag or logical_date is invalid
     """
     return _set_dag_run_terminal_state(
@@ -349,7 +362,7 @@ def set_dag_run_state_to_failed(
     run_id: str | None = None,
     commit: bool = False,
     session: SASession = NEW_SESSION,
-) -> list[TaskInstance]:
+) -> tuple[list[TaskInstance], list[TaskInstance]]:
     """
     Set the dag run's state to failed.
 
@@ -359,8 +372,15 @@ def set_dag_run_state_to_failed(
     :param run_id: the Dag run_id to start looking from
     :param commit: commit Dag and tasks to be altered to the database
     :param session: database session
-    :return: If commit is true, list of tasks that have been updated,
-             otherwise list of tasks that will be updated
+    :return: A tuple ``(all_updated_tis, killed_tis)``.
+        ``all_updated_tis`` is the combined list of task instances that were updated:
+        previously-running ones (now FAILED) and non-finished pending ones (now SKIPPED).
+        If ``commit`` is ``False``, this is the list of task instances *that would be* updated.
+        ``killed_tis`` contains only the non-teardown task instances that were in an active
+        running state (RUNNING, DEFERRED, UP_FOR_RESCHEDULE, AWAITING_INPUT) and were
+        forcefully terminated — teardown tasks are intentionally excluded because they are
+        left running to finish their own cleanup. Use ``killed_tis`` for firing terminal
+        listener hooks.
     """
     return _set_dag_run_terminal_state(
         dag=dag,

airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py

@@ -57,6 +57,7 @@
 )
 from airflow.api_fastapi.core_api.datamodels.task_instances import NewTaskResponse
 from airflow.api_fastapi.core_api.services.public.common import BulkService
+from airflow.api_fastapi.core_api.services.public.task_instances import _emit_state_listener_hooks
 from airflow.listeners.listener import get_listener_manager
 from airflow.models.dagrun import DagRun, clear_partition_runs
 from airflow.models.taskinstance import TaskInstance
@@ -193,7 +194,10 @@ def patch_dag_run_state(
 ) -> None:
     """Set a Dag Run's state (success/queued/failed), firing the matching listener hooks."""
     if state == DagRunMutableStates.SUCCESS:
-        set_dag_run_state_to_success(dag=dag, run_id=dag_run.run_id, commit=True, session=session)
+        _, killed_tis = set_dag_run_state_to_success(
+            dag=dag, run_id=dag_run.run_id, commit=True, session=session
+        )
+        _emit_state_listener_hooks(killed_tis, TaskInstanceState.SUCCESS)
         try:
             get_listener_manager().hook.on_dag_run_success(dag_run=dag_run, msg="")
         except Exception:
@@ -204,7 +208,10 @@ def patch_dag_run_state(
         # Not notifying on queued - only notifying on RUNNING, which happens in the scheduler.
         set_dag_run_state_to_queued(dag=dag, run_id=dag_run.run_id, commit=True, session=session)
     elif state == DagRunMutableStates.FAILED:
-        set_dag_run_state_to_failed(dag=dag, run_id=dag_run.run_id, commit=True, session=session)
+        _, killed_tis = set_dag_run_state_to_failed(
+            dag=dag, run_id=dag_run.run_id, commit=True, session=session
+        )
+        _emit_state_listener_hooks(killed_tis, TaskInstanceState.FAILED)
         try:
             get_listener_manager().hook.on_dag_run_failed(dag_run=dag_run, msg="")
         except Exception:

airflow-core/tests/unit/api/common/test_mark_tasks.py

@@ -50,9 +50,10 @@ def test_set_dag_run_state_to_failed(dag_maker: DagMaker[SerializedDAG]):
             ti.set_state(TaskInstanceState.RUNNING)
     dag_maker.session.flush()
 
-    updated_tis: list[TaskInstance] = set_dag_run_state_to_failed(
+    result: tuple[list[TaskInstance], list[TaskInstance]] = set_dag_run_state_to_failed(
         dag=dag, run_id=dr.run_id, commit=True, session=dag_maker.session
     )
+    updated_tis, _ = result
     assert len(updated_tis) == 2
     task_dict = {ti.task_id: ti for ti in updated_tis}
     assert task_dict["running"].state == TaskInstanceState.FAILED
@@ -82,9 +83,10 @@ def test_set_dag_run_state_to_success_unfinished_teardown(
     dag_maker.session.flush()
     assert dr.state == DagRunState.RUNNING
 
-    updated_tis: list[TaskInstance] = set_dag_run_state_to_success(
+    result: tuple[list[TaskInstance], list[TaskInstance]] = set_dag_run_state_to_success(
         dag=dag, run_id=dr.run_id, commit=True, session=dag_maker.session
     )
+    updated_tis, _ = result
     run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
     assert run is not None
     assert run.state != DagRunState.SUCCESS
@@ -111,9 +113,10 @@ def test_set_dag_run_state_to_success_keeps_finished_task_states(
     dag_maker.session.flush()
     dr.set_state(DagRunState.FAILED)
 
-    updated_tis: list[TaskInstance] = set_dag_run_state_to_success(
+    result: tuple[list[TaskInstance], list[TaskInstance]] = set_dag_run_state_to_success(
         dag=dag, run_id=dr.run_id, commit=True, session=dag_maker.session
     )
+    updated_tis, _ = result
     run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
     assert run is not None
     assert run.state == DagRunState.SUCCESS

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

@@ -48,7 +48,7 @@
 from airflow.timetables.simple import PartitionedAssetTimetable, PartitionedAtRuntime
 from airflow.timetables.trigger import CronPartitionTimetable
 from airflow.utils.session import provide_session
-from airflow.utils.state import DagRunState, State
+from airflow.utils.state import DagRunState, State, TaskInstanceState
 from airflow.utils.types import DagRunTriggeredByType, DagRunType
 
 from tests_common.test_utils.api_fastapi import _check_dag_run_note, _check_last_log
@@ -1594,20 +1594,139 @@ def test_patch_dag_run_bad_request(self, test_client):
         assert body["detail"][0]["msg"] == "Input should be 'queued', 'success' or 'failed'"
 
     @pytest.mark.parametrize(
-        ("state", "listener_state"),
+        ("state", "expected_dagrun_state"),
         [
-            ("queued", []),
-            ("success", [DagRunState.SUCCESS]),
-            ("failed", [DagRunState.FAILED]),
+            ("queued", None),
+            ("success", DagRunState.SUCCESS),
+            ("failed", DagRunState.FAILED),
         ],
     )
     @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
-    def test_patch_dag_run_notifies_listeners(self, test_client, state, listener_state, listener_manager):
+    def test_patch_dag_run_notifies_listeners(
+        self, test_client, state, expected_dagrun_state, listener_manager
+    ):
         listener = ClassBasedListener()
         listener_manager(listener)
         response = test_client.patch(f"/dags/{DAG1_ID}/dagRuns/{DAG1_RUN1_ID}", json={"state": state})
         assert response.status_code == 200
-        assert listener.state == listener_state
+        assert listener.state == ([expected_dagrun_state] if expected_dagrun_state else [])
+
+    @pytest.mark.parametrize(
+        ("dag_run_state", "expected_ti_state"),
+        [
+            ("success", TaskInstanceState.SUCCESS),
+            ("failed", TaskInstanceState.FAILED),
+        ],
+    )
+    @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+    def test_patch_dag_run_notifies_ti_listeners_for_running_tasks(
+        self,
+        test_client,
+        dag_maker,
+        session,
+        listener_manager,
+        dag_run_state,
+        expected_ti_state,
+    ):
+        with dag_maker(dag_id="test_ti_listeners", schedule=None, serialized=True):
+            EmptyOperator(task_id="t1")
+
+        dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+        ti = session.scalar(
+            select(TaskInstance).where(
+                TaskInstance.dag_id == dr.dag_id,
+                TaskInstance.run_id == dr.run_id,
+                TaskInstance.task_id == "t1",
+            )
+        )
+        ti.state = TaskInstanceState.RUNNING
+        dag_maker.sync_dagbag_to_db()
+        session.commit()
+
+        listener = ClassBasedListener()
+        listener_manager(listener)
+
+        response = test_client.patch(
+            f"/dags/test_ti_listeners/dagRuns/{dr.run_id}", json={"state": dag_run_state}
+        )
+        assert response.status_code == 200
+        # Use `is` for type-safe comparison between TaskInstanceState and DagRunState.
+        assert listener.state[0] is expected_ti_state
+        assert listener.state[1] is DagRunState(dag_run_state)
+        assert len(listener.state) == 2
+
+    @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+    def test_patch_dag_run_does_not_notify_ti_listeners_for_non_running_tasks(
+        self,
+        test_client,
+        dag_maker,
+        session,
+        listener_manager,
+    ):
+        with dag_maker(dag_id="test_ti_listeners_queued", schedule=None, serialized=True):
+            EmptyOperator(task_id="t1")
+
+        dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+        ti = session.scalar(
+            select(TaskInstance).where(
+                TaskInstance.dag_id == dr.dag_id,
+                TaskInstance.run_id == dr.run_id,
+                TaskInstance.task_id == "t1",
+            )
+        )
+        ti.state = TaskInstanceState.QUEUED
+        dag_maker.sync_dagbag_to_db()
+        session.commit()
+
+        listener = ClassBasedListener()
+        listener_manager(listener)
+
+        response = test_client.patch(
+            f"/dags/test_ti_listeners_queued/dagRuns/{dr.run_id}", json={"state": "success"}
+        )
+        assert response.status_code == 200
+        # Only the dagrun-level hook should have fired; no TI hooks for a non-running task.
+        # The list length check distinguishes "only dagrun fired" from "both fired".
+        assert listener.state == [DagRunState.SUCCESS]
+
+    @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+    def test_patch_dag_run_does_not_notify_ti_listeners_for_running_teardown_tasks(
+        self,
+        test_client,
+        dag_maker,
+        session,
+        listener_manager,
+    ):
+        with dag_maker(dag_id="test_ti_listeners_teardown", schedule=None, serialized=True):
+            normal = EmptyOperator(task_id="normal")
+            teardown = EmptyOperator(task_id="teardown").as_teardown(setups=normal)
+            normal >> teardown
+
+        dr = dag_maker.create_dagrun(state=DagRunState.RUNNING)
+        for task_id in ("normal", "teardown"):
+            ti = session.scalar(
+                select(TaskInstance).where(
+                    TaskInstance.dag_id == dr.dag_id,
+                    TaskInstance.run_id == dr.run_id,
+                    TaskInstance.task_id == task_id,
+                )
+            )
+            ti.state = TaskInstanceState.RUNNING
+        dag_maker.sync_dagbag_to_db()
+        session.commit()
+
+        listener = ClassBasedListener()
+        listener_manager(listener)
+
+        response = test_client.patch(
+            f"/dags/test_ti_listeners_teardown/dagRuns/{dr.run_id}", json={"state": "success"}
+        )
+        assert response.status_code == 200
+        # Normal task was killed — its TI listener fires. Teardown task is intentionally skipped.
+        # Use `is` for type-safe comparison between TaskInstanceState and DagRunState.
+        assert len(listener.state) == 2
+        assert listener.state[0] is TaskInstanceState.SUCCESS
+        assert listener.state[1] is DagRunState.SUCCESS
 
 
 class TestDeleteDagRun: