From 1d8201e3b9a5d651831f52c18fed7093a779a1a0 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Wed, 29 May 2024 18:09:58 +0800
Subject: [PATCH 01/18] feat(dagrun): add start_from_trigger support to mapped
 operator

---
 airflow/models/dagrun.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index d4ef937e9d797..2312a133bff33 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -57,6 +57,7 @@
 from airflow.models.abstractoperator import NotMapped
 from airflow.models.base import Base, StringID
 from airflow.models.expandinput import NotFullyPopulated
+from airflow.models.mappedoperator import MappedOperator
 from airflow.models.taskinstance import TaskInstance as TI
 from airflow.models.tasklog import LogTemplate
 from airflow.stats import Stats
@@ -1577,6 +1578,18 @@ def schedule_tis(
                 and not ti.task.outlets
             ):
                 dummy_ti_ids.append((ti.task_id, ti.map_index))
+            elif isinstance(ti.task, MappedOperator) and ti.task.start_trigger_args is not None:
+                context = ti.get_template_context()
+                mapped_kwargs, _ = ti.task._expand_mapped_kwargs(context, session)
+                start_from_trigger = mapped_kwargs.get("start_from_trigger", ti.task.start_trigger_args)
+                if start_from_trigger is True and "trigger_kwargs" in mapped_kwargs:
+                    start_trigger_args = ti.task.start_trigger_args
+                    start_trigger_args.trigger_kwargs = mapped_kwargs.get("trigger_kwargs")
+                    if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
+                        ti.try_number += 1
+                    ti.defer_task_from_start_trigger(session=session, start_trigger_args=start_trigger_args)
+                else:
+                    schedulable_ti_ids.append((ti.task_id, ti.map_index))
             elif ti.task.start_from_trigger is True and ti.task.start_trigger_args is not None:
                 ti.start_date = timezone.utcnow()
                 if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:

From 88a41b66f6c16e8830ca5cab7de1bb4b2a458d79 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Thu, 30 May 2024 15:59:24 +0800
Subject: [PATCH 02/18] refactor(dagrun): reduplicate task scheduling if-else
 blocks

---
 airflow/models/dagrun.py | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 2312a133bff33..0001dcf678cfd 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1578,23 +1578,28 @@ def schedule_tis(
                 and not ti.task.outlets
             ):
                 dummy_ti_ids.append((ti.task_id, ti.map_index))
-            elif isinstance(ti.task, MappedOperator) and ti.task.start_trigger_args is not None:
-                context = ti.get_template_context()
-                mapped_kwargs, _ = ti.task._expand_mapped_kwargs(context, session)
-                start_from_trigger = mapped_kwargs.get("start_from_trigger", ti.task.start_trigger_args)
-                if start_from_trigger is True and "trigger_kwargs" in mapped_kwargs:
-                    start_trigger_args = ti.task.start_trigger_args
-                    start_trigger_args.trigger_kwargs = mapped_kwargs.get("trigger_kwargs")
+            # check whether the operator supports start execution from triggerer
+            elif ti.task.start_trigger_args is not None:
+                start_trigger_args = ti.task.start_trigger_args
+                if isinstance(ti.task, MappedOperator):
+                    context = ti.get_template_context()
+                    mapped_kwargs, _ = ti.task._expand_mapped_kwargs(context, session)
+                    start_from_trigger = mapped_kwargs.get("start_from_trigger", ti.task.start_from_trigger)
+
+                    # update the trigger_kwargs if it's in expanded kwargs
+                    start_trigger_args.trigger_kwargs = mapped_kwargs.get(
+                        "trigger_kwargs", start_trigger_args.trigger_kwargs
+                    )
+                else:
+                    start_from_trigger = ti.task.start_from_trigger
+
+                if start_from_trigger is True:
+                    ti.start_date = timezone.utcnow()
                     if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
                         ti.try_number += 1
-                    ti.defer_task_from_start_trigger(session=session, start_trigger_args=start_trigger_args)
+                    ti.defer_task(exception=None, session=session)
                 else:
                     schedulable_ti_ids.append((ti.task_id, ti.map_index))
-            elif ti.task.start_from_trigger is True and ti.task.start_trigger_args is not None:
-                ti.start_date = timezone.utcnow()
-                if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
-                    ti.try_number += 1
-                ti.defer_task(exception=None, session=session)
             else:
                 schedulable_ti_ids.append((ti.task_id, ti.map_index))
 

From b3e29bac0d1dcd1d95b7a3e8af82f7d34a2e1d84 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 31 May 2024 17:57:04 +0800
Subject: [PATCH 03/18] docs(deferring): add dynamic task mapping support to
 "Triggering Deferral from Start"

---
 .../authoring-and-scheduling/deferring.rst    | 46 +++++++++++++++++--
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index a9b26703bce63..6cfda203d4080 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -182,9 +182,6 @@ This is particularly useful when deferring is the only thing the ``execute`` met
 
 ``start_from_trigger`` and ``trigger_kwargs`` can also be modified at the instance level for more flexible configuration.
 
-.. warning::
-    Dynamic task mapping is not supported when ``trigger_kwargs`` is modified at instance level.
-
 .. code-block:: python
 
     from datetime import timedelta
@@ -213,6 +210,49 @@ This is particularly useful when deferring is the only thing the ``execute`` met
             # We have no more work to do here. Mark as complete.
             return
 
+To enable Dynamic Task Mapping support, you can define ``start_from_trigger`` and ``trigger_kwargs`` in the parameter of "__init__". Note that you don't need to define both of them to use this feature, but you do need to use the exact same parameter name. For example, if you define an argument as ``t_kwargs`` and assign this value to ``self.start_trigger_args.trigger_kwargs``, it will not work.
+
+.. code-block:: python
+
+    from datetime import timedelta
+    from typing import Any
+
+    from airflow.sensors.base import BaseSensorOperator
+    from airflow.triggers.temporal import TimeDeltaTrigger
+    from airflow.utils.context import Context
+
+
+    class WaitTwoHourSensor(BaseSensorOperator):
+        start_trigger_args = StartTriggerArgs(
+            trigger_cls="airflow.triggers.temporal.TimeDeltaTrigger",
+            trigger_kwargs={},
+            next_method="execute_complete",
+            timeout=None,
+        )
+
+        def __init__(
+            self,
+            *args: list[Any],
+            trigger_kwargs: dict[str, Any] | None,
+            start_from_trigger: bool,
+            **kwargs: dict[str, Any],
+        ) -> None:
+            super().__init__(*args, **kwargs)
+            self.start_trigger_args.trigger_kwargs = trigger_kwargs
+            self.start_from_trigger = start_from_trigger
+
+        def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+            # We have no more work to do here. Mark as complete.
+            return
+
+These parameters can be mapped using the ``expand`` method.
+
+.. code-block:: python
+
+    WaitTwoHourSensor.partial(task_id="transform").expand(
+        trigger_kwargs=[{"moment": timedelta(hours=1)}], start_from_trigger=True
+    )
+
 Writing Triggers
 ~~~~~~~~~~~~~~~~
 

