fix: Record history for orphaned tasks during K8s executor failures

Move orphaned task detection before end_date assignment to ensure
TaskInstanceHistory is recorded for tasks that become detached
during scheduler restarts due to Kubernetes API 429 errors.

Author: HsiuChuanHsu

airflow-core/src/airflow/models/taskinstance.py

@@ -1512,6 +1512,35 @@ def fetch_handle_failure_context(
         if not test_mode:
             ti.refresh_from_db(session)
 
+        # Check for orphaned task before setting end_date
+        if (
+            ti.is_eligible_to_retry()
+            and ti.state is None
+            and ti.start_date is not None
+            and ti.end_date is None
+        ):
+            # If the task instance state is None but has a start_date without end_date,
+            # it likely means the task was running but became orphaned and its state was reset.
+            # This can happen during scheduler restarts when executors fail to adopt running tasks
+            # (e.g., due to Kubernetes API 429 errors). We should still record the task instance
+            # history to maintain complete log history for troubleshooting.
+            from airflow.models.taskinstancehistory import TaskInstanceHistory
+
+            log.info(
+                "Recording task instance history for orphaned task %s that was previously running "
+                "(start_date: %s, state reset to None)",
+                ti.key,
+                ti.start_date,
+            )
+            # Temporarily set state to RUNNING to trigger proper history recording
+            original_state = ti.state
+            ti.state = TaskInstanceState.RUNNING
+            try:
+                TaskInstanceHistory.record_ti(ti, session=session)
+            finally:
+                # Restore the original state
+                ti.state = original_state
+
         ti.end_date = timezone.utcnow()
         ti.set_duration()
 
@@ -1547,28 +1576,6 @@ def fetch_handle_failure_context(
                 # about to retry so we record the task instance history. For other states, the task
                 # instance was cleared and already recorded in the task instance history.
                 ti.prepare_db_for_next_try(session)
-            elif ti.state is None and ti.start_date is not None and ti.end_date is None:
-                # If the task instance state is None but has a start_date without end_date,
-                # it likely means the task was running but became orphaned and its state was reset.
-                # This can happen during scheduler restarts when executors fail to adopt running tasks
-                # (e.g., due to Kubernetes API 429 errors). We should still record the task instance
-                # history to maintain complete log history for troubleshooting.
-                from airflow.models.taskinstancehistory import TaskInstanceHistory
-
-                log.info(
-                    "Recording task instance history for orphaned task %s that was previously running "
-                    "(start_date: %s, state reset to None)",
-                    ti.key,
-                    ti.start_date,
-                )
-                # Temporarily set state to RUNNING to trigger proper history recording
-                original_state = ti.state
-                ti.state = TaskInstanceState.RUNNING
-                try:
-                    TaskInstanceHistory.record_ti(ti, session=session)
-                finally:
-                    # Restore the original state
-                    ti.state = original_state
 
             ti.state = State.UP_FOR_RETRY
 

airflow-core/tests/unit/models/test_taskinstance.py

@@ -2424,6 +2424,98 @@ def test_handle_failure_task_undefined(self, create_task_instance):
         del ti.task
         ti.handle_failure("test ti.task undefined")
 
+    @patch("airflow.models.taskinstancehistory.TaskInstanceHistory.record_ti")
+    def test_fetch_handle_failure_context_orphaned_task_records_history(
+        self, mock_record_ti, dag_maker, session
+    ):
+        """
+        Test that orphaned tasks (state=None, start_date!=None, end_date=None) get their history recorded.
+        This scenario occurs when tasks are running but become orphaned due to executor failures
+        (e.g., Kubernetes API 429 errors causing scheduler restarts and pod adoption failures).
+        """
+        with dag_maker(dag_id="test_orphaned_task"):
+            task = EmptyOperator(task_id="orphaned_task", retries=2)  # Allow 2 retries
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.get_task_instance(task.task_id, session=session)
+        ti.task = task
+
+        # Simulate an orphaned task: state=None but has start_date (was running) and no end_date
+        start_time = timezone.utcnow() - datetime.timedelta(minutes=5)
+        ti.state = None  # State was reset during scheduler restart
+        ti.start_date = start_time  # Task had started previously
+        ti.end_date = None  # Task was still running when it became orphaned
+        ti.try_number = 1  # First attempt
+        ti.max_tries = 3  # 1 original + 2 retries = 3 total attempts
+
+        session.merge(ti)
+        session.commit()
+
+        # Call fetch_handle_failure_context which should detect and handle orphaned tasks
+        failure_context = TaskInstance.fetch_handle_failure_context(
+            ti=ti,
+            error="Test orphaned task error",
+            test_mode=False,
+            session=session,
+            fail_fast=False,
+        )
+
+        # Verify that TaskInstanceHistory.record_ti was called for the orphaned task
+        mock_record_ti.assert_called_once()
+        call_args = mock_record_ti.call_args
+        recorded_ti = call_args[0][0]  # First positional argument (ti)
+
+        # Verify the correct TaskInstance was recorded
+        assert recorded_ti.task_id == ti.task_id
+        assert recorded_ti.dag_id == ti.dag_id
+        assert recorded_ti.run_id == ti.run_id
+        assert recorded_ti.start_date == start_time
+
+        # Verify the task instance state is set to UP_FOR_RETRY after failure handling
+        assert ti.state == State.UP_FOR_RETRY
+        assert failure_context["ti"] == ti
+
+    @patch("airflow.models.taskinstancehistory.TaskInstanceHistory.record_ti")
+    def test_fetch_handle_failure_context_orphaned_task_without_start_date_no_history(
+        self, mock_record_ti, dag_maker, session
+    ):
+        """
+        Test that tasks with state=None but no start_date do NOT trigger orphaned task history recording.
+        This ensures we only record history for tasks that were actually running.
+        """
+        with dag_maker(dag_id="test_not_orphaned_task"):
+            task = EmptyOperator(task_id="not_orphaned_task", retries=2)
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.get_task_instance(task.task_id, session=session)
+        ti.task = task
+
+        # Simulate a task that was never started: state=None and no start_date
+        ti.state = None
+        ti.start_date = None  # Task never started
+        ti.end_date = None
+        ti.try_number = 1
+        ti.max_tries = 3  # Allow retries
+
+        session.merge(ti)
+        session.commit()
+
+        # Call fetch_handle_failure_context
+        failure_context = TaskInstance.fetch_handle_failure_context(
+            ti=ti,
+            error="Test non-orphaned task error",
+            test_mode=False,
+            session=session,
+            fail_fast=False,
+        )
+
+        # Verify that TaskInstanceHistory.record_ti was NOT called
+        mock_record_ti.assert_not_called()
+
+        # Verify the task instance state is set to UP_FOR_RETRY after failure handling
+        assert ti.state == State.UP_FOR_RETRY
+        assert failure_context["ti"] == ti
+
     def test_handle_failure_fail_fast(self, dag_maker, session):
         start_date = timezone.datetime(2016, 6, 1)