Fix task callbacks being skipped when TriggerDagRunOperator gets a 404

[kaxil]

Fixes #70683

Problem

When TriggerDagRunOperator targets a Dag that does not exist, the task dies without running on_failure_callback / on_retry_callback, without firing task-instance listeners, and without evaluating the task's retry_policy. To the user the task just goes red, and none of their failure handling runs.

The cause is not specific to TriggerDagRunOperator. run() maps the task's outcome through a flat chain of except clauses, and several of those clauses do real work. An exception raised inside an except clause is not offered to that clause's siblings, so it escapes run() entirely, skipping the retry decision in _handle_current_task_failed() and everything in finalize().

For the reported case the chain is:

1. The execution API returns 404 when no non-stale Dag matches (dag_runs.py#L107).
2. DagRunOperations.trigger only special-cases 409, so the 404 re-raises (client.py#L927).
3. The supervisor converts it to an API_SERVER_ERROR response (supervisor.py#L945).
4. CommsDecoder._from_frame raises AirflowRuntimeError (comms.py#L359).

That raise happens inside except DagRunTriggerException, 39 lines above the clause that already handles AirflowRuntimeError. Python never consults it.

It is not limited to the API server. _defer_task and _await_input_task run serde_serialize over user-supplied kwargs, and serde raises for any value it has no serializer for:

TypeError: cannot serialize object of type <class 'object'>

So self.defer(..., kwargs={"client": some_hook_client}) loses its callbacks the same way, with no 404 anywhere in sight. A file handle, a generator or a lambda all behave identically.

Solution

Split the function where deciding the outcome ends and reporting it begins:

• _run_task_and_map_outcome() holds the existing chain verbatim and returns (state, msg, error).
• run() calls it, and its except now covers the handlers as well as the task body, flowing into the same terminal-state finally as before.

The handler chain and the finally block are not edited, and nothing is reindented: git diff and git diff -w are identical, so there is no whitespace churn to read past.

Catching the handlers is deliberately structural rather than a list of the four that can raise today. This chain has grown handlers over time (HITL, DagRunTrigger), which is how the defect arose in the first place; a fifth risky handler should not have to remember to opt in.

Preserving the chain's own semantics

A catch-all boundary must not flatten the distinctions the chain already makes, so _handle_handler_failure():

• Keeps non-retryable exceptions non-retryable. AirflowFailException, AirflowSensorTimeout and AirflowTaskTerminated mean "do not retry" wherever they are raised. Routing them through the retry-count check would hand back UP_FOR_RETRY.
• Catches Exception, not BaseException. A KeyboardInterrupt still reaches main()'s exit-code-2 path. The supervisor's default termination signal is SIGINT, escalating SIGINT → SIGTERM → SIGKILL, so swallowing it would turn an operator-initiated kill into an ordinary retry.
• Fails closed if the failure path itself fails. This is reachable from user code today: RetryDecision is an unvalidated dataclass, so a retry_policy returning retry_delay=30 rather than a timedelta makes _finalize_task_failure raise AttributeError on .total_seconds(), outside _evaluate_retry_policy's own error handling. Re-entering that path with its own exception would fail the same way and escape, losing the callbacks this change exists to preserve, so it falls back to a plain FAILED state. The reported error is the exception that actually ended the run, so a callback sees the broken policy; the task's own failure stays reachable on the implicit exception chain.

That path also runs _finalize_task_failure twice, so the failure counters are now emitted once per task run rather than once per pass. Without that, a single failure would report ti_failures=2.

What changes for users

On main today the supervisor's exit-code fallback already recovers UP_FOR_RETRY from the non-zero exit, so the task does still retry. What was being lost, and now works:

• on_failure_callback and on_retry_callback run
• on_task_instance_failed listeners fire
• email_on_failure / email_on_retry are sent
• A retry_policy is evaluated, so it can force FAIL or supply a custom delay and reason
• The RetryTask message carries end_date, retry delay and retry reason instead of being skipped
• ti_failures, operator_failures and ti.finish stats are recorded
• dag.test() and run_task_in_process no longer crash out of the CLI, since InProcessTestSupervisor.start calls run() unguarded

The issue reports retries dropped entirely, which was accurate for 3.1.0: the _should_retry exit-code fallback landed later in #55767 and first shipped in 3.1.2.

Gotchas

If the first trigger attempt partially succeeds (the Dag run is created, then a later call in the same handler fails), the task is now retried where it previously died. On retry TriggerDagRunOperator mints a run id from run_after or utcnow() when no explicit trigger_run_id is set, so that retry can create a second Dag run. This behaviour is unchanged by this PR: the supervisor's exit-code fallback already produced the same retry on main. Pinning trigger_run_id makes the retry idempotent via the existing 409 handling.

[Copilot]

Pull request overview

This PR fixes a Task SDK task-runner control-flow bug where exceptions raised inside run()’s exception handlers could escape the function entirely, skipping retry decisions and all failure/finalization callbacks (notably affecting TriggerDagRunOperator when the Execution API returns 404, and deferral paths when kwargs are not serializable).

Changes:

• Refactors the task execution flow by extracting the existing exception-mapping chain into _run_task_and_map_outcome() and wrapping it in run() so handler-thrown exceptions still flow through the same terminal-state reporting and finalize() path.
• Adds _handle_handler_failure() / _terminal_failure() to preserve “do not retry” semantics for non-retryable exceptions and to “fail closed” if the failure path itself errors.
• Prevents double-counting failure metrics when the failure path runs twice due to an internal error.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
task-sdk/src/airflow/sdk/execution_time/task_runner.py	Restructures run() to catch handler failures, adds handler-failure outcome mapping, and guards failure metrics from double emission.
task-sdk/tests/task_sdk/execution_time/test_task_runner.py	Adds regression tests for retries/callbacks on handler failures (missing Dag 404 path and unserializable deferral kwargs), plus coverage for edge cases (non-retryable, KeyboardInterrupt, failure-path failure).

Comments suppressed due to low confidence (1)

task-sdk/tests/task_sdk/execution_time/test_task_runner.py:5419

• Test docstring includes a GitHub issue URL/number. Airflow guidelines recommend not embedding issue numbers in test docstrings; keep the regression explanation but remove the issue reference.

        Regression test for https://github.com/apache/airflow/issues/70683 -- the
        ``AirflowRuntimeError`` was raised from inside ``run()``'s
        ``except DagRunTriggerException`` handler, so it escaped ``run()`` without
        evaluating retries or running ``on_failure_callback`` / ``on_retry_callback``.
        """

[github-actions]

Backport successfully created: v3-3-test

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
✅	v3-3-test