From 7ada27dc6bbfe7f8212f27ae9192114d444a974f Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 31 May 2024 18:25:00 +0800
Subject: [PATCH 04/18] test(dagrun): add test to start_trigger with mapped op

---
 tests/models/test_dagrun.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/tests/models/test_dagrun.py b/tests/models/test_dagrun.py
index 021d09e889420..7a791aa365a82 100644
--- a/tests/models/test_dagrun.py
+++ b/tests/models/test_dagrun.py
@@ -2043,6 +2043,38 @@ def test_schedule_tis_empty_operator_try_number(dag_maker, session: Session):
     assert empty_ti.try_number == 1
 
 
+def test_schedule_tis_start_trigger_through_expand(dag_maker, session):
+    """
+    Test that an operator with start_trigger_args set can be directly deferred during scheduling.
+    """
+
+    class TestOperator(BaseOperator):
+        start_trigger_args = StartTriggerArgs(
+            trigger_cls="airflow.triggers.testing.SuccessTrigger",
+            trigger_kwargs={},
+            next_method="execute_complete",
+            timeout=None,
+        )
+        start_from_trigger = False
+
+        def __init__(self, *args, start_from_trigger: bool = False, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.start_from_trigger = start_from_trigger
+
+        def execute_complete(self):
+            pass
+
+    with dag_maker(session=session):
+        TestOperator.partial(task_id="test_task").expand(start_from_trigger=[True, False])
+
+    dr: DagRun = dag_maker.create_dagrun()
+
+    dr.schedule_tis(dr.task_instances, session=session)
+    tis = [(ti.state, ti.map_index) for ti in dr.task_instances]
+    assert tis[0] == (TaskInstanceState.DEFERRED, 0)
+    assert tis[1] == (None, 1)
+
+
 def test_mapped_expand_kwargs(dag_maker):
     with dag_maker():
 

From 5304bcf3d39b0ab1efcae03a407ab9b6c1e8dca7 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 21 Jun 2024 17:50:04 +0800
Subject: [PATCH 05/18] refactor(mappedoperator): extract start trigger args
 expansion to mappedoperator

---
 airflow/models/dagrun.py         | 16 ++++------------
 airflow/models/mappedoperator.py | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 0001dcf678cfd..d0966e15e5f22 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -57,7 +57,6 @@
 from airflow.models.abstractoperator import NotMapped
 from airflow.models.base import Base, StringID
 from airflow.models.expandinput import NotFullyPopulated
-from airflow.models.mappedoperator import MappedOperator
 from airflow.models.taskinstance import TaskInstance as TI
 from airflow.models.tasklog import LogTemplate
 from airflow.stats import Stats
@@ -1580,20 +1579,13 @@ def schedule_tis(
                 dummy_ti_ids.append((ti.task_id, ti.map_index))
             # check whether the operator supports start execution from triggerer
             elif ti.task.start_trigger_args is not None:
-                start_trigger_args = ti.task.start_trigger_args
+                from airflow.models.mappedoperator import MappedOperator
+
                 if isinstance(ti.task, MappedOperator):
                     context = ti.get_template_context()
-                    mapped_kwargs, _ = ti.task._expand_mapped_kwargs(context, session)
-                    start_from_trigger = mapped_kwargs.get("start_from_trigger", ti.task.start_from_trigger)
-
-                    # update the trigger_kwargs if it's in expanded kwargs
-                    start_trigger_args.trigger_kwargs = mapped_kwargs.get(
-                        "trigger_kwargs", start_trigger_args.trigger_kwargs
-                    )
-                else:
-                    start_from_trigger = ti.task.start_from_trigger
+                    ti.task._expand_start_trigger(context=context, session=session)
 
-                if start_from_trigger is True:
+                if ti.task.start_from_trigger is True:
                     ti.start_date = timezone.utcnow()
                     if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
                         ti.try_number += 1
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index 45a6ae1ac176b..9755e2e42e68e 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -729,6 +729,24 @@ def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -
             "params": params,
         }
 
+    def _expand_start_trigger(self, *, context: Context, session: Session) -> None:
+        """Get the kwargs to create the unmapped start_from_trigger and start_trigger_args.
+
+        This method is for allowing mapped operator to start execution from triggerer.
+        """
+        if not self.start_trigger_args:
+            return
+
+        mapped_kwargs, _ = self._expand_mapped_kwargs(context, session)
+        self.start_from_trigger = mapped_kwargs.get("start_from_trigger", self.start_from_trigger)
+
+        self.start_trigger_args.trigger_kwargs = mapped_kwargs.get(
+            "trigger_kwargs", self.start_trigger_args.trigger_kwargs
+        )
+        self.start_trigger_args.timeout = mapped_kwargs.get(
+            "trigger_timeout", self.start_trigger_args.timeout
+        )
+
     def unmap(self, resolve: None | Mapping[str, Any] | tuple[Context, Session]) -> BaseOperator:
         """
         Get the "normal" Operator after applying the current mapping.

From 61cfb897e0feb13e1815e07157496978033be5a9 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Sat, 22 Jun 2024 11:43:48 +0800
Subject: [PATCH 06/18] feat(mapped_operator): add partial support to
 start_trigger_args

---
 airflow/models/mappedoperator.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index 9755e2e42e68e..d48dc6c377c51 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -738,13 +738,22 @@ def _expand_start_trigger(self, *, context: Context, session: Session) -> None:
             return
 
         mapped_kwargs, _ = self._expand_mapped_kwargs(context, session)
-        self.start_from_trigger = mapped_kwargs.get("start_from_trigger", self.start_from_trigger)
+        if self._disallow_kwargs_override:
+            prevent_duplicates(
+                self.partial_kwargs,
+                mapped_kwargs,
+                fail_reason="unmappable or already specified",
+            )
+
+        self.start_from_trigger = self.partial_kwargs.get(
+            "start_from_trigger", mapped_kwargs.get("start_from_trigger", self.start_from_trigger)
+        )
 
-        self.start_trigger_args.trigger_kwargs = mapped_kwargs.get(
-            "trigger_kwargs", self.start_trigger_args.trigger_kwargs
+        self.start_trigger_args.trigger_kwargs = self.partial_kwargs.get(
+            "trigger_kwargs", mapped_kwargs.get("trigger_kwargs", self.start_trigger_args.trigger_kwargs)
         )
-        self.start_trigger_args.timeout = mapped_kwargs.get(
-            "trigger_timeout", self.start_trigger_args.timeout
+        self.start_trigger_args.timeout = self.partial_kwargs.get(
+            "trigger_timeout", mapped_kwargs.get("trigger_timeout", self.start_trigger_args.timeout)
         )
 
     def unmap(self, resolve: None | Mapping[str, Any] | tuple[Context, Session]) -> BaseOperator:

From 7a94a9c7c27e871eec8a5fbcb7a0354056cf4403 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Mon, 24 Jun 2024 09:15:11 +0800
Subject: [PATCH 07/18] fix(mappedoperator): make mapped kwargs override
 partial kwargs when expand start trigger

---
 airflow/models/mappedoperator.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index d48dc6c377c51..bb1dd7a6e5569 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -745,15 +745,18 @@ def _expand_start_trigger(self, *, context: Context, session: Session) -> None:
                 fail_reason="unmappable or already specified",
             )
 
-        self.start_from_trigger = self.partial_kwargs.get(
-            "start_from_trigger", mapped_kwargs.get("start_from_trigger", self.start_from_trigger)
+        # Ordering is significant; mapped kwargs should override partial ones.
+
+        self.start_from_trigger = mapped_kwargs.get(
+            "start_from_trigger", self.partial_kwargs.get("start_from_trigger", self.start_from_trigger)
         )
 
-        self.start_trigger_args.trigger_kwargs = self.partial_kwargs.get(
-            "trigger_kwargs", mapped_kwargs.get("trigger_kwargs", self.start_trigger_args.trigger_kwargs)
+        self.start_trigger_args.trigger_kwargs = mapped_kwargs.get(
+            "trigger_kwargs",
+            self.partial_kwargs.get("trigger_kwargs", self.start_trigger_args.trigger_kwargs),
         )
-        self.start_trigger_args.timeout = self.partial_kwargs.get(
-            "trigger_timeout", mapped_kwargs.get("trigger_timeout", self.start_trigger_args.timeout)
+        self.start_trigger_args.timeout = mapped_kwargs.get(
+            "trigger_timeout", self.partial_kwargs.get("trigger_timeout", self.start_trigger_args.timeout)
         )
 
     def unmap(self, resolve: None | Mapping[str, Any] | tuple[Context, Session]) -> BaseOperator:

From f9c57d6d7181a8b324e257b282ecd43bee8a75ca Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Mon, 24 Jun 2024 09:20:15 +0800
Subject: [PATCH 08/18] docs(deferring): add partial example to start_trigger

---
 docs/apache-airflow/authoring-and-scheduling/deferring.rst | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index 6cfda203d4080..1a4985fbac875 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -245,12 +245,12 @@ To enable Dynamic Task Mapping support, you can define ``start_from_trigger`` an
             # We have no more work to do here. Mark as complete.
             return
 
-These parameters can be mapped using the ``expand`` method.
+These parameters can be mapped using the ``expand`` and ``partial`` methods.
 
 .. code-block:: python
 
-    WaitTwoHourSensor.partial(task_id="transform").expand(
-        trigger_kwargs=[{"moment": timedelta(hours=1)}], start_from_trigger=True
+    WaitTwoHourSensor.partial(task_id="transform").partial(start_from_trigger=True).expand(
+        trigger_kwargs=[{"moment": timedelta(hours=2)}, {"moment": timedelta(hours=2)}]
     )
 
 Writing Triggers

From 1a22d0bade4ad2fe7a014a9982e2cb92565e285a Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Mon, 24 Jun 2024 09:23:25 +0800
Subject: [PATCH 09/18] style(dagrun): change is True checking style

---
 airflow/models/dagrun.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index d0966e15e5f22..d4cea6fc01290 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1585,7 +1585,7 @@ def schedule_tis(
                     context = ti.get_template_context()
                     ti.task._expand_start_trigger(context=context, session=session)
 
-                if ti.task.start_from_trigger is True:
+                if ti.task.start_from_trigger:
                     ti.start_date = timezone.utcnow()
                     if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
                         ti.try_number += 1

From 2c3c64ad5530576c1d7689f116467a4749ed1343 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Mon, 24 Jun 2024 16:42:42 +0800
Subject: [PATCH 10/18] refactor(mappedoperator): split _expand_start_trigger
 into _expand_start_from_trigger and _expand_start_trigger_args

---
 airflow/models/dagrun.py         |  6 +++--
 airflow/models/mappedoperator.py | 45 ++++++++++++++++++++++++++------
 airflow/models/taskinstance.py   | 15 ++++++++---
 3 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index d4cea6fc01290..2c6451bab552d 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1583,9 +1583,11 @@ def schedule_tis(
 
                 if isinstance(ti.task, MappedOperator):
                     context = ti.get_template_context()
-                    ti.task._expand_start_trigger(context=context, session=session)
+                    start_from_trigger = ti.task._expand_start_from_trigger(context=context, session=session)
+                else:
+                    start_from_trigger = ti.task.start_from_trigger
 
-                if ti.task.start_from_trigger:
+                if start_from_trigger:
                     ti.start_date = timezone.utcnow()
                     if ti.state != TaskInstanceState.UP_FOR_RESCHEDULE:
                         ti.try_number += 1
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index bb1dd7a6e5569..af74dad5f6c80 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -51,6 +51,7 @@
 from airflow.serialization.enums import DagAttributeTypes
 from airflow.task.priority_strategy import PriorityWeightStrategy, validate_and_load_priority_weight_strategy
 from airflow.ti_deps.deps.mapped_task_expanded import MappedTaskIsExpanded
+from airflow.triggers.base import StartTriggerArgs
 from airflow.typing_compat import Literal
 from airflow.utils.context import context_update_for_unmapped
 from airflow.utils.helpers import is_container, prevent_duplicates
@@ -81,7 +82,6 @@
     from airflow.models.param import ParamsDict
     from airflow.models.xcom_arg import XComArg
     from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
-    from airflow.triggers.base import StartTriggerArgs
     from airflow.utils.context import Context
     from airflow.utils.operator_resources import Resources
     from airflow.utils.task_group import TaskGroup
@@ -729,13 +729,13 @@ def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -
             "params": params,
         }
 
-    def _expand_start_trigger(self, *, context: Context, session: Session) -> None:
-        """Get the kwargs to create the unmapped start_from_trigger and start_trigger_args.
+    def _expand_start_from_trigger(self, *, context: Context, session: Session) -> bool:
+        """Get the kwargs to create the unmapped start_from_trigger.
 
         This method is for allowing mapped operator to start execution from triggerer.
         """
         if not self.start_trigger_args:
-            return
+            return False
 
         mapped_kwargs, _ = self._expand_mapped_kwargs(context, session)
         if self._disallow_kwargs_override:
@@ -746,18 +746,47 @@ def _expand_start_trigger(self, *, context: Context, session: Session) -> None:
             )
 
         # Ordering is significant; mapped kwargs should override partial ones.
