Honor AirflowFailException raised inside on_retry_callback

[1fanwang]

Closes #60172.

_run_task_state_change_callbacks in task-sdk/src/airflow/sdk/execution_time/task_runner.py catches every exception from a callback and logs it. That's the right default for noisy cleanup work, but it also swallows the explicit AirflowFailException signal — a user raising it inside on_retry_callback to say "fail without retrying" had no way to actually fail the task. The state stayed UP_FOR_RETRY and another attempt was scheduled.

The fix narrows the catch in the retry-callback path:

• AirflowFailException re-raises so the caller can react.
• Every other Exception is still logged + swallowed — Wei's concern on the prior attempt (PR Fix AirflowFailException in on_retry_callback not preventing retries #64198) about cleanup callbacks that may fail is preserved.

run() now defers the supervisor RetryTask message until after the retry callback runs. finalize() gained an optional msg parameter: when the retry callback raises AirflowFailException, finalize promotes the state to FAILED, replaces the pending RetryTask with TaskState(FAILED), and runs the failure-path finalizers (on_failure_callback, listener hook, email_on_failure).

Two prior attempts (#60415 closed 2026-03-05, #64198 closed 2026-05-05) tried to solve this; this PR picks up the Wei-shaped design and ships it with a regression test.

Tests

task-sdk/tests/task_sdk/execution_time/test_task_runner.py::TestTaskRunnerCallsCallbacks::test_airflow_fail_exception_in_on_retry_callback_fails_task exercises the full path: retry callback raises AirflowFailException, the failure callback runs, the supervisor receives TaskState(FAILED) instead of RetryTask. The test fails on main and passes with this PR; existing callback tests in the file behave unchanged (generic Exception from a retry callback still swallowed + retry still scheduled).

Reproducer

Reverting just the production changes (task-sdk/src/airflow/sdk/execution_time/task_runner.py and supervisor.py) back to upstream/main while keeping the new test surfaces the bug directly. With the call site adapted to the older finalize() signature (no msg kwarg), runtime_ti.state lands on UP_FOR_RETRY instead of FAILED:

task-sdk/tests/task_sdk/execution_time/test_task_runner.py::TestTaskRunnerCallsCallbacks::test_airflow_fail_exception_in_on_retry_callback_fails_task FAILED [100%]

=================================== FAILURES ===================================
_ TestTaskRunnerCallsCallbacks.test_airflow_fail_exception_in_on_retry_callback_fails_task _
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:4542: in test_airflow_fail_exception_in_on_retry_callback_fails_task
    assert runtime_ti.state == TaskInstanceState.FAILED
E   AssertionError: assert <TaskInstanceState.UP_FOR_RETRY: 'up_for_retry'> == <TaskInstanceState.FAILED: 'failed'>
E
E     - failed
E     + up_for_retry
========================= 1 failed, 1 warning in 9.30s =========================

Restoring the production code (the full PR diff) flips the same test to PASSED:

task-sdk/tests/task_sdk/execution_time/test_task_runner.py::TestTaskRunnerCallsCallbacks::test_airflow_fail_exception_in_on_retry_callback_fails_task PASSED [100%]

========================= 1 passed, 1 warning in 1.72s =========================

The UP_FOR_RETRY line is the literal symptom from #60172 — the user raised AirflowFailException inside on_retry_callback and the task still ended up scheduled for another attempt. With the fix, the same callback raise promotes the task to FAILED and the failure-path finalizers fire under that state.

[SameerMesiah97]

Approach looks fine to me but someone more familiar with this area should weigh in. I have left some comments.

[SameerMesiah97]

nit: this could be clearer:

# Delay reporting UP_FOR_RETRY to the supervisor until after
# on_retry_callback runs so AirflowFailException can promote
# the task to FAILED and suppress the retry.

I don't think the issue reference is needed. But this is a more subjective point.

[1fanwang]

Done in 193d2d6 — applied the shorter wording and dropped the issue reference.

[SameerMesiah97]

Right now, this is a bit hard to follow. I would recommend extracting the listener and email notification into a helper like this:

def _handle_failure_notifications(
    *,
    task,
    ti,
    context,
    error,
    log,
    send_email: bool,
) -> None:
    try:
        get_listener_manager().hook.on_task_instance_failed(
            previous_state=TaskInstanceState.RUNNING,
            task_instance=ti,
            error=error,
        )
    except Exception:
        log.exception("error calling listener")

    if send_email and task.email:
        _send_error_email_notification(task, ti, context, error, log)

Then plug _handle_failure_notifications in so that lines 1960-1976 are replaced by this:

        _handle_failure_notifications(
            task=task,
            ti=ti,
            context=context,
            error=error,
            log=log,
            send_email=task.email_on_failure,
        )

    else:
        _handle_failure_notifications(
            task=task,
            ti=ti,
            context=context,
            error=error,
            log=log,
            send_email=bool(error and task.email_on_retry),
        )

[1fanwang]

Done in 193d2d6 — added _handle_failure_notifications and applied it across all three failure-path branches in finalize (AirflowFailException, retry, and FAILED) so the listener+email pattern only lives in one place.

[SameerMesiah97]

Remove this.

[1fanwang]

Done in f59182e — dropped the docstring; the test name already describes the behavior.

[SameerMesiah97]

Maybe it would be better to assert the state transition here instead of the number of calls:

assert retry_callback_calls == [TaskInstanceState.UP_FOR_RETRY]
assert failure_callback_calls == [TaskInstanceState.FAILED]

[1fanwang]

Done in f59182e — retry_callback_calls == [UP_FOR_RETRY] and failure_callback_calls == [FAILED] pin both the count and the state at callback time.