Repository: apache/airflow
Pull request #48819: Fix WorkflowTrigger to work with TaskSDK
Author: login: gopidesupavan
Labels: name: area:providers
name: area:Triggerer
name: provider:standard
Created: 2025-04-04T23:25:08Z | Merged: 2025-04-09T16:50:52Z

closes: #48425


<img width="1709" alt="image" src="https://github.com/user-attachments/assets/dfbe062a-bc10-4340-a05a-aa42bb18fbf9" />


<!--
 Licensed to the Apache Software Foundation (ASF) under one
 or more contributor license agreements.  See the NOTICE file
 distributed with this work for additional information
 regarding copyright ownership.  The ASF licenses this file
 to you under the Apache License, Version 2.0 (the
 "License"); you may not use this file except in compliance
 with the License.  You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing,
 software distributed under the License is distributed on an
 "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 KIND, either express or implied.  See the License for the
 specific language governing permissions and limitations
 under the License.
 -->

<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.

Feel free to ping committers for the review!

In case of an existing issue, reference it using one of the following:

closes: #ISSUE
related: #ISSUE

How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->



<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).



Diff:
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 520699e9e5093..be7fe23004b1c 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -50,8 +50,10 @@
     GetConnection,
     GetDagRunState,
     GetDRCount,
+    GetTICount,
     GetVariable,
     GetXCom,
+    TICount,
     VariableResult,
     XComResult,
 )
@@ -222,6 +224,7 @@ class TriggerStateSync(BaseModel):
         XComResult,
         DagRunStateResult,
         DRCount,
+        TICount,
         ErrorResponse,
     ],
     Field(discriminator="type"),