-
-        self.start_from_trigger = mapped_kwargs.get(
+        return mapped_kwargs.get(
             "start_from_trigger", self.partial_kwargs.get("start_from_trigger", self.start_from_trigger)
         )
 
-        self.start_trigger_args.trigger_kwargs = mapped_kwargs.get(
+    def _expand_start_trigger_args(self, *, context: Context, session: Session) -> StartTriggerArgs:
+        """Get the kwargs to create the unmapped start_trigger_args.
+
+        This method is for allowing mapped operator to start execution from triggerer.
+        """
+        if not self.start_trigger_args:
+            raise AirflowException(
+                "Cannot expand start_trigger_args for mapped operator without class level start_trigger_args"
+            )
+
+        mapped_kwargs, _ = self._expand_mapped_kwargs(context, session)
+        if self._disallow_kwargs_override:
+            prevent_duplicates(
+                self.partial_kwargs,
+                mapped_kwargs,
+                fail_reason="unmappable or already specified",
+            )
+
+        # Ordering is significant; mapped kwargs should override partial ones.
+        trigger_kwargs = mapped_kwargs.get(
             "trigger_kwargs",
             self.partial_kwargs.get("trigger_kwargs", self.start_trigger_args.trigger_kwargs),
         )
-        self.start_trigger_args.timeout = mapped_kwargs.get(
+        next_kwargs = mapped_kwargs.get(
+            "next_kwargs",
+            self.partial_kwargs.get("next_kwargs", self.start_trigger_args.next_kwargs),
+        )
+        timeout = mapped_kwargs.get(
             "trigger_timeout", self.partial_kwargs.get("trigger_timeout", self.start_trigger_args.timeout)
         )
+        return StartTriggerArgs(
+            trigger_cls=self.start_trigger_args.trigger_cls,
+            trigger_kwargs=trigger_kwargs,
+            next_method=self.start_trigger_args.next_method,
+            next_kwargs=next_kwargs,
+            timeout=timeout,
+        )
 
     def unmap(self, resolve: None | Mapping[str, Any] | tuple[Context, Session]) -> BaseOperator:
         """
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 393da63c74a8d..a20175967c312 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1617,13 +1617,20 @@ def _defer_task(
         next_kwargs = exception.kwargs
         timeout = exception.timeout
     elif ti.task is not None and ti.task.start_trigger_args is not None:
+        if isinstance(ti.task, MappedOperator):
+            context = ti.get_template_context()
+            start_trigger_args = ti.task._expand_start_trigger_args(context=context, session=session)
+        else:
+            start_trigger_args = ti.task.start_trigger_args
+
+        trigger_kwargs = start_trigger_args.trigger_kwargs or {}
+        next_kwargs = start_trigger_args.next_kwargs
+        next_method = start_trigger_args.next_method
+        timeout = start_trigger_args.timeout
         trigger_row = Trigger(
             classpath=ti.task.start_trigger_args.trigger_cls,
-            kwargs=ti.task.start_trigger_args.trigger_kwargs or {},
+            kwargs=trigger_kwargs,
         )
-        next_kwargs = ti.task.start_trigger_args.next_kwargs
-        next_method = ti.task.start_trigger_args.next_method
-        timeout = ti.task.start_trigger_args.timeout
     else:
         raise AirflowException("exception and ti.task.start_trigger_args cannot both be None")
 

From a94bfe235f8328e841f711da78410630a939a03e Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Mon, 24 Jun 2024 17:06:34 +0800
Subject: [PATCH 11/18] docs(dagrun): improve the deferring section of the
 schedule_tis docstring

---
 airflow/models/dagrun.py                                   | 5 ++++-
 docs/apache-airflow/authoring-and-scheduling/deferring.rst | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 2c6451bab552d..cc22016dcd819 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1577,7 +1577,10 @@ def schedule_tis(
                 and not ti.task.outlets
             ):
                 dummy_ti_ids.append((ti.task_id, ti.map_index))
-            # check whether the operator supports start execution from triggerer
+            # check "start_trigger_args" to see whether the operator supports start execution from triggerer
+            # if so, we'll then check "start_from_trigger" to see whether this feature is turned on and defer
+            # this task.
+            # if not, we'll add this "ti" into "schedulable_ti_ids" and later execute it to run in the worker
             elif ti.task.start_trigger_args is not None:
                 from airflow.models.mappedoperator import MappedOperator
 
diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index 1a4985fbac875..e2a4f25f14ec9 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -249,7 +249,7 @@ These parameters can be mapped using the ``expand`` and ``partial`` methods.
 
 .. code-block:: python
 
-    WaitTwoHourSensor.partial(task_id="transform").partial(start_from_trigger=True).expand(
+    WaitTwoHourSensor.partial(task_id="transform", start_from_trigger=True).expand(
         trigger_kwargs=[{"moment": timedelta(hours=2)}, {"moment": timedelta(hours=2)}]
     )
 

From 340c6e369e7e5d23e2285f4b2dd702712e2d99f5 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 28 Jun 2024 14:34:25 +0800
Subject: [PATCH 12/18] feat(mappedoperator): do not include xcom when
 expanding start trigger args and flag

---
 airflow/decorators/base.py                    |  6 ++++--
 airflow/models/expandinput.py                 | 20 ++++++++++++-------
 airflow/models/mappedoperator.py              | 16 +++++++++------
 .../authoring-and-scheduling/deferring.rst    |  2 +-
 4 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index 5a20fa55dd284..284e4eb208a84 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -550,11 +550,13 @@ def __attrs_post_init__(self):
         super(DecoratedMappedOperator, DecoratedMappedOperator).__attrs_post_init__(self)
         XComArg.apply_upstream_relationship(self, self.op_kwargs_expand_input.value)
 
-    def _expand_mapped_kwargs(self, context: Context, session: Session) -> tuple[Mapping[str, Any], set[int]]:
+    def _expand_mapped_kwargs(
+        self, context: Context, session: Session, *, include_xcom: bool = True
+    ) -> tuple[Mapping[str, Any], set[int]]:
         # We only use op_kwargs_expand_input so this must always be empty.
         if self.expand_input is not EXPAND_INPUT_EMPTY:
             raise AssertionError(f"unexpected expand_input: {self.expand_input}")
-        op_kwargs, resolved_oids = super()._expand_mapped_kwargs(context, session)
+        op_kwargs, resolved_oids = super()._expand_mapped_kwargs(context, session, include_xcom=include_xcom)
         return {"op_kwargs": op_kwargs}, resolved_oids
 
     def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -> dict[str, Any]:
diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py
index 4673eb4960064..96a4658f0b981 100644
--- a/airflow/models/expandinput.py
+++ b/airflow/models/expandinput.py
@@ -69,8 +69,8 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
         yield from self._input.iter_references()
 
     @provide_session
-    def resolve(self, context: Context, *, session: Session = NEW_SESSION) -> Any:
-        data, _ = self._input.resolve(context, session=session)
+    def resolve(self, context: Context, *, session: Session = NEW_SESSION, include_xcom: bool = True) -> Any:
+        data, _ = self._input.resolve(context, session=session, include_xcom=include_xcom)
         return data[self._key]
 
 
@@ -165,8 +165,10 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int:
         lengths = self._get_map_lengths(run_id, session=session)
         return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1)
 
-    def _expand_mapped_field(self, key: str, value: Any, context: Context, *, session: Session) -> Any:
-        if _needs_run_time_resolution(value):
+    def _expand_mapped_field(
+        self, key: str, value: Any, context: Context, *, session: Session, include_xcom: bool = True
+    ) -> Any:
+        if include_xcom and _needs_run_time_resolution(value):
             value = value.resolve(context, session=session)
         map_index = context["ti"].map_index
         if map_index < 0:
@@ -203,7 +205,9 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
             if isinstance(x, XComArg):
                 yield from x.iter_references()
 
-    def resolve(self, context: Context, session: Session) -> tuple[Mapping[str, Any], set[int]]:
+    def resolve(
+        self, context: Context, session: Session, *, include_xcom: bool = True
+    ) -> tuple[Mapping[str, Any], set[int]]:
         data = {k: self._expand_mapped_field(k, v, context, session=session) for k, v in self.value.items()}
         literal_keys = {k for k, _ in self._iter_parse_time_resolved_kwargs()}
         resolved_oids = {id(v) for k, v in data.items() if k not in literal_keys}
@@ -248,7 +252,9 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
                 if isinstance(x, XComArg):
                     yield from x.iter_references()
 
-    def resolve(self, context: Context, session: Session) -> tuple[Mapping[str, Any], set[int]]:
+    def resolve(
+        self, context: Context, session: Session, *, include_xcom: bool = True
+    ) -> tuple[Mapping[str, Any], set[int]]:
         map_index = context["ti"].map_index
         if map_index < 0:
             raise RuntimeError("can't resolve task-mapping argument without expanding")
@@ -258,7 +264,7 @@ def resolve(self, context: Context, session: Session) -> tuple[Mapping[str, Any]
             mapping = self.value[map_index]
             if not isinstance(mapping, collections.abc.Mapping):
                 mapping = mapping.resolve(context, session)
-        else:
+        elif include_xcom:
             mappings = self.value.resolve(context, session)
             if not isinstance(mappings, collections.abc.Sequence):
                 raise ValueError(f"expand_kwargs() expects a list[dict], not {_describe_type(mappings)}")
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index af74dad5f6c80..edf95c2541ea2 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -688,14 +688,16 @@ def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
         """Implement DAGNode."""
         return DagAttributeTypes.OP, self.task_id
 
-    def _expand_mapped_kwargs(self, context: Context, session: Session) -> tuple[Mapping[str, Any], set[int]]:
+    def _expand_mapped_kwargs(
+        self, context: Context, session: Session, *, include_xcom: bool = True
+    ) -> tuple[Mapping[str, Any], set[int]]:
         """
         Get the kwargs to create the unmapped operator.
 
         This exists because taskflow operators expand against op_kwargs, not the
         entire operator kwargs dict.
         """
-        return self._get_specified_expand_input().resolve(context, session)
+        return self._get_specified_expand_input().resolve(context, session, include_xcom=include_xcom)
 
     def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -> dict[str, Any]:
         """
@@ -730,14 +732,15 @@ def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -
         }
 
     def _expand_start_from_trigger(self, *, context: Context, session: Session) -> bool:
-        """Get the kwargs to create the unmapped start_from_trigger.
+        """
+        Get the kwargs to create the unmapped start_from_trigger.
 
         This method is for allowing mapped operator to start execution from triggerer.
         """
         if not self.start_trigger_args:
             return False
 
-        mapped_kwargs, _ = self._expand_mapped_kwargs(context, session)
+        mapped_kwargs, _ = self._expand_mapped_kwargs(context, session, include_xcom=False)
         if self._disallow_kwargs_override:
             prevent_duplicates(
                 self.partial_kwargs,
@@ -751,7 +754,8 @@ def _expand_start_from_trigger(self, *, context: Context, session: Session) -> b
         )
 
     def _expand_start_trigger_args(self, *, context: Context, session: Session) -> StartTriggerArgs:
-        """Get the kwargs to create the unmapped start_trigger_args.
+        """
+        Get the kwargs to create the unmapped start_trigger_args.
 
         This method is for allowing mapped operator to start execution from triggerer.
         """
@@ -760,7 +764,7 @@ def _expand_start_trigger_args(self, *, context: Context, session: Session) -> S
                 "Cannot expand start_trigger_args for mapped operator without class level start_trigger_args"
             )
 
-        mapped_kwargs, _ = self._expand_mapped_kwargs(context, session)
+        mapped_kwargs, _ = self._expand_mapped_kwargs(context, session, include_xcom=False)
         if self._disallow_kwargs_override:
             prevent_duplicates(
                 self.partial_kwargs,
diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index e2a4f25f14ec9..c89d671523c2d 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -245,7 +245,7 @@ To enable Dynamic Task Mapping support, you can define ``start_from_trigger`` an
             # We have no more work to do here. Mark as complete.
             return
 
-These parameters can be mapped using the ``expand`` and ``partial`` methods.
+These parameters can be mapped using the ``expand`` and ``partial`` methods. Note that XCom values won't be resolved at this stage.
 
 .. code-block:: python
 

From 061b8e12f0f811008fa47c47f368d91cf2dd9e65 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 12 Jul 2024 15:40:38 +0800
Subject: [PATCH 13/18] refactor: make include_xcom arg in
 _expand_mapped_kwargs a required argument

---
 airflow/decorators/base.py       | 2 +-
 airflow/models/mappedoperator.py | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index 284e4eb208a84..d743acbe50b2b 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -551,7 +551,7 @@ def __attrs_post_init__(self):
         XComArg.apply_upstream_relationship(self, self.op_kwargs_expand_input.value)
 
     def _expand_mapped_kwargs(
-        self, context: Context, session: Session, *, include_xcom: bool = True
+        self, context: Context, session: Session, *, include_xcom: bool
     ) -> tuple[Mapping[str, Any], set[int]]:
         # We only use op_kwargs_expand_input so this must always be empty.
         if self.expand_input is not EXPAND_INPUT_EMPTY:
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index edf95c2541ea2..9744247c5bd50 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -689,7 +689,7 @@ def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
         return DagAttributeTypes.OP, self.task_id
 
     def _expand_mapped_kwargs(
-        self, context: Context, session: Session, *, include_xcom: bool = True
+        self, context: Context, session: Session, *, include_xcom: bool
     ) -> tuple[Mapping[str, Any], set[int]]:
         """
         Get the kwargs to create the unmapped operator.
@@ -812,7 +812,7 @@ def unmap(self, resolve: None | Mapping[str, Any] | tuple[Context, Session]) ->
             if isinstance(resolve, collections.abc.Mapping):
                 kwargs = resolve
             elif resolve is not None:
-                kwargs, _ = self._expand_mapped_kwargs(*resolve)
+                kwargs, _ = self._expand_mapped_kwargs(*resolve, include_xcom=True)
             else:
                 raise RuntimeError("cannot unmap a non-serialized operator without context")
             kwargs = self._get_unmap_kwargs(kwargs, strict=self._disallow_kwargs_override)
@@ -907,7 +907,7 @@ def render_template_fields(
         # set_current_task_session context manager to store the session in the current task.
         session = get_current_task_instance_session()
 
-        mapped_kwargs, seen_oids = self._expand_mapped_kwargs(context, session)
+        mapped_kwargs, seen_oids = self._expand_mapped_kwargs(context, session, include_xcom=True)
         unmapped_task = self.unmap(mapped_kwargs)
         context_update_for_unmapped(context, unmapped_task)
 

From 62bef82819840c7c8336d888ea4aad249823dd17 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 12 Jul 2024 15:47:57 +0800
Subject: [PATCH 14/18] refactor(expandinput): make include_com in resolve a
 required argument

---
 airflow/models/expandinput.py | 19 +++++++++++--------
 airflow/models/param.py       |  2 +-
 airflow/models/xcom_arg.py    | 16 ++++++++--------
 airflow/template/templater.py |  4 ++--
 airflow/utils/mixins.py       |  2 +-
 5 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py
index 96a4658f0b981..417c3bd0c1c9c 100644
--- a/airflow/models/expandinput.py
+++ b/airflow/models/expandinput.py
@@ -69,7 +69,7 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
         yield from self._input.iter_references()
 
     @provide_session
-    def resolve(self, context: Context, *, session: Session = NEW_SESSION, include_xcom: bool = True) -> Any:
+    def resolve(self, context: Context, *, include_xcom: bool, session: Session = NEW_SESSION) -> Any:
         data, _ = self._input.resolve(context, session=session, include_xcom=include_xcom)
         return data[self._key]
 
@@ -166,10 +166,10 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int:
         return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1)
 
     def _expand_mapped_field(
-        self, key: str, value: Any, context: Context, *, session: Session, include_xcom: bool = True
+        self, key: str, value: Any, context: Context, *, session: Session, include_xcom: bool
     ) -> Any:
         if include_xcom and _needs_run_time_resolution(value):
-            value = value.resolve(context, session=session)
+            value = value.resolve(context, session=session, include_xcom=include_xcom)
         map_index = context["ti"].map_index
         if map_index < 0:
             raise RuntimeError("can't resolve task-mapping argument without expanding")
@@ -206,9 +206,12 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
                 yield from x.iter_references()
 
     def resolve(
-        self, context: Context, session: Session, *, include_xcom: bool = True
+        self, context: Context, session: Session, *, include_xcom: bool
     ) -> tuple[Mapping[str, Any], set[int]]:
-        data = {k: self._expand_mapped_field(k, v, context, session=session) for k, v in self.value.items()}
+        data = {
+            k: self._expand_mapped_field(k, v, context, session=session, include_xcom=include_xcom)
+            for k, v in self.value.items()
+        }
         literal_keys = {k for k, _ in self._iter_parse_time_resolved_kwargs()}
         resolved_oids = {id(v) for k, v in data.items() if k not in literal_keys}
         return data, resolved_oids
@@ -253,7 +256,7 @@ def iter_references(self) -> Iterable[tuple[Operator, str]]:
                     yield from x.iter_references()
 
     def resolve(
-        self, context: Context, session: Session, *, include_xcom: bool = True
+        self, context: Context, session: Session, *, include_xcom: bool
     ) -> tuple[Mapping[str, Any], set[int]]:
         map_index = context["ti"].map_index
         if map_index < 0:
@@ -263,9 +266,9 @@ def resolve(
         if isinstance(self.value, collections.abc.Sized):
             mapping = self.value[map_index]
             if not isinstance(mapping, collections.abc.Mapping):
-                mapping = mapping.resolve(context, session)
+                mapping = mapping.resolve(context, session, include_xcom=include_xcom)
         elif include_xcom:
-            mappings = self.value.resolve(context, session)
+            mappings = self.value.resolve(context, session, include_xcom=include_xcom)
             if not isinstance(mappings, collections.abc.Sequence):
                 raise ValueError(f"expand_kwargs() expects a list[dict], not {_describe_type(mappings)}")
             mapping = mappings[map_index]
diff --git a/airflow/models/param.py b/airflow/models/param.py
index 3f56e50cd19e6..a5b9504a06bb6 100644
--- a/airflow/models/param.py
+++ b/airflow/models/param.py
@@ -329,7 +329,7 @@ def __init__(self, current_dag: DAG, name: str, default: Any = NOTSET):
     def iter_references(self) -> Iterable[tuple[Operator, str]]:
         return ()
 
-    def resolve(self, context: Context) -> Any:
+    def resolve(self, context: Context, *, include_xcom: bool) -> Any:
         """Pull DagParam value from DagRun context. This method is run during ``op.execute()``."""
         with contextlib.suppress(KeyError):
             return context["dag_run"].conf[self._name]
diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py
index bf7d8323e8fa0..108634bb7e0f9 100644
--- a/airflow/models/xcom_arg.py
+++ b/airflow/models/xcom_arg.py
@@ -208,7 +208,7 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
         raise NotImplementedError()
 
     @provide_session
-    def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
+    def resolve(self, context: Context, session: Session = NEW_SESSION, *, include_xcom: bool) -> Any:
         """
         Pull XCom value.
 
@@ -437,7 +437,7 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
         )
 
     @provide_session
-    def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
+    def resolve(self, context: Context, session: Session = NEW_SESSION, *, include_xcom: bool) -> Any:
         ti = context["ti"]
         if TYPE_CHECKING:
             assert isinstance(ti, TaskInstance)
@@ -551,8 +551,8 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
         return self.arg.get_task_map_length(run_id, session=session)
 
     @provide_session
-    def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
-        value = self.arg.resolve(context, session=session)
+    def resolve(self, context: Context, session: Session = NEW_SESSION, *, include_xcom: bool) -> Any:
+        value = self.arg.resolve(context, session=session, include_xcom=include_xcom)
         if not isinstance(value, (Sequence, dict)):
             raise ValueError(f"XCom map expects sequence or dict, not {type(value).__name__}")
         return _MapResult(value, self.callables)
@@ -632,8 +632,8 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
         return max(ready_lengths)
 
     @provide_session
-    def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
-        values = [arg.resolve(context, session=session) for arg in self.args]
+    def resolve(self, context: Context, session: Session = NEW_SESSION, *, include_xcom: bool) -> Any:
+        values = [arg.resolve(context, session=session, include_xcom=include_xcom) for arg in self.args]
         for value in values:
             if not isinstance(value, (Sequence, dict)):
                 raise ValueError(f"XCom zip expects sequence or dict, not {type(value).__name__}")
@@ -707,8 +707,8 @@ def get_task_map_length(self, run_id: str, *, session: Session) -> int | None:
         return sum(ready_lengths)
 
     @provide_session
-    def resolve(self, context: Context, session: Session = NEW_SESSION) -> Any:
-        values = [arg.resolve(context, session=session) for arg in self.args]
+    def resolve(self, context: Context, session: Session = NEW_SESSION, *, include_xcom: bool) -> Any:
+        values = [arg.resolve(context, session=session, include_xcom=include_xcom) for arg in self.args]
         for value in values:
             if not isinstance(value, (Sequence, dict)):
                 raise ValueError(f"XCom concat expects sequence or dict, not {type(value).__name__}")
diff --git a/airflow/template/templater.py b/airflow/template/templater.py
index bfbb6f5722c56..3dc10d36ff546 100644
--- a/airflow/template/templater.py
+++ b/airflow/template/templater.py
@@ -46,7 +46,7 @@ class LiteralValue(ResolveMixin):
     def iter_references(self) -> Iterable[tuple[Operator, str]]:
         return ()
 
-    def resolve(self, context: Context) -> Any:
+    def resolve(self, context: Context, *, include_xcom: bool) -> Any:
         return self.value
 
 
@@ -172,7 +172,7 @@ def render_template(
         if isinstance(value, ObjectStoragePath):
             return self._render_object_storage_path(value, context, jinja_env)
         if isinstance(value, ResolveMixin):
-            return value.resolve(context)
+            return value.resolve(context, include_xcom=True)
 
         # Fast path for common built-in collections.
         if value.__class__ is tuple:
diff --git a/airflow/utils/mixins.py b/airflow/utils/mixins.py
index eb0a1e81e1346..324299b3a6e59 100644
--- a/airflow/utils/mixins.py
+++ b/airflow/utils/mixins.py
@@ -64,7 +64,7 @@ def iter_references(self) -> typing.Iterable[tuple[Operator, str]]:
         """
         raise NotImplementedError
 
-    def resolve(self, context: Context) -> typing.Any:
+    def resolve(self, context: Context, *, include_xcom: bool) -> typing.Any:
         """
         Resolve this value for runtime.
 

From b53b12bd886cb343fec668174525a62bfa5c5077 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Fri, 12 Jul 2024 18:05:12 +0800
Subject: [PATCH 15/18] docs(deferring): add more instruciotn for mapped op
 with start_from_trigger support

---
 docs/apache-airflow/authoring-and-scheduling/deferring.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/apache-airflow/authoring-and-scheduling/deferring.rst b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
index c89d671523c2d..d30a665c83aa9 100644
--- a/docs/apache-airflow/authoring-and-scheduling/deferring.rst
+++ b/docs/apache-airflow/authoring-and-scheduling/deferring.rst
@@ -210,7 +210,7 @@ This is particularly useful when deferring is the only thing the ``execute`` met
             # We have no more work to do here. Mark as complete.
             return
 
-To enable Dynamic Task Mapping support, you can define ``start_from_trigger`` and ``trigger_kwargs`` in the parameter of "__init__". Note that you don't need to define both of them to use this feature, but you do need to use the exact same parameter name. For example, if you define an argument as ``t_kwargs`` and assign this value to ``self.start_trigger_args.trigger_kwargs``, it will not work.
+To enable Dynamic Task Mapping support, you can define ``start_from_trigger`` and ``trigger_kwargs`` in the parameter of "__init__". Note that you don't need to define both of them to use this feature, but you do need to use the exact same parameter name. For example, if you define an argument as ``t_kwargs`` and assign this value to ``self.start_trigger_args.trigger_kwargs``, it will not work. Note that this works different from mapping an operator without ``start_from_trigger`` support. The whole ``__init__`` method is skipped when mapping an operator whose ``start_from_trigger`` is set to True. Only argument ``trigger_kwargs`` is used and passed into ``trigger_cls``.
 
 .. code-block:: python
 

From 0c296298c22e7b127f76a393061dd5dce6c41f33 Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Sat, 20 Jul 2024 00:03:21 +0800
Subject: [PATCH 16/18] refactor(operator): make expand_start_from_trigger part
 of the abstract operator to unify how base op and mapped op retrieve
 start_from_trigger

---
 airflow/models/abstractoperator.py | 11 +++++++++++
 airflow/models/baseoperator.py     | 13 +++++++++++++
 airflow/models/dagrun.py           |  8 ++------
 airflow/models/mappedoperator.py   | 10 +++++++---
 4 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index 89e8b6cc723c0..1e9d434ee387f 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -427,6 +427,17 @@ def unmap(self, resolve: None | dict[str, Any] | tuple[Context, Session]) -> Bas
         """
         raise NotImplementedError()
 
+    def expand_start_from_trigger(self, *, context: Context, session: Session) -> bool:
+        """
+        Get the start_from_trigger value of the current abstract operator.
+
+        MappedOperator uses this to unmap start_from_trigger to decide whether to start the task
+        execution directly from triggerer.
+
+        :meta private:
+        """
+        raise NotImplementedError()
+
     @property
     def priority_weight_total(self) -> int:
         """
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 30ab5918671b6..7d194389817e5 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -1795,6 +1795,19 @@ def unmap(self, resolve: None | dict[str, Any] | tuple[Context, Session]) -> Bas
         """
         return self
 
+    def expand_start_from_trigger(self, *, context: Context, session: Session) -> bool:
+        """
+        Get the start_from_trigger value of the current abstract operator.
+
+        Since a BaseOperator is not mapped to begin with, this simply returns
+        the original value of start_from_trigger.
+        MappedOperator uses this to unmap start_from_trigger to decide whether to start the task
+        execution directly from triggerer.
+
+        :meta private:
+        """
+        return self.start_from_trigger
+
 
 # TODO: Deprecate for Airflow 3.0
 Chainable = Union[DependencyMixin, Sequence[DependencyMixin]]
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index cc22016dcd819..4eb40eb4b359e 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1582,13 +1582,9 @@ def schedule_tis(
             # this task.
             # if not, we'll add this "ti" into "schedulable_ti_ids" and later execute it to run in the worker
             elif ti.task.start_trigger_args is not None:
-                from airflow.models.mappedoperator import MappedOperator
 
-                if isinstance(ti.task, MappedOperator):
-                    context = ti.get_template_context()
-                    start_from_trigger = ti.task._expand_start_from_trigger(context=context, session=session)
-                else:
-                    start_from_trigger = ti.task.start_from_trigger
+                context = ti.get_template_context()
+                start_from_trigger = ti.task.expand_start_from_trigger(context=context, session=session)
 
                 if start_from_trigger:
                     ti.start_date = timezone.utcnow()
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index 9744247c5bd50..31b7c9ccc523c 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -731,12 +731,16 @@ def _get_unmap_kwargs(self, mapped_kwargs: Mapping[str, Any], *, strict: bool) -
             "params": params,
         }
 
-    def _expand_start_from_trigger(self, *, context: Context, session: Session) -> bool:
+    def expand_start_from_trigger(self, *, context: Context, session: Session) -> bool:
         """
-        Get the kwargs to create the unmapped start_from_trigger.
+        Get the start_from_trigger value of the current abstract operator.
 
-        This method is for allowing mapped operator to start execution from triggerer.
+        MappedOperator uses this to unmap start_from_trigger to decide whether to start the task
+        execution directly from triggerer.
+
+        :meta private:
         """
+        # start_from_trigger only makes sense when start_trigger_args exists.
         if not self.start_trigger_args:
             return False
 

From 3ee46d267b33d8976a9ca044efacdf47a47c90bb Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Sat, 20 Jul 2024 00:22:31 +0800
Subject: [PATCH 17/18] refactor(operator): make expand_start_trigger_args part
 of the abstract operator to unify how base op and mapped op retrieve
 start_trigger_args

---
 airflow/models/abstractoperator.py | 11 +++++++++++
 airflow/models/baseoperator.py     | 13 +++++++++++--
 airflow/models/dagrun.py           |  1 -
 airflow/models/mappedoperator.py   |  6 ++----
 airflow/models/taskinstance.py     | 12 +++++++-----
 5 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
index 1e9d434ee387f..9cf1830bb4389 100644
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -55,6 +55,7 @@
     from airflow.models.operator import Operator
     from airflow.models.taskinstance import TaskInstance
     from airflow.task.priority_strategy import PriorityWeightStrategy
+    from airflow.triggers.base import StartTriggerArgs
     from airflow.utils.task_group import TaskGroup
 
 DEFAULT_OWNER: str = conf.get_mandatory_value("operators", "default_owner")
@@ -438,6 +439,16 @@ def expand_start_from_trigger(self, *, context: Context, session: Session) -> bo
         """
         raise NotImplementedError()
 
+    def expand_start_trigger_args(self, *, context: Context, session: Session) -> StartTriggerArgs | None:
+        """
+        Get the start_trigger_args value of the current abstract operator.
+
+        MappedOperator uses this to unmap start_trigger_args to decide how to start a task from triggerer.
+
+        :meta private:
+        """
+        raise NotImplementedError()
+
     @property
     def priority_weight_total(self) -> int:
         """
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 7d194389817e5..8525b78f60de5 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -1801,13 +1801,22 @@ def expand_start_from_trigger(self, *, context: Context, session: Session) -> bo
 
         Since a BaseOperator is not mapped to begin with, this simply returns
         the original value of start_from_trigger.
-        MappedOperator uses this to unmap start_from_trigger to decide whether to start the task
-        execution directly from triggerer.
 
         :meta private:
         """
         return self.start_from_trigger
 
+    def expand_start_trigger_args(self, *, context: Context, session: Session) -> StartTriggerArgs | None:
+        """
+        Get the start_trigger_args value of the current abstract operator.
+
+        Since a BaseOperator is not mapped to begin with, this simply returns
+        the original value of start_trigger_args.
+
+        :meta private:
+        """
+        return self.start_trigger_args
+
 
 # TODO: Deprecate for Airflow 3.0
 Chainable = Union[DependencyMixin, Sequence[DependencyMixin]]
diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py
index 4eb40eb4b359e..d47be6e74bab1 100644
--- a/airflow/models/dagrun.py
+++ b/airflow/models/dagrun.py
@@ -1582,7 +1582,6 @@ def schedule_tis(
             # this task.
             # if not, we'll add this "ti" into "schedulable_ti_ids" and later execute it to run in the worker
             elif ti.task.start_trigger_args is not None:
-
                 context = ti.get_template_context()
                 start_from_trigger = ti.task.expand_start_from_trigger(context=context, session=session)
 
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index 31b7c9ccc523c..2377fdab00756 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -757,16 +757,14 @@ def expand_start_from_trigger(self, *, context: Context, session: Session) -> bo
             "start_from_trigger", self.partial_kwargs.get("start_from_trigger", self.start_from_trigger)
         )
 
-    def _expand_start_trigger_args(self, *, context: Context, session: Session) -> StartTriggerArgs:
+    def expand_start_trigger_args(self, *, context: Context, session: Session) -> StartTriggerArgs | None:
         """
         Get the kwargs to create the unmapped start_trigger_args.
 
         This method is for allowing mapped operator to start execution from triggerer.
         """
         if not self.start_trigger_args:
-            raise AirflowException(
-                "Cannot expand start_trigger_args for mapped operator without class level start_trigger_args"
-            )
+            return None
 
         mapped_kwargs, _ = self._expand_mapped_kwargs(context, session, include_xcom=False)
         if self._disallow_kwargs_override:
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index a20175967c312..96d9a3df5ca26 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -83,6 +83,7 @@
     AirflowTaskTimeout,
     DagRunNotFound,
     RemovedInAirflow3Warning,
+    TaskDeferralError,
     TaskDeferred,
     UnmappableXComLengthPushed,
     UnmappableXComTypePushed,
@@ -1617,11 +1618,12 @@ def _defer_task(
         next_kwargs = exception.kwargs
         timeout = exception.timeout
     elif ti.task is not None and ti.task.start_trigger_args is not None:
-        if isinstance(ti.task, MappedOperator):
-            context = ti.get_template_context()
-            start_trigger_args = ti.task._expand_start_trigger_args(context=context, session=session)
-        else:
-            start_trigger_args = ti.task.start_trigger_args
+        context = ti.get_template_context()
+        start_trigger_args = ti.task.expand_start_trigger_args(context=context, session=session)
+        if start_trigger_args is None:
+            raise TaskDeferralError(
+                "A none 'None' start_trigger_args has been change to 'None' during expandion"
+            )
 
         trigger_kwargs = start_trigger_args.trigger_kwargs or {}
         next_kwargs = start_trigger_args.next_kwargs

From 636265e11961b673a70cde8c951f7c2759cfc13f Mon Sep 17 00:00:00 2001
From: Wei Lee <weilee.rx@gmail.com>
Date: Sat, 20 Jul 2024 00:24:53 +0800
Subject: [PATCH 18/18] refactor(taskinstance): change the exception type of
 "no exception and validate start_trigger_args passed to _defer_task" to
 TaskDeferralError

---
 airflow/models/taskinstance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 96d9a3df5ca26..086e89b0ebf55 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1634,7 +1634,7 @@ def _defer_task(
             kwargs=trigger_kwargs,
         )
     else:
-        raise AirflowException("exception and ti.task.start_trigger_args cannot both be None")
+        raise TaskDeferralError("exception and ti.task.start_trigger_args cannot both be None")
 
     # First, make the trigger entry
     session.add(trigger_row)