From 49a6aa72899463fdd9e5ec8cbc0c06e67a5d8e30 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Thu, 4 Jun 2026 15:23:21 +0530 Subject: [PATCH 1/4] Fix rescheduled sensors hanging before poke closes: #68010 --- .../execution_api/datamodels/taskinstance.py | 3 + .../execution_api/routes/task_instances.py | 10 ++++ .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_06_30.py | 10 ++++ .../versions/head/test_task_instances.py | 54 +++++++++++++++++ .../v2026_06_30/test_task_instances.py | 58 ++++++++++++++++++- .../airflow/sdk/api/datamodels/_generated.py | 3 + .../sdk/execution_time/schema/schema.json | 13 +++++ .../airflow/sdk/execution_time/task_runner.py | 6 ++ task-sdk/tests/conftest.py | 6 ++ .../execution_time/test_task_runner.py | 18 +++++- 11 files changed, 180 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ad051b3e6d340..4117b06f7f06a 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -402,6 +402,9 @@ class TIRunContext(BaseModel): task_reschedule_count: int = 0 """How many times the task has been rescheduled.""" + first_task_reschedule_start_date: UtcDateTime | None = None + """The first reschedule start date for the task instance, if it has been rescheduled.""" + max_tries: int """Maximum number of tries for the task instance (from DB).""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 3b142a4e6b1b7..81bcf209e2d69 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -296,6 +296,14 @@ def ti_run( ) or 0 ) + first_task_reschedule_start_date = None + if task_reschedule_count > 0: + first_task_reschedule_start_date = session.scalar( + select(TaskReschedule.start_date) + .where(TaskReschedule.ti_id == task_instance_id) + .order_by(TaskReschedule.id.asc()) + .limit(1) + ) dr.team_name = get_team_name_for_ti(task_instance_id, session) @@ -309,6 +317,8 @@ def ti_run( xcom_keys_to_clear=xcom_keys, should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) + if first_task_reschedule_start_date is not None: + context.first_task_reschedule_start_date = first_task_reschedule_start_date # Only set if they are non-null if ti.next_method: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index dc7035d31e3c9..5a6fb793a7c85 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -44,6 +44,7 @@ AddAssetsByAliasEndpoint, AddAwaitingInputStatePayload, AddConnectionTestEndpoint, + AddFirstTaskRescheduleStartDateField, AddPartitionDateField, AddRetryPolicyFields, AddTaskAndAssetStateStoreEndpoints, @@ -56,6 +57,7 @@ HeadVersion(), Version( "2026-06-30", + AddFirstTaskRescheduleStartDateField, AddVariableKeysEndpoint, AddConnectionTestEndpoint, AddAwaitingInputStatePayload, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py index e89e2ed04cc5d..2deecc8995fe3 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py @@ -53,6 +53,16 @@ class AddConnectionTestEndpoint(VersionChange): ) +class AddFirstTaskRescheduleStartDateField(VersionChange): + """Add first_task_reschedule_start_date field to TIRunContext.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("first_task_reschedule_start_date").didnt_exist, + ) + + class AddTaskInstanceQueueField(VersionChange): """Add the `queue` field to the TaskInstance model.""" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 4920fc001f9d8..39f632d5d11a3 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -322,6 +322,60 @@ def test_ti_run_state_to_running( ) assert response.status_code == 409 + def test_ti_run_state_includes_first_task_reschedule_start_date( + self, + client, + session, + create_task_instance, + ): + """Test that running a rescheduled Task Instance includes its first reschedule start date.""" + instant_str = "2024-09-30T12:00:00Z" + instant = timezone.parse(instant_str) + first_reschedule_start_date = timezone.datetime(2024, 9, 30, 10) + second_reschedule_start_date = timezone.datetime(2024, 9, 30, 11) + + ti = create_task_instance( + task_id="test_ti_run_state_includes_first_task_reschedule_start_date", + state=State.QUEUED, + dagrun_state=DagRunState.RUNNING, + session=session, + start_date=instant, + dag_id=str(uuid4()), + ) + session.add_all( + [ + TaskReschedule( + ti_id=ti.id, + start_date=first_reschedule_start_date, + end_date=timezone.datetime(2024, 9, 30, 10, 1), + reschedule_date=timezone.datetime(2024, 9, 30, 10, 2), + ), + TaskReschedule( + ti_id=ti.id, + start_date=second_reschedule_start_date, + end_date=timezone.datetime(2024, 9, 30, 11, 1), + reschedule_date=timezone.datetime(2024, 9, 30, 11, 2), + ), + ] + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": instant_str, + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result["task_reschedule_count"] == 2 + assert result["first_task_reschedule_start_date"] == "2024-09-30T10:00:00Z" + def test_ti_run_returns_execution_token( self, client, exec_app, session, create_task_instance, time_machine ): diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py index b2cc016be84d7..f273276befe85 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_06_30/test_task_instances.py @@ -17,12 +17,15 @@ from __future__ import annotations +from uuid import uuid4 + import pytest from airflow._shared.timezones import timezone +from airflow.models import TaskReschedule from airflow.utils.state import DagRunState, State -from tests_common.test_utils.db import clear_db_runs +from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_serialized_dags pytestmark = pytest.mark.db_test @@ -41,11 +44,62 @@ @pytest.fixture def old_ver_client(client): - """Execution API version immediately before ``partition_date`` was added.""" + """Execution API version immediately before ``partition_date`` and ``first_task_reschedule_start_date`` were added.""" client.headers["Airflow-API-Version"] = "2026-06-16" return client +@pytest.fixture(autouse=True) +def setup_teardown(): + clear_db_runs() + clear_db_serialized_dags() + clear_db_dags() + yield + clear_db_runs() + clear_db_serialized_dags() + clear_db_dags() + + +def test_first_task_reschedule_start_date_removed_from_previous_version( + old_ver_client, + session, + create_task_instance, +): + ti = create_task_instance( + task_id="test_first_task_reschedule_start_date_removed_from_previous_version", + state=State.QUEUED, + dagrun_state=DagRunState.RUNNING, + session=session, + start_date=timezone.datetime(2024, 9, 30, 12), + dag_id=str(uuid4()), + ) + session.add( + TaskReschedule( + ti_id=ti.id, + start_date=timezone.datetime(2024, 9, 30, 10), + end_date=timezone.datetime(2024, 9, 30, 10, 1), + reschedule_date=timezone.datetime(2024, 9, 30, 10, 2), + ) + ) + session.commit() + + response = old_ver_client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + }, + ) + + assert response.status_code == 200 + result = response.json() + assert result["task_reschedule_count"] == 1 + assert "first_task_reschedule_start_date" not in result + + class TestPartitionDateFieldBackwardCompat: @pytest.fixture(autouse=True) def _freeze_time(self, time_machine): diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 72aebeab3e49b..b7a8befd3df86 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -789,6 +789,9 @@ class TIRunContext(BaseModel): dag_run: DagRun task_reschedule_count: Annotated[int | None, Field(title="Task Reschedule Count")] = 0 + first_task_reschedule_start_date: Annotated[ + AwareDatetime | None, Field(title="First Task Reschedule Start Date") + ] = None max_tries: Annotated[int, Field(title="Max Tries")] variables: Annotated[list[VariableResponse] | None, Field(title="Variables")] = None connections: Annotated[list[ConnectionResponse] | None, Field(title="Connections")] = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index e6ce8aa3d066e..847fb61352cd7 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4739,6 +4739,19 @@ "title": "Task Reschedule Count", "type": "integer" }, + "first_task_reschedule_start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "First Task Reschedule Start Date" + }, "max_tries": { "title": "Max Tries", "type": "integer" 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 b6616dc30cfd7..0b139222ba7ed 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -668,6 +668,12 @@ def get_first_reschedule_date(self, context: Context) -> AwareDatetime | None: # If the task has not been rescheduled, there is no need to ask the supervisor return None + first_task_reschedule_start_date = getattr( + self._ti_context_from_server, "first_task_reschedule_start_date", None + ) + if first_task_reschedule_start_date is not None: + return first_task_reschedule_start_date + max_tries: int = self.max_tries retries: int = self.task.retries or 0 first_try_number = max_tries - retries + 1 diff --git a/task-sdk/tests/conftest.py b/task-sdk/tests/conftest.py index c1ef3b72c92f4..321dbc7c349de 100644 --- a/task-sdk/tests/conftest.py +++ b/task-sdk/tests/conftest.py @@ -229,6 +229,7 @@ def __call__( run_after: str | datetime = ..., run_type: str = ..., task_reschedule_count: int = ..., + first_task_reschedule_start_date: str | datetime | None = ..., conf: dict[str, Any] | None = ..., should_retry: bool = ..., max_tries: int = ..., @@ -249,6 +250,7 @@ def __call__( run_after: str | datetime = ..., run_type: str = ..., task_reschedule_count: int = ..., + first_task_reschedule_start_date: str | datetime | None = ..., conf=None, consumed_asset_events: Sequence[AssetEventDagRunReference] = ..., ) -> dict[str, Any]: ... @@ -271,6 +273,7 @@ def _make_context( run_after: str | datetime = "2024-12-01T01:00:00Z", run_type: str = "manual", task_reschedule_count: int = 0, + first_task_reschedule_start_date: str | datetime | None = None, conf: dict[str, Any] | None = None, should_retry: bool = False, max_tries: int = 0, @@ -292,6 +295,7 @@ def _make_context( consumed_asset_events=list(consumed_asset_events), ), task_reschedule_count=task_reschedule_count, + first_task_reschedule_start_date=first_task_reschedule_start_date, max_tries=max_tries, should_retry=should_retry, ) @@ -314,6 +318,7 @@ def _make_context_dict( run_after: str | datetime = "2024-12-01T00:00:00Z", run_type: str = "manual", task_reschedule_count: int = 0, + first_task_reschedule_start_date: str | datetime | None = None, conf=None, consumed_asset_events: Sequence[AssetEventDagRunReference] = (), ) -> dict[str, Any]: @@ -329,6 +334,7 @@ def _make_context_dict( run_type=run_type, conf=conf, task_reschedule_count=task_reschedule_count, + first_task_reschedule_start_date=first_task_reschedule_start_date, consumed_asset_events=consumed_asset_events, ) return context.model_dump(exclude_unset=True, mode="json") diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 1aad1a02d5eab..28f8c1663388b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2763,7 +2763,7 @@ def __init__(self, command, *args, **kwargs): def test_get_first_reschedule_date( self, create_runtime_ti, mock_supervisor_comms, task_reschedule_count, expected_date ): - """Test that the first reschedule date is fetched from the Supervisor.""" + """Test that the first reschedule date falls back to the Supervisor.""" task = BaseOperator(task_id="hello") runtime_ti = create_runtime_ti(task=task, task_reschedule_count=task_reschedule_count) @@ -2774,6 +2774,22 @@ def test_get_first_reschedule_date( context = runtime_ti.get_template_context() assert runtime_ti.get_first_reschedule_date(context=context) == expected_date + def test_get_first_reschedule_date_uses_context_from_server( + self, create_runtime_ti, make_ti_context, mock_supervisor_comms + ): + """Test that first reschedule date from server context avoids a Supervisor request.""" + first_reschedule_date = timezone.datetime(2025, 1, 1) + task = BaseOperator(task_id="hello") + runtime_ti = create_runtime_ti(task=task, task_reschedule_count=1) + runtime_ti._ti_context_from_server = make_ti_context( + task_reschedule_count=1, + first_task_reschedule_start_date=first_reschedule_date, + ) + + context = runtime_ti.get_template_context() + assert runtime_ti.get_first_reschedule_date(context=context) == first_reschedule_date + mock_supervisor_comms.send.assert_not_called() + def test_get_ti_count(self, mock_supervisor_comms): """Test that get_ti_count sends the correct request and returns the count.""" mock_supervisor_comms.send.return_value = TICount(count=2) From 87f99bb803353e646fee4c6c470fd63b1a3d67af Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Fri, 3 Jul 2026 11:02:19 +0530 Subject: [PATCH 2/4] Address review feedback on PR #68012 Fix the CI failures and review comments from PR #68012: - Remove the dead module-level VersionBundle construction in schema/versions/__init__.py that referenced names only imported under TYPE_CHECKING, breaking the lazy cadwyn import on any access to the module and causing NameError in ~30 task-sdk tests. - Remove a duplicated old_ver_client fixture and import block that had been pasted into test_task_instances.py, which triggered mypy no-redef and ruff F811/E402 failures in Static checks. - Use func.min(start_date) instead of ordering by TaskReschedule.id when computing the first reschedule start date, since insertion order does not guarantee chronological order. --- .../api_fastapi/execution_api/routes/task_instances.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 81bcf209e2d69..10a1a0d139b3d 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -299,10 +299,7 @@ def ti_run( first_task_reschedule_start_date = None if task_reschedule_count > 0: first_task_reschedule_start_date = session.scalar( - select(TaskReschedule.start_date) - .where(TaskReschedule.ti_id == task_instance_id) - .order_by(TaskReschedule.id.asc()) - .limit(1) + select(func.min(TaskReschedule.start_date)).where(TaskReschedule.ti_id == task_instance_id) ) dr.team_name = get_team_name_for_ti(task_instance_id, session) From fea9229db772cd362b514a53fea494ed10db40ae Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Fri, 3 Jul 2026 13:12:01 +0530 Subject: [PATCH 3/4] Regenerate ts-sdk supervisor schema for first_task_reschedule_start_date A new prek hook (check-ts-sdk-supervisor-schema) landed on main after this branch was created, requiring the TypeScript SDK's generated supervisor types to stay in sync with the Python schema.json. Run `pnpm run generate:supervisor` to add the missing first_task_reschedule_start_date field to ts-sdk/src/generated/supervisor.ts. --- ts-sdk/src/generated/supervisor.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index fe61155f39e4c..309ab2a43bb69 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -216,6 +216,7 @@ export type TaskInstanceState = | "deferred" | "awaiting_input"; export type TaskRescheduleCount = number; +export type FirstTaskRescheduleStartDate = string | null; export type MaxTries = number; export type Key = string; export type Value = string | null; @@ -974,6 +975,7 @@ export interface TaskCallbackRequest { export interface TIRunContext { dag_run: DagRun; task_reschedule_count?: TaskRescheduleCount; + first_task_reschedule_start_date?: FirstTaskRescheduleStartDate; max_tries: MaxTries; variables?: Variables; connections?: Connections; From 1b6042916e395107c67225610b86630eafbf8098 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Sun, 5 Jul 2026 11:19:50 +0530 Subject: [PATCH 4/4] Reduce reschedule startup queries in execution API The run endpoint was doing two TaskReschedule round trips to compute values that can be fetched together. Collapsing them into one aggregate query keeps the rescheduled-sensor context fix intact while trimming overhead on task startup, and the test now proves we return the earliest reschedule timestamp instead of depending on insertion order. --- .../execution_api/routes/task_instances.py | 17 ++++++----------- .../versions/head/test_task_instances.py | 12 ++++++------ 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 10a1a0d139b3d..38ff851322353 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -290,17 +290,12 @@ def ti_run( xcom_query = xcom_query.where(XComModel.map_index == map_index) xcom_keys = list(session.scalars(xcom_query)) - task_reschedule_count = ( - session.scalar( - select(func.count(TaskReschedule.id)).where(TaskReschedule.ti_id == task_instance_id) - ) - or 0 - ) - first_task_reschedule_start_date = None - if task_reschedule_count > 0: - first_task_reschedule_start_date = session.scalar( - select(func.min(TaskReschedule.start_date)).where(TaskReschedule.ti_id == task_instance_id) - ) + task_reschedule_count, first_task_reschedule_start_date = session.execute( + select( + func.count(TaskReschedule.id), + func.min(TaskReschedule.start_date), + ).where(TaskReschedule.ti_id == task_instance_id) + ).one() dr.team_name = get_team_name_for_ti(task_instance_id, session) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 39f632d5d11a3..9a4e8fc2430dc 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -344,18 +344,18 @@ def test_ti_run_state_includes_first_task_reschedule_start_date( ) session.add_all( [ - TaskReschedule( - ti_id=ti.id, - start_date=first_reschedule_start_date, - end_date=timezone.datetime(2024, 9, 30, 10, 1), - reschedule_date=timezone.datetime(2024, 9, 30, 10, 2), - ), TaskReschedule( ti_id=ti.id, start_date=second_reschedule_start_date, end_date=timezone.datetime(2024, 9, 30, 11, 1), reschedule_date=timezone.datetime(2024, 9, 30, 11, 2), ), + TaskReschedule( + ti_id=ti.id, + start_date=first_reschedule_start_date, + end_date=timezone.datetime(2024, 9, 30, 10, 1), + reschedule_date=timezone.datetime(2024, 9, 30, 10, 2), + ), ] ) session.commit()