Add Listener notification in patch_task_instance api endpoint

airflow-core/docs/administration-and-deployment/listeners.rst

@@ -40,6 +40,8 @@ DagRun State Change Events
 --------------------------
 
 DagRun state change events occur when a :class:`~airflow.models.dagrun.DagRun` changes state.
+Beginning with Airflow 3, listeners are also notified whenever a state change is triggered through the API
+(for ``on_dag_run_success`` and ``on_dag_run_failed``) e.g., when a DagRun is marked as success from the Airflow UI.
 
 - ``on_dag_run_running``
 
@@ -66,8 +68,12 @@ DagRun state change events occur when a :class:`~airflow.models.dagrun.DagRun` c
 TaskInstance State Change Events
 --------------------------------
 
-TaskInstance state change events occur when a :class:`~airflow.models.taskinstance.TaskInstance` changes state.
+TaskInstance state change events occur when a :class:`~airflow.sdk.execution_time.task_runner.RuntimeTaskInstance` changes state.
 You can use these events to react to ``LocalTaskJob`` state changes.
+Starting with Airflow 3, listeners are also notified when a state change is triggered through the API
+(for ``on_task_instance_success`` and ``on_task_instance_failed``) e.g., when marking a task instance as success from the Airflow UI.
+In such cases, the listener will receive a :class:`~airflow.models.taskinstance.TaskInstance` instance instead
+of a :class:`~airflow.sdk.execution_time.task_runner.RuntimeTaskInstance` instance.
 
 - ``on_task_instance_running``
 
@@ -182,12 +188,14 @@ For example if you want to implement a listener that uses the ``error`` field in
 List of changes in the listener interfaces since 2.8.0 when they were introduced:
 
 
-+-----------------+--------------------------------------------+-------------------------------------------------------------------------+
-| Airflow Version | Affected method                            | Change                                                                  |
-+=================+============================================+=========================================================================+
-| 2.10.0          | ``on_task_instance_failed``                | An error field added to the interface                                   |
-+-----------------+--------------------------------------------+-------------------------------------------------------------------------+
-| 3.0.0           | ``on_task_instance_running``,              | ``session`` argument removed from task instance listeners,              |
-|                 | ``on_task_instance_success``,              | ``task_instance`` object is now an instance of ``RuntimeTaskInstance``  |
-|                 | ``on_task_instance_failed``                |                                                                         |
-+-----------------+--------------------------------------------+-------------------------------------------------------------------------+
++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+
+| Airflow Version | Affected method                            | Change                                                                                                                        |
++=================+============================================+===============================================================================================================================+
+| 2.10.0          | ``on_task_instance_failed``                | An error field added to the interface                                                                                         |
++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+
+| 3.0.0           | ``on_task_instance_running``               | ``session`` argument removed from task instance listeners,                                                                    |
+|                 |                                            | ``task_instance`` object is now an instance of ``RuntimeTaskInstance``                                                        |
++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+
+| 3.0.0           | ``on_task_instance_failed``,               | ``session`` argument removed from task instance listeners,                                                                    |
+|                 | ``on_task_instance_success``               | ``task_instance`` object is now an instance of ``RuntimeTaskInstance`` when on worker and ``TaskInstance`` when on API server |
++-----------------+--------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------+

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

@@ -19,6 +19,7 @@
 
 from typing import Annotated, Literal, cast
 
+import structlog
 from fastapi import Depends, HTTPException, Query, Request, status
 from fastapi.exceptions import RequestValidationError
 from pydantic import ValidationError
@@ -64,6 +65,7 @@
 from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag
 from airflow.api_fastapi.logging.decorators import action_logging
 from airflow.exceptions import TaskNotFound
+from airflow.listeners.listener import get_listener_manager
 from airflow.models import Base, DagRun
 from airflow.models.dag import DAG
 from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances
@@ -73,6 +75,8 @@
 from airflow.utils.db import get_query_count
 from airflow.utils.state import DagRunState, TaskInstanceState
 
+log = structlog.get_logger(__name__)
+
 task_instances_router = AirflowRouter(tags=["Task Instance"], prefix="/dags/{dag_id}")
 task_instances_prefix = "/dagRuns/{dag_run_id}/taskInstances"
 
@@ -882,6 +886,20 @@ def patch_task_instance(
                     status.HTTP_409_CONFLICT, f"Task id {task_id} is already in {data['new_state']} state"
                 )
             ti = tis[0] if isinstance(tis, list) else tis
+            try:
+                if data["new_state"] == TaskInstanceState.SUCCESS:
+                    get_listener_manager().hook.on_task_instance_success(
+                        previous_state=None, task_instance=ti
+                    )
+                elif data["new_state"] == TaskInstanceState.FAILED:
+                    get_listener_manager().hook.on_task_instance_failed(
+                        previous_state=None,
+                        task_instance=ti,
+                        error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.",
+                    )
+            except Exception:
+                log.exception("error calling listener")
+
         elif key == "note":
             if update_mask or body.note is not None:
                 if ti.task_instance_note is None:

airflow-core/src/airflow/example_dags/plugins/event_listener.py

@@ -20,6 +20,7 @@
 from typing import TYPE_CHECKING
 
 from airflow.listeners import hookimpl
+from airflow.models.taskinstance import TaskInstance
 
 if TYPE_CHECKING:
     from airflow.models.dagrun import DagRun
@@ -31,10 +32,10 @@
 @hookimpl
 def on_task_instance_running(previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance):
     """
-    This method is called when task state changes to RUNNING.
-    Through callback, parameters like previous_task_state, task_instance object can be accessed.
-    This will give more information about current task_instance that is running its dag_run,
-    task and dag information.
+    Called when task state changes to RUNNING.
+
+    previous_task_state and task_instance object can be used to retrieve more information about current
+    task_instance that is running, its dag_run, task and dag information.
     """
     print("Task instance is in running state")
     print(" Previous state of the Task instance:", previous_state)