@@ -233,7 +236,15 @@ class TriggerStateSync(BaseModel):
 
 
 ToTriggerSupervisor = Annotated[
-    Union[messages.TriggerStateChanges, GetConnection, GetVariable, GetXCom, GetDagRunState, GetDRCount],
+    Union[
+        messages.TriggerStateChanges,
+        GetConnection,
+        GetVariable,
+        GetXCom,
+        GetTICount,
+        GetDagRunState,
+        GetDRCount,
+    ],
     Field(discriminator="type"),
 ]
 """
@@ -411,6 +422,16 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -
             dr_resp = self.client.dag_runs.get_state(msg.dag_id, msg.run_id)
             resp = DagRunStateResult.from_api_response(dr_resp).model_dump_json().encode()
 
+        elif isinstance(msg, GetTICount):
+            ti_count = self.client.task_instances.get_count(
+                dag_id=msg.dag_id,
+                task_ids=msg.task_ids,
+                task_group_id=msg.task_group_id,
+                logical_dates=msg.logical_dates,
+                run_ids=msg.run_ids,
+                states=msg.states,
+            )
+            resp = ti_count.model_dump_json().encode()
         else:
             raise ValueError(f"Unknown message type {type(msg)}")
 
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index af197d6f35bfa..50320681df1de 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -767,3 +767,103 @@ async def test_trigger_can_fetch_trigger_dag_run_count_in_deferrable(session, da
     assert task_instance.state == TaskInstanceState.SCHEDULED
     assert task_instance.next_method != "__fail__"
     assert task_instance.next_kwargs == {"event": {"count": 1}}
+
+
+class CustomTriggerWorkflowStateTrigger(BaseTrigger):
+    """Custom Trigger to check the triggerer can access the get_ti_count and get_dr_count."""
+
+    def __init__(self, external_dag_id, execution_dates, external_task_ids, allowed_states, run_ids):
+        self.external_dag_id = external_dag_id
+        self.execution_dates = execution_dates
+        self.external_task_ids = external_task_ids
+        self.allowed_states = allowed_states
+        self.run_ids = run_ids
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            f"{type(self).__module__}.{type(self).__qualname__}",
+            {
+                "external_dag_id": self.external_dag_id,
+                "execution_dates": self.execution_dates,
+                "external_task_ids": self.external_task_ids,
+                "allowed_states": self.allowed_states,
+                "run_ids": self.run_ids,
+            },
+        )
+
+    async def run(self, **args) -> AsyncIterator[TriggerEvent]:
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        ti_count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+            dag_id=self.external_dag_id,
+            task_ids=self.external_task_ids,
+            task_group_id=None,
+            run_ids=self.run_ids,
+            logical_dates=self.execution_dates,
+            states=self.allowed_states,
+        )
+        dr_count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+            dag_id=self.external_dag_id,
+            run_ids=self.run_ids,
+            logical_dates=self.execution_dates,
+            states=["running"],
+        )
+        yield TriggerEvent({"ti_count": ti_count, "dr_count": dr_count})
+
+
+@pytest.mark.xfail(
+    reason="We know that test is flaky and have no time to fix it before 3.0. "
+    "We should fix it later. TODO: AIP-72"
+)
+@pytest.mark.asyncio
+@pytest.mark.flaky(reruns=2, reruns_delay=10)
+@pytest.mark.execution_timeout(30)
+async def test_trigger_can_fetch_dag_run_count_ti_count_in_deferrable(session, dag_maker):
+    """Checks that the trigger will successfully fetch the count of trigger DAG runs."""
+    # Create the test DAG and task
+    with dag_maker(dag_id="parent_dag", session=session):
+        EmptyOperator(task_id="parent_task")
+    parent_dag_run = dag_maker.create_dagrun()
+    parent_task = parent_dag_run.task_instances[0]
+    parent_task.state = TaskInstanceState.SUCCESS
+
+    with dag_maker(dag_id="trigger_can_fetch_dag_run_count_ti_count_in_deferrable", session=session):
+        EmptyOperator(task_id="dummy1")
+    dr = dag_maker.create_dagrun()
+    task_instance = dr.task_instances[0]
+    task_instance.state = TaskInstanceState.DEFERRED
+
+    # Use the same dag run with states deferred to fetch the count
+    trigger = CustomTriggerWorkflowStateTrigger(
+        external_dag_id=parent_task.dag_id,
+        execution_dates=[parent_task.logical_date],
+        external_task_ids=[parent_task.task_id],
+        allowed_states=[State.SUCCESS],
+        run_ids=[parent_task.run_id],
+    )
+    trigger_orm = Trigger(
+        classpath=trigger.serialize()[0],
+        kwargs={
+            "external_dag_id": parent_dag_run.dag_id,
+            "execution_dates": [parent_dag_run.logical_date],
+            "external_task_ids": [parent_task.task_id],
+            "allowed_states": [State.SUCCESS],
+            "run_ids": [parent_dag_run.run_id],
+        },
+    )
+    session.add(trigger_orm)
+    session.commit()
+    task_instance.trigger_id = trigger_orm.id
+
+    job = Job()
+    session.add(job)
+    session.commit()
+
+    supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=1, logger=None)
+    supervisor.run()
+
+    parent_task.refresh_from_db()
+    task_instance.refresh_from_db()
+    assert task_instance.state == TaskInstanceState.SCHEDULED
+    assert task_instance.next_method != "__fail__"
+    assert task_instance.next_kwargs == {"event": {"ti_count": 1, "dr_count": 1}}
diff --git a/providers/standard/src/airflow/providers/standard/sensors/external_task.py b/providers/standard/src/airflow/providers/standard/sensors/external_task.py
index dd5ad6c4e7a3a..e64eb0d6763b0 100644
--- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py
+++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py
@@ -346,8 +346,6 @@ def _calculate_count(self, count: int, dttm_filter: list[datetime.datetime]) ->
         """Calculate the normalized count based on the type of check."""
         if self.external_task_ids:
             return count / len(self.external_task_ids)
-        elif self.external_task_group_id:
-            return count / len(dttm_filter)
         else:
             return count
 
@@ -421,16 +419,22 @@ def execute(self, context: Context) -> None:
         if not self.deferrable:
             super().execute(context)
         else:
+            dttm_filter = self._get_dttm_filter(context)
+            logical_or_execution_dates = (
+                {"logical_dates": dttm_filter} if AIRFLOW_V_3_0_PLUS else {"execution_date": dttm_filter}
+            )
             self.defer(
                 timeout=self.execution_timeout,
                 trigger=WorkflowTrigger(
                     external_dag_id=self.external_dag_id,
                     external_task_group_id=self.external_task_group_id,
                     external_task_ids=self.external_task_ids,
-                    logical_dates=self._get_dttm_filter(context),
                     allowed_states=self.allowed_states,
+                    failed_states=self.failed_states,
+                    skipped_states=self.skipped_states,
                     poke_interval=self.poll_interval,
                     soft_fail=self.soft_fail,
+                    **logical_or_execution_dates,
                 ),
                 method_name="execute_complete",
             )
diff --git a/providers/standard/src/airflow/providers/standard/triggers/external_task.py b/providers/standard/src/airflow/providers/standard/triggers/external_task.py
index adf3a535e3ef2..a5f8b67f54972 100644
--- a/providers/standard/src/airflow/providers/standard/triggers/external_task.py
+++ b/providers/standard/src/airflow/providers/standard/triggers/external_task.py
@@ -50,6 +50,7 @@ class WorkflowTrigger(BaseTrigger):
     :param allowed_states: States considered as successful for external tasks.
     :param poke_interval: The interval (in seconds) for poking the external tasks.
     :param soft_fail: If True, the trigger will not fail the entire dag on external task failure.
+    :param logical_dates: A list of logical dates for the external dag.
     """
 
     def __init__(
@@ -57,6 +58,7 @@ def __init__(
         external_dag_id: str,
         run_ids: list[str] | None = None,
         execution_dates: list[datetime] | None = None,
+        logical_dates: list[datetime] | None = None,
         external_task_ids: typing.Collection[str] | None = None,
         external_task_group_id: str | None = None,
         failed_states: typing.Iterable[str] | None = None,
@@ -76,6 +78,7 @@ def __init__(
         self.poke_interval = poke_interval
         self.soft_fail = soft_fail
         self.execution_dates = execution_dates
+        self.logical_dates = logical_dates
         super().__init__(**kwargs)
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
@@ -92,6 +95,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
+            data["logical_dates"] = self.logical_dates
         else:
             data["execution_dates"] = self.execution_dates
 
@@ -99,9 +103,16 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        if AIRFLOW_V_3_0_PLUS:
+            get_count_func = self._get_count_af_3
+            run_id_or_dates = (self.run_ids or self.logical_dates) or []
+        else:
+            get_count_func = self._get_count
+            run_id_or_dates = self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
@@ -109,18 +120,43 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.logical_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.logical_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        else:
+            return count
+
     @sync_to_async
     def _get_count(self, states: typing.Iterable[str] | None) -> int:
         """
diff --git a/providers/standard/tests/unit/standard/triggers/test_external_task.py b/providers/standard/tests/unit/standard/triggers/test_external_task.py
index 247085a407ff4..3aa581257d79e 100644
--- a/providers/standard/tests/unit/standard/triggers/test_external_task.py
+++ b/providers/standard/tests/unit/standard/triggers/test_external_task.py
@@ -39,11 +39,315 @@
 key, value = next(iter(_DATES.items()))
 
 
+@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 3")
 class TestWorkflowTrigger:
     DAG_ID = "external_task"
     TASK_ID = "external_task_op"
     RUN_ID = "external_task_run_id"
     STATES = ["success", "fail"]
+    LOGICAL_DATE = timezone.datetime(2022, 1, 1)
+
+    @pytest.mark.flaky(reruns=5)
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count")
+    @pytest.mark.asyncio
+    async def test_task_workflow_trigger_success(self, mock_get_count):
+        """check the db count get called correctly."""
+        mock_get_count.side_effect = mocked_get_count
+
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=[self.LOGICAL_DATE],
+            external_task_ids=[self.TASK_ID],
+            allowed_states=self.STATES,
+            poke_interval=0.2,
+        )
+
+        gen = trigger.run()
+        trigger_task = asyncio.create_task(gen.__anext__())
+        fake_task = asyncio.create_task(fake_async_fun())
+        await trigger_task
+        await fake_task
+        assert fake_task.done()  # confirm that get_count is done in an async fashion
+        assert trigger_task.done()
+        result = trigger_task.result()
+        assert result.payload == {"status": "success"}
+
+        mock_get_count.assert_called_once_with(
+            dag_id="external_task",
+            task_ids=["external_task_op"],
+            task_group_id=None,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=None,
+            states=["success", "fail"],
+        )
+        # test that it returns after yielding
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+    @pytest.mark.flaky(reruns=5)
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count")
+    @pytest.mark.asyncio
+    async def test_task_workflow_trigger_failed(self, mock_get_count):
+        mock_get_count.side_effect = mocked_get_count
+
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=[self.RUN_ID],
+            external_task_ids=[self.TASK_ID],
+            failed_states=self.STATES,
+            poke_interval=0.2,
+        )
+
+        gen = trigger.run()
+        trigger_task = asyncio.create_task(gen.__anext__())
+        fake_task = asyncio.create_task(fake_async_fun())
+        await trigger_task
+        await fake_task
+        assert fake_task.done()  # confirm that get_count is done in an async fashion
+        assert trigger_task.done()
+        result = trigger_task.result()
+        assert isinstance(result, TriggerEvent)
+        assert result.payload == {"status": "failed"}
+        mock_get_count.assert_called_once_with(
+            dag_id="external_task",
+            task_ids=["external_task_op"],
+            task_group_id=None,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=[self.RUN_ID],
+            states=["success", "fail"],
+        )
+        # test that it returns after yielding
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+    @pytest.mark.asyncio
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count")
+    async def test_task_workflow_trigger_fail_count_eq_0(self, mock_get_count):
+        mock_get_count.return_value = 0
+
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=[self.RUN_ID],
+            external_task_ids=[self.TASK_ID],
+            failed_states=self.STATES,
+            poke_interval=0.2,
+        )
+
+        gen = trigger.run()
+        trigger_task = asyncio.create_task(gen.__anext__())
+        await trigger_task
+        assert trigger_task.done()
+        result = trigger_task.result()
+        assert isinstance(result, TriggerEvent)
+        assert result.payload == {"status": "success"}
+        mock_get_count.assert_called_once_with(
+            dag_id="external_task",
+            task_ids=["external_task_op"],
+            task_group_id=None,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=[self.RUN_ID],
+            states=["success", "fail"],
+        )
+        # test that it returns after yielding
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+    @pytest.mark.flaky(reruns=5)
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count")
+    @pytest.mark.asyncio
+    async def test_task_workflow_trigger_skipped(self, mock_get_count):
+        mock_get_count.side_effect = mocked_get_count
+
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=[self.LOGICAL_DATE],
+            external_task_ids=[self.TASK_ID],
+            skipped_states=self.STATES,
+            poke_interval=0.2,
+        )
+
+        gen = trigger.run()
+        trigger_task = asyncio.create_task(gen.__anext__())
+        fake_task = asyncio.create_task(fake_async_fun())
+        await trigger_task
+        await fake_task
+        assert fake_task.done()  # confirm that get_count is done in an async fashion
+        assert trigger_task.done()
+        result = trigger_task.result()
+        assert isinstance(result, TriggerEvent)
+        assert result.payload == {"status": "skipped"}
+        mock_get_count.assert_called_once_with(
+            dag_id="external_task",
+            task_ids=["external_task_op"],
+            task_group_id=None,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=None,
+            states=["success", "fail"],
+        )
+
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count")
+    @mock.patch("asyncio.sleep")
+    @pytest.mark.asyncio
+    async def test_task_workflow_trigger_sleep_success(self, mock_sleep, mock_get_count):
+        mock_get_count.side_effect = [0, 1]
+
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=[self.LOGICAL_DATE],
+            external_task_ids=[self.TASK_ID],
+            poke_interval=0.2,
+        )
+
+        gen = trigger.run()
+        trigger_task = asyncio.create_task(gen.__anext__())
+        await trigger_task
+        assert trigger_task.done()
+        result = trigger_task.result()
+        assert isinstance(result, TriggerEvent)
+        assert result.payload == {"status": "success"}
+        mock_get_count.assert_called()
+        assert mock_get_count.call_count == 2
+
+        # test that it returns after yielding
+        with pytest.raises(StopAsyncIteration):
+            await gen.__anext__()
+
+        mock_sleep.assert_awaited()
+        assert mock_sleep.await_count == 1
+
+    @pytest.mark.parametrize(
+        "task_ids, task_group_id, states, logical_dates, mock_ti_count, mock_dag_count, expected",
+        [
+            (
+                ["task_id_one", "task_id_two"],
+                None,
+                ["success"],
+                [
+                    timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc),
+                    timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc),
+                ],
+                4,
+                2,
+                2,
+            ),
+            (
+                [],
+                "task_group_id",
+                ["success"],
+                [
+                    timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc),
+                    timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc),
+                ],
+                2,
+                2,
+                2,
+            ),
+            (
+                [],
+                None,
+                ["success"],
+                [
+                    timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc),
+                    timezone.datetime(2020, 7, 6, 13, tzinfo=timezone.utc),
+                ],
+                2,
+                2,
+                2,
+            ),
+        ],
+        ids=[
+            "with_task_ids",
+            "with task_group_id only",
+            "no task_ids or task_group_id",
+        ],
+    )
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_ti_count")
+    @mock.patch("airflow.sdk.execution_time.task_runner.RuntimeTaskInstance.get_dr_count")
+    @pytest.mark.asyncio
+    async def test_get_count_af_3(
+        self,
+        mock_get_dr_count,
+        mock_get_ti_count,
+        task_ids,
+        task_group_id,
+        states,
+        logical_dates,
+        mock_ti_count,
+        mock_dag_count,
+        expected,
+    ):
+        """
+        case1: when provided two task_ids, and two dag runs, the get_ti_count should return 4(each dag run returns two tasks)
+                and normalized count becomes 2
+        case2: when provided task_group_id, and two dag runs, the get_ti_count should return 2(each dag run returns 1 task group)
+        case3: when not provided any task_ids or task_group_id, the get_dr_count should return 2(total dag runs 2)
+        """
+
+        mock_get_ti_count.return_value = mock_ti_count
+        mock_get_dr_count.return_value = mock_dag_count
+
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=logical_dates,
+            external_task_ids=task_ids,
+            external_task_group_id=task_group_id,
+            allowed_states=states,
+            poke_interval=0.2,
+        )
+
+        get_count_af_3 = await trigger._get_count_af_3(states)
+        assert get_count_af_3 == expected
+
+        if task_ids or task_group_id:
+            mock_get_ti_count.assert_called_once()
+            assert mock_get_ti_count.call_count == 1
+            mock_get_dr_count.assert_not_called()
+            assert mock_get_dr_count.call_count == 0
+
+        if not task_ids and not task_group_id:
+            mock_get_dr_count.assert_called_once()
+            assert mock_get_dr_count.call_count == 1
+            mock_get_ti_count.assert_not_called()
+            assert mock_get_ti_count.call_count == 0
+
+    def test_serialization(self):
+        """
+        Asserts that the WorkflowTrigger correctly serializes its arguments and classpath.
+        """
+        trigger = WorkflowTrigger(
+            external_dag_id=self.DAG_ID,
+            logical_dates=[self.LOGICAL_DATE],
+            run_ids=[self.RUN_ID],
+            failed_states=["failed"],
+            skipped_states=["skipped"],
+            external_task_ids=[self.TASK_ID],
+            allowed_states=self.STATES,
+            poke_interval=5,
+        )
+        classpath, kwargs = trigger.serialize()
+        assert classpath == "airflow.providers.standard.triggers.external_task.WorkflowTrigger"
+        assert kwargs == {
+            "external_dag_id": self.DAG_ID,
+            "logical_dates": [self.LOGICAL_DATE],
+            "run_ids": [self.RUN_ID],
+            "external_task_ids": [self.TASK_ID],
+            "external_task_group_id": None,
+            "failed_states": ["failed"],
+            "skipped_states": ["skipped"],
+            "allowed_states": self.STATES,
+            "poke_interval": 5,
+            "soft_fail": False,
+        }
+
+
+@pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test only for Airflow 2")
+class TestWorkflowTriggerAF2:
+    DAG_ID = "external_task"
+    TASK_ID = "external_task_op"
+    RUN_ID = "external_task_run_id"
+    STATES = ["success", "fail"]
 
     @pytest.mark.flaky(reruns=5)
     @mock.patch("airflow.providers.standard.triggers.external_task._get_count")
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 9cecdc5547d44..2d5a3add593bd 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -413,18 +413,19 @@ def get_ti_count(
         """Return the number of task instances matching the given criteria."""
         log = structlog.get_logger(logger_name="task")
 
-        SUPERVISOR_COMMS.send_request(
-            log=log,
-            msg=GetTICount(
-                dag_id=dag_id,
-                task_ids=task_ids,
-                task_group_id=task_group_id,
-                logical_dates=logical_dates,
-                run_ids=run_ids,
-                states=states,
-            ),
-        )
-        response = SUPERVISOR_COMMS.get_message()
+        with SUPERVISOR_COMMS.lock:
+            SUPERVISOR_COMMS.send_request(
+                log=log,
+                msg=GetTICount(
+                    dag_id=dag_id,
+                    task_ids=task_ids,
+                    task_group_id=task_group_id,
+                    logical_dates=logical_dates,
+                    run_ids=run_ids,
+                    states=states,
+                ),
+            )
+            response = SUPERVISOR_COMMS.get_message()
 
         if TYPE_CHECKING:
             assert isinstance(response, TICount)
@@ -441,16 +442,17 @@ def get_dr_count(
         """Return the number of DAG runs matching the given criteria."""
         log = structlog.get_logger(logger_name="task")
 
-        SUPERVISOR_COMMS.send_request(
-            log=log,
-            msg=GetDRCount(
-                dag_id=dag_id,
-                logical_dates=logical_dates,
-                run_ids=run_ids,
-                states=states,
-            ),
-        )
-        response = SUPERVISOR_COMMS.get_message()
+        with SUPERVISOR_COMMS.lock:
+            SUPERVISOR_COMMS.send_request(
+                log=log,
+                msg=GetDRCount(
+                    dag_id=dag_id,
+                    logical_dates=logical_dates,
+                    run_ids=run_ids,
+                    states=states,
+                ),
+            )
+            response = SUPERVISOR_COMMS.get_message()
 
         if TYPE_CHECKING:
             assert isinstance(response, DRCount)


Reviews:
user: login: ashb
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T21:32:54Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T21:46:38Z
user: login: kaxil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T17:01:36Z
user: login: kaxil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T17:02:34Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T17:05:08Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T17:59:53Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T18:01:48Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T18:14:21Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T20:33:40Z
user: login: kaxil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T12:38:25Z
user: login: kaxil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T12:40:24Z
user: login: kaxil
body: Not sure of the task group thing, rest lgtm
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-09T12:40:49Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T16:29:47Z
user: login: kaxil
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T16:45:23Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T16:50:03Z
user: login: gopidesupavan
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T16:50:29Z

Review comments:
path: task-sdk/src/airflow/sdk/execution_time/context.py
user: login: ashb
body: Nit: would you mind undoing this change -- it's best not to touch a file if the change isn't related to the PR, and this one is def more of a style nit
created_at: 2025-04-07T21:32:54Z
author_association: MEMBER
path: task-sdk/src/airflow/sdk/execution_time/context.py
user: login: gopidesupavan
body: Yeah, my bad updated now..
created_at: 2025-04-07T21:46:38Z
author_association: MEMBER
diff_hunk: @@ -427,8 +427,10 @@ def execute(self, context: Context) -> None:
                     external_dag_id=self.external_dag_id,
                     external_task_group_id=self.external_task_group_id,
                     external_task_ids=self.external_task_ids,
-                    logical_dates=self._get_dttm_filter(context),
+                    execution_dates=self._get_dttm_filter(context),
path: providers/standard/src/airflow/providers/standard/sensors/external_task.py
user: login: kaxil
body: Shouldn't we have `logical_dates` instead as `execution_dates` is old term
created_at: 2025-04-08T17:01:36Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: kaxil
body: Does this work?
created_at: 2025-04-08T17:02:34Z
author_association: MEMBER
diff_hunk: @@ -427,8 +427,10 @@ def execute(self, context: Context) -> None:
                     external_dag_id=self.external_dag_id,
                     external_task_group_id=self.external_task_group_id,
                     external_task_ids=self.external_task_ids,
-                    logical_dates=self._get_dttm_filter(context),
+                    execution_dates=self._get_dttm_filter(context),
path: providers/standard/src/airflow/providers/standard/sensors/external_task.py
user: login: gopidesupavan
body: Yeah was checking the confluence page about this, will push the changes..
created_at: 2025-04-08T17:05:08Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: gopidesupavan
body: Yeah its working, I ran one of the the example_dag.

```
from __future__ import annotations

import pendulum

from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.sensors.external_task import ExternalTaskMarker, ExternalTaskSensor
from airflow.sdk import DAG

start_date = pendulum.datetime(2021, 1, 1, tz="UTC")

with DAG(
    dag_id="example_external_task_marker_parent",
    start_date=start_date,
    catchup=False,
    schedule=None,
    tags=["example2"],
) as parent_dag:
    # [START howto_operator_external_task_marker]
    parent_task = ExternalTaskMarker(
        task_id="parent_task",
        external_dag_id="example_external_task_marker_child",
        external_task_id="child_task1",
    )
    # [END howto_operator_external_task_marker]

with DAG(
    dag_id="example_external_task_marker_child",
    start_date=start_date,
    schedule=None,
    catchup=False,
    tags=["example2"],
) as child_dag:
    # [START howto_operator_external_task_sensor]
    child_task1 = ExternalTaskSensor(
        task_id="child_task1",
        external_dag_id=parent_dag.dag_id,
        external_task_id=parent_task.task_id,
        timeout=600,
        allowed_states=["success"],
        failed_states=["failed", "skipped"],
        deferrable=True,
    )
    # [END howto_operator_external_task_sensor]

    # [START howto_operator_external_task_sensor_with_task_group]
    child_task2 = ExternalTaskSensor(
        task_id="child_task2",
        external_dag_id=parent_dag.dag_id,
        external_task_group_id="parent_dag_task_group_id",
        timeout=600,
        allowed_states=["success"],
        failed_states=["failed", "skipped"],
        deferrable=True,
    )
    # [END howto_operator_external_task_sensor_with_task_group]

    child_task3 = EmptyOperator(task_id="child_task3")
    child_task1 >> child_task2 >> child_task3
```

<img width="1721" alt="image" src="https://github.com/user-attachments/assets/f318824e-9894-4248-87e6-87870945e65f" />

created_at: 2025-04-08T17:59:53Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: gopidesupavan
body: Or  you have doubt about these calls? `RuntimeTaskInstance.get_dr_count` or  `RuntimeTaskInstance.get_ti_count`. as these are working i have used them as is. if you think any problem, i am happy to create new async methods for these may be in context file?
created_at: 2025-04-08T18:01:48Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: gopidesupavan
body: Let me double check..
created_at: 2025-04-08T18:14:21Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: gopidesupavan
body: Hey @kaxil good shout, fixed it. also updated in existing logic.

If the task_group_id provides we no need to divide that we can just assert with the length of dates.

example: 

execution_dates = [date1, date2], task_group_id='task_group_id1' -> so the get_ti_count returns 2 in this case when success criteria met. 

WDYT?
created_at: 2025-04-08T20:33:40Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: kaxil
body: a task group might have more tasks, and we would want all of them to succeed though
created_at: 2025-04-09T12:38:24Z
author_association: MEMBER
diff_hunk: @@ -346,8 +346,6 @@ def _calculate_count(self, count: int, dttm_filter: list[datetime.datetime]) ->
         """Calculate the normalized count based on the type of check."""
         if self.external_task_ids:
             return count / len(self.external_task_ids)
-        elif self.external_task_group_id:
-            return count / len(dttm_filter)
path: providers/standard/src/airflow/providers/standard/sensors/external_task.py
user: login: kaxil
body: not sure on this one
created_at: 2025-04-09T12:40:23Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: gopidesupavan
body: Yes agree with you, Thinking about this one more time, task_group condition will never met in the current logic. i think we would have to return count of succeeded task_groups.
 
For example, if a task group contains 3 tasks and the user provides 2 execution dates, the get_ti_count function returns 6, assuming all tasks succeed. However, we can't directly validate this against the number of dates provided (2). Even if we divide the total count (6) by the number of dates (2) and get 3, this value still can't be reliably used to validate against the dates alone.

So for the task_group_id case what would be the correct validation not sure now i am getting doubt 🤔 







created_at: 2025-04-09T16:29:47Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: kaxil
body: Feel free to merge it and handle that in a separate PR.

I think the logic in 2.x was correct, we should fix that for both sync and async

https://github.com/apache/airflow/blob/546675500e0953fdd3d832ab6263b2e18dfc1037/providers/standard/src/airflow/providers/standard/utils/sensor_helper.py#L62-L77
created_at: 2025-04-09T16:45:23Z
author_association: MEMBER
diff_hunk: @@ -89,38 +89,69 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
             "allowed_states": self.allowed_states,
             "poke_interval": self.poke_interval,
             "soft_fail": self.soft_fail,
+            "execution_dates": self.execution_dates,
         }
         if AIRFLOW_V_3_0_PLUS:
             data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
 
         return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
 
     async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
+        get_count_func = self._get_count_af_3 if AIRFLOW_V_3_0_PLUS else self._get_count
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
         while True:
             if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
+                failed_count = await get_count_func(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
+                skipped_count = await get_count_func(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_dates):  # type: ignore[arg-type]
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    async def _get_count_af_3(self, states):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        run_id_or_dates = self.run_ids or self.execution_dates or []
+
+        if self.external_task_ids or self.external_task_group_id:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                dag_id=self.external_dag_id,
+                task_ids=self.external_task_ids,
+                task_group_id=self.external_task_group_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.execution_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        elif self.external_task_group_id:
+            return count / len(run_id_or_dates)
+        else:
+            return count
path: providers/standard/src/airflow/providers/standard/triggers/external_task.py
user: login: gopidesupavan
body: sure, thanks kaxil.
created_at: 2025-04-09T16:50:03Z
author_association: MEMBER
diff_hunk: @@ -346,8 +346,6 @@ def _calculate_count(self, count: int, dttm_filter: list[datetime.datetime]) ->
         """Calculate the normalized count based on the type of check."""
         if self.external_task_ids:
             return count / len(self.external_task_ids)
-        elif self.external_task_group_id:
-            return count / len(dttm_filter)
path: providers/standard/src/airflow/providers/standard/sensors/external_task.py
user: login: gopidesupavan
body: Will be fixing this separate PR
created_at: 2025-04-09T16:50:29Z
author_association: MEMBER

PR comments:
user: login: gopidesupavan
created_at: 2025-04-04T23:29:42Z
body: @ashb @kaxil i have see that even for WorkflowTrigger, when multiple tasks in deferrable, the messages are mixing up. This happens only when continuously tasks entering into trigger.

<img width="1263" alt="image" src="https://github.com/user-attachments/assets/e7c8f23f-8caa-440f-bcee-fe48f08abddf" />

author_association: MEMBER
user: login: gopidesupavan
created_at: 2025-04-07T21:18:32Z
body: Also this is ready for review..
author_association: MEMBER
user: login: kaxil
created_at: 2025-04-08T08:37:40Z
body: @gopidesupavan Could you resolve the coflicts please?
author_association: MEMBER
user: login: gopidesupavan
created_at: 2025-04-08T16:18:50Z
body: > @gopidesupavan Could you resolve the coflicts please?

Yeah rebased now..
author_association: MEMBER
user: login: kaxil
created_at: 2025-04-09T13:37:44Z
body: (I rebased on main since the failures are fixed in main already _I think_)
author_association: MEMBER