[v3-3-test] Fix scheduler firing on_failure_callback for heartbeat-timed-out retries (#66767)

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -3548,12 +3548,21 @@ def _purge_task_instances_without_heartbeats(
             # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID).
             if not _ensure_ti_has_dag_version_id(ti, session, self.log):
                 continue
+            # ti.task isn't loaded in this purge path, so is_eligible_to_retry() uses its
+            # no-task fallback (``try_number <= max_tries``), which skips the retries-configured
+            # check its task-loaded branch applies; guard with ``max_tries > 0`` so a task
+            # declared with retries=0 isn't treated as retry-eligible here.
+            if ti.max_tries > 0 and ti.is_eligible_to_retry():
+                task_callback_type = TaskInstanceState.UP_FOR_RETRY
+            else:
+                task_callback_type = TaskInstanceState.FAILED
             request = TaskCallbackRequest(
                 filepath=ti.dag_model.relative_fileloc or "",
                 bundle_name=_hb_bundle_name,
                 bundle_version=_hb_bundle_version,
                 ti=ti,
                 msg=str(task_instance_heartbeat_timeout_message_details),
+                task_callback_type=task_callback_type,
                 context_from_server=TIRunContext(
                     dag_run=DRDataModel.model_validate(ti.dag_run, from_attributes=True),
                     max_tries=ti.max_tries,

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -8608,6 +8608,121 @@ def test_scheduler_passes_context_from_server_on_heartbeat_timeout(self, dag_mak
         assert callback_request.context_from_server.dag_run.logical_date == dag_run.logical_date
         assert callback_request.context_from_server.max_tries == ti.max_tries
 
+    @pytest.mark.parametrize(
+        ("state", "retries", "try_number", "expected_callback_type", "expected_dispatched_callback"),
+        [
+            pytest.param(
+                TaskInstanceState.RUNNING,
+                0,
+                1,
+                TaskInstanceState.FAILED,
+                "on_failure_callback",
+                id="no_retries",
+            ),
+            pytest.param(
+                TaskInstanceState.RUNNING,
+                2,
+                1,
+                TaskInstanceState.UP_FOR_RETRY,
+                "on_retry_callback",
+                id="retries_available_first_attempt",
+            ),
+            pytest.param(
+                TaskInstanceState.RUNNING,
+                2,
+                2,
+                TaskInstanceState.UP_FOR_RETRY,
+                "on_retry_callback",
+                id="retries_available_mid_chain",
+            ),
+            pytest.param(
+                TaskInstanceState.RUNNING,
+                2,
+                3,
+                TaskInstanceState.FAILED,
+                "on_failure_callback",
+                id="retries_exhausted",
+            ),
+            pytest.param(
+                TaskInstanceState.RESTARTING,
+                1,
+                5,
+                TaskInstanceState.UP_FOR_RETRY,
+                "on_retry_callback",
+                id="restarting_stays_eligible_past_max_tries",
+            ),
+        ],
+    )
+    def test_heartbeat_timeout_sets_callback_type_by_retry_eligibility(
+        self,
+        dag_maker,
+        session,
+        state,
+        retries,
+        try_number,
+        expected_callback_type,
+        expected_dispatched_callback,
+    ):
+        """Heartbeat-timeout cleanup must populate ``task_callback_type`` so the Dag processor
+        fires ``on_retry_callback`` when the task still has retries left, not
+        ``on_failure_callback``.
+
+        Reproduces the bug end-to-end through the actual scheduler purge path:
+
+        1. A TI is ``RUNNING`` (or ``RESTARTING``) with a stale ``last_heartbeat_at`` (worker
+           OOMKilled, node evicted, scheduler restarted, etc.).
+        2. ``_find_and_purge_task_instances_without_heartbeats`` builds a
+           ``TaskCallbackRequest`` and hands it to the executor's ``send_callback``.
+        3. The Dag processor branches on ``request.task_callback_type``:
+           ``UP_FOR_RETRY`` -> ``task.on_retry_callback``; anything else (including ``None``)
+           -> ``task.on_failure_callback``. See
+           ``airflow-core/src/airflow/dag_processing/processor.py``::``_execute_task_callbacks``.
+
+        Before the fix, step 2 left ``task_callback_type`` as ``None``, so step 3 always fell
+        into the ``else`` branch and ``on_failure_callback`` fired even when the task still had
+        retries left -- producing spurious failure alerts for tasks that ultimately succeeded on
+        retry.
+
+        The parametrized cases cover the full ``max_tries`` / ``try_number`` matrix for a
+        ``RUNNING`` TI -- no retries, retries available (first attempt and mid-chain), and
+        retries exhausted (``try_number > max_tries``) -- plus a ``RESTARTING`` TI (cleared
+        while running), which ``is_eligible_to_retry`` keeps retry-eligible even past
+        ``max_tries``. The ``expected_dispatched_callback`` column mirrors the Dag processor's
+        branch so the assertion captures the user-visible outcome, not just the field value.
+        """
+        with dag_maker(dag_id=f"hb_timeout_r{retries}_t{try_number}", session=session):
+            EmptyOperator(task_id="test_task", retries=retries)
+
+        dag_run = dag_maker.create_dagrun(run_id="test_run", state=DagRunState.RUNNING)
+
+        mock_executor = MagicMock()
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor])
+
+        ti = dag_run.get_task_instance(task_id="test_task")
+        ti.state = state
+        ti.try_number = try_number
+        ti.queued_by_job_id = scheduler_job.id
+        ti.last_heartbeat_at = timezone.utcnow() - timedelta(seconds=600)
+        session.merge(ti)
+        session.commit()
+
+        self.job_runner._find_and_purge_task_instances_without_heartbeats()
+
+        mock_executor.send_callback.assert_called_once()
+        request = mock_executor.send_callback.call_args[0][0]
+        assert isinstance(request, TaskCallbackRequest)
+        assert request.task_callback_type == expected_callback_type
+        # Mirror processor._execute_task_callbacks: UP_FOR_RETRY -> on_retry_callback, else
+        # on_failure_callback. Asserting the dispatched callback closes the loop on the
+        # user-visible behaviour, not just the field value.
+        dispatched_callback = (
+            "on_retry_callback"
+            if request.task_callback_type is TaskInstanceState.UP_FOR_RETRY
+            else "on_failure_callback"
+        )
+        assert dispatched_callback == expected_dispatched_callback
+
     @pytest.mark.parametrize(
         ("retries", "callback_kind", "expected"),
         [