@@ -61,16 +62,27 @@ def on_task_instance_running(previous_state: TaskInstanceState, task_instance: R
 
 # [START howto_listen_ti_success_task]
 @hookimpl
-def on_task_instance_success(previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance):
+def on_task_instance_success(
+    previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance | TaskInstance
+):
     """
-    This method is called when task state changes to SUCCESS.
-    Through callback, parameters like previous_task_state, task_instance object can be accessed.
-    This will give more information about current task_instance that has succeeded its
-    dag_run, task and dag information.
+    Called when task state changes to SUCCESS.
+
+    previous_task_state and task_instance object can be used to retrieve more information about current
+    task_instance that has succeeded, its dag_run, task and dag information.
+
+    A RuntimeTaskInstance is provided in most cases, except when the task's state change is triggered
+    through the API. In that case, the TaskInstance available on the API server will be provided instead.
     """
     print("Task instance in success state")
     print(" Previous state of the Task instance:", previous_state)
 
+    if isinstance(task_instance, TaskInstance):
+        print("Task instance's state was changed through the API.")
+
+        print(f"Task operator:{task_instance.operator}")
+        return
+
     context = task_instance.get_template_context()
     operator = context["task"]
 
@@ -83,16 +95,29 @@ def on_task_instance_success(previous_state: TaskInstanceState, task_instance: R
 # [START howto_listen_ti_failure_task]
 @hookimpl
 def on_task_instance_failed(
-    previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance, error: None | str | BaseException
+    previous_state: TaskInstanceState,
+    task_instance: RuntimeTaskInstance | TaskInstance,
+    error: None | str | BaseException,
 ):
     """
-    This method is called when task state changes to FAILED.
-    Through callback, parameters like previous_task_state, task_instance object can be accessed.
-    This will give more information about current task_instance that has failed its dag_run,
-    task and dag information.
+    Called when task state changes to FAILED.
+
+    previous_task_state, task_instance object and error can be used to retrieve more information about current
+    task_instance that has failed, its dag_run, task and dag information.
+
+    A RuntimeTaskInstance is provided in most cases, except when the task's state change is triggered
+    through the API. In that case, the TaskInstance available on the API server will be provided instead.
     """
     print("Task instance in failure state")
 
+    if isinstance(task_instance, TaskInstance):
+        print("Task instance's state was changed through the API.")
+
+        print(f"Task operator:{task_instance.operator}")
+        if error:
+            print(f"Failure caused by {error}")
+        return
+
     context = task_instance.get_template_context()
     task = context["task"]
 

airflow-core/src/airflow/listeners/spec/taskinstance.py

@@ -22,6 +22,7 @@
 from pluggy import HookspecMarker
 
 if TYPE_CHECKING:
+    from airflow.models.taskinstance import TaskInstance
     from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
     from airflow.utils.state import TaskInstanceState
 
@@ -34,14 +35,16 @@ def on_task_instance_running(previous_state: TaskInstanceState | None, task_inst
 
 
 @hookspec
-def on_task_instance_success(previous_state: TaskInstanceState | None, task_instance: RuntimeTaskInstance):
+def on_task_instance_success(
+    previous_state: TaskInstanceState | None, task_instance: RuntimeTaskInstance | TaskInstance
+):
     """Execute when task state changes to SUCCESS. previous_state can be None."""
 
 
 @hookspec
 def on_task_instance_failed(
     previous_state: TaskInstanceState | None,
-    task_instance: RuntimeTaskInstance,
+    task_instance: RuntimeTaskInstance | TaskInstance,
     error: None | str | BaseException,
 ):
     """Execute when task state changes to FAIL. previous_state can be None."""

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

@@ -30,6 +30,7 @@
 from airflow.dag_processing.bundles.manager import DagBundlesManager
 from airflow.jobs.job import Job
 from airflow.jobs.triggerer_job_runner import TriggererJobRunner
+from airflow.listeners.listener import get_listener_manager
 from airflow.models import DagRun, TaskInstance
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.dagbag import DagBag
@@ -3032,6 +3033,39 @@ class TestPatchTaskInstance(TestTaskInstanceEndpoint):
     TASK_ID = "print_the_context"
     RUN_ID = "TEST_DAG_RUN_ID"
 
+    @pytest.fixture(autouse=True)
+    def clean_listener_manager(self):
+        get_listener_manager().clear()
+        yield
+        get_listener_manager().clear()
+
+    @pytest.mark.parametrize(
+        "state, listener_state",
+        [
+            ("success", [TaskInstanceState.SUCCESS]),
+            ("failed", [TaskInstanceState.FAILED]),
+            ("skipped", []),
+        ],
+    )
+    def test_patch_task_instance_notifies_listeners(self, test_client, session, state, listener_state):
+        from unit.listeners.class_listener import ClassBasedListener
+
+        self.create_task_instances(session)
+
+        listener = ClassBasedListener()
+        get_listener_manager().add_listener(listener)
+        test_client.patch(
+            self.ENDPOINT_URL,
+            json={
+                "new_state": state,
+            },
+        )
+
+        response2 = test_client.get(self.ENDPOINT_URL)
+        assert response2.status_code == 200
+        assert response2.json()["state"] == state
+        assert listener.state == listener_state
+
     @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
     def test_should_call_mocked_api(self, mock_set_ti_state, test_client, session):
         self.create_task_instances(session)