diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 31d0831bbadec..6da2531febdad 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1623,7 +1623,16 @@ def _on_term(signum, frame): msg, state = _handle_current_task_success(context, ti) except DagRunTriggerException as drte: log.info("::group::Post Execute") - msg, state = _handle_trigger_dag_run(drte, context, ti, log) + try: + msg, state = _handle_trigger_dag_run(drte, context, ti, log) + except (AirflowTaskTimeout, AirflowException, AirflowRuntimeError) as e: + # Errors returned by the API server (for example a 404 for a missing + # target DAG) are raised while handling the trigger exception. They + # must still use the normal failure path so retries and callbacks run. + log.exception("Task failed while triggering DagRun") + log.info("::group::Post Execute") + msg, state = _handle_current_task_failed(ti, e, log, context) + error = e except TaskDeferred as defer: log.info("::group::Post Execute") msg, state = _defer_task(defer, ti, log) diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 0031974b86646..da8db6a22bffd 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -5180,6 +5180,37 @@ def _send(msg=None, **kwargs): with pytest.raises(_TriggerSendError): run(ti, ti.get_template_context(), log) + @time_machine.travel("2025-01-01 00:00:00", tick=False) + def test_handle_trigger_dag_run_api_error_uses_failure_path( + self, create_runtime_ti, mock_supervisor_comms + ): + """API errors while triggering a DAG should honor the task retry policy.""" + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="missing_dag", + trigger_run_id="test_run_id", + retries=1, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_api_error", + run_id="test_run", + task=task, + should_retry=True, + ) + + def _send(msg=None, **kwargs): + if isinstance(msg, TriggerDagRun): + raise AirflowRuntimeError("target DAG was not found") + return mock.DEFAULT + + mock_supervisor_comms.send.side_effect = _send + + state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock()) + + assert state == TaskInstanceState.UP_FOR_RETRY + assert msg.state == TaskInstanceState.UP_FOR_RETRY + assert isinstance(error, AirflowRuntimeError) + @pytest.mark.parametrize( ("allowed_states", "failed_states", "target_dr_state", "expected_task_state"), [