From a88e6b1b46b93291f8017c1705fd5c82d5d2edc4 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Sun, 23 Aug 2026 13:10:34 +0530 Subject: [PATCH 1/5] Include first task reschedule start date in TI run context Rescheduled tasks fetched their first task reschedule start date from the supervisor when RuntimeTaskInstance.get_first_reschedule_date() was called. For sensors this happens before poke() starts, so a worker blocked on the supervisor/API round trip before emitting the sensor's usual poke log line. The ti_run endpoint already queries TaskReschedule for the reschedule count; widen that to a single count + min(start_date) query and return the first reschedule start date in the run context. The Task SDK uses the value when present and keeps the supervisor request as a fallback for older servers. closes: #68010 --- .../execution_api/datamodels/taskinstance.py | 3 ++ .../execution_api/routes/task_instances.py | 14 ++--- .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_10_30.py | 10 ++++ .../versions/head/test_task_instances.py | 54 +++++++++++++++++++ .../v2026_10_30/test_task_instances.py | 45 +++++++++++++++- .../airflow/sdk/api/datamodels/_generated.py | 3 ++ .../sdk/execution_time/schema/schema.json | 13 +++++ .../schema/versions/__init__.py | 2 + .../schema/versions/v2026_10_30.py | 16 ++++++ .../airflow/sdk/execution_time/task_runner.py | 6 +++ task-sdk/tests/conftest.py | 6 +++ .../execution_time/test_task_runner.py | 18 ++++++- ts-sdk/src/generated/supervisor.ts | 2 + 14 files changed, 185 insertions(+), 9 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 2d158476ada6d..81000c67d305a 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 @@ -448,6 +448,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 cb42639c98461..3f462882ed91a 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 @@ -321,12 +321,12 @@ def ti_run( read = build_xcom_read_query(producer_ids=select(TI.id).where(TI.id == task_instance_id)) entity = xcom_entity(read) xcom_keys = list(session.scalars(read.with_only_columns(entity.key))) - task_reschedule_count = ( - session.scalar( - select(func.count(TaskReschedule.id)).where(TaskReschedule.ti_id == task_instance_id) - ) - or 0 - ) + 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() context = TIRunContext( dag_run=dr, @@ -352,6 +352,8 @@ def ti_run( terminal_task_id=group.terminal_task_id, terminal_is_mapped=terminal.get_needs_expansion(), ) + if first_task_reschedule_start_date is not None: + context.first_task_reschedule_start_date = first_task_reschedule_start_date # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg # spec; the route excludes unset fields, keeping regular responses lean. 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 f23972db0045d..aa4d1028d2fa3 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 @@ -56,6 +56,7 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddDagRunNoteUpdateEndpoint, + AddFirstTaskRescheduleStartDateField, AddLoopContext, AddMultiTeamToTIRunContext, AddStoppedTaskReport, @@ -81,6 +82,7 @@ AddStoppedTaskReport, AddTaskInstanceRegionCoordinates, IdentifyArchivedTaskStateUpdates, + AddFirstTaskRescheduleStartDateField, ), Version( "2026-06-30", diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index a75c7591982b2..ec4b6ddf12529 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -133,3 +133,13 @@ class AddMultiTeamToTIRunContext(VersionChange): def remove_multi_team_field(response: ResponseInfo) -> None: # type: ignore[misc] """Strip ``multi_team`` from the run context for older clients.""" response.body.pop("multi_team", None) + + +class AddFirstTaskRescheduleStartDateField(VersionChange): + """Add the ``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, + ) 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 584b112f213ff..983465d2296f2 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 @@ -877,6 +877,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=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() + + 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_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index 55280830e1496..3b4fcd392aa8b 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -17,6 +17,8 @@ from __future__ import annotations +from uuid import uuid4 + import pytest from sqlalchemy import select @@ -28,7 +30,7 @@ from airflow.models.trigger import Trigger from airflow.models.xcom import XComModelV2 from airflow.sdk import task -from airflow.utils.state import State +from airflow.utils.state import DagRunState, State from tests_common.test_utils.db import clear_db_runs @@ -224,7 +226,7 @@ def test_legacy_worker_heartbeat_rejects_restarting_task(client, session, create @pytest.fixture def old_ver_client(client): - """Execution API version immediately before ``arg_bindings`` was added.""" + """Execution API version immediately before the ``2026-10-30`` fields were added.""" client.headers["Airflow-API-Version"] = "2026-06-30" return client @@ -329,3 +331,42 @@ def test_mutations_after_retry_response_by_version( stored = XComModelV2.get_for_attempt(ti.id, "extra_link", session=session) assert (stored is not None) is xcom_kept assert XComModelV2.get_for_attempt(successor.id, "extra_link", session=session) is None + + +class TestFirstTaskRescheduleStartDateBackwardCompat: + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + def test_old_version_strips_first_task_reschedule_start_date( + self, + old_ver_client, + session, + create_task_instance, + ): + ti = create_task_instance( + task_id="test_first_task_reschedule_start_date_compat", + 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=RUN_PATCH_BODY) + + assert response.status_code == 200 + result = response.json() + assert result["task_reschedule_count"] == 1 + assert "first_task_reschedule_start_date" not in result diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index d6e9e45f0d24c..b534358814897 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -856,6 +856,9 @@ class TIRunContext(BaseModel): dag_run: DagRun loop: LoopContext | None = None 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 7e09a736de793..77888c5adfdea 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -5226,6 +5226,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/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 647e5dc919317..0f35e9f0e6f2b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -40,6 +40,7 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, AddDagDefinitionsToDagFileParsingResult, + AddFirstTaskRescheduleStartDateToSupervisorTIRunContext, AddLoopRuntime, AddRegionCoordinatesToSupervisorTaskInstance, AddRetryReasonToTaskState, @@ -54,6 +55,7 @@ def get_bundle() -> VersionBundle: AddRetryReasonToTaskState, AddDagDefinitionsToDagFileParsingResult, AddRegionCoordinatesToSupervisorTaskInstance, + AddFirstTaskRescheduleStartDateToSupervisorTIRunContext, ), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index baffc8c9d0f1b..d1efcf0b8c380 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -95,3 +95,19 @@ class AddDagDefinitionsToDagFileParsingResult(VersionChange): schema(DagFileParsingResult).field("parsed_definitions").didnt_exist, schema(DagFileParsingResult).field("dag_source_codes").didnt_exist, ) + + +class AddFirstTaskRescheduleStartDateToSupervisorTIRunContext(VersionChange): + """ + Add the ``first_task_reschedule_start_date`` field carrying the earliest reschedule start date. + + The supervisor-schema mirror of the execution API's + ``AddFirstTaskRescheduleStartDateField``, named apart so the two migrations + are not confused. + """ + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("first_task_reschedule_start_date").didnt_exist, + ) 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 652225b6673aa..88485e8c0f639 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -730,6 +730,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 + if from_server := self._ti_context_from_server: + # Servers older than the 2026-10-30 API version omit this, so fall through to + # asking the supervisor below. + if (first_reschedule_date := from_server.first_task_reschedule_start_date) is not None: + return first_reschedule_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 a7d14ef1b54fb..1fee940436cb5 100644 --- a/task-sdk/tests/conftest.py +++ b/task-sdk/tests/conftest.py @@ -239,6 +239,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 = ..., @@ -259,6 +260,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]: ... @@ -281,6 +283,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, @@ -304,6 +307,7 @@ def _make_context( partition_key=None, ), task_reschedule_count=task_reschedule_count, + first_task_reschedule_start_date=first_task_reschedule_start_date, max_tries=max_tries, should_retry=should_retry, ) @@ -326,6 +330,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]: @@ -341,6 +346,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 8d4ade1a3401b..2cd63cce594c7 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 @@ -3594,7 +3594,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) @@ -3605,6 +3605,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) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 31196f868a5cb..7faf8c4877029 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -237,6 +237,7 @@ export type MaxIterations = number; export type TerminalTaskId = string; export type TerminalIsMapped = boolean; export type TaskRescheduleCount = number; +export type FirstTaskRescheduleStartDate = string | null; export type MaxTries = number; export type Key = string; export type Value = string | null; @@ -1049,6 +1050,7 @@ export interface TIRunContext { dag_run: DagRun; loop?: LoopContext | null; task_reschedule_count?: TaskRescheduleCount; + first_task_reschedule_start_date?: FirstTaskRescheduleStartDate; max_tries: MaxTries; variables?: Variables; connections?: Connections; From c938c9f582b7f4cb8097765872cee3c4166c7833 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Sun, 4 Oct 2026 18:10:40 +0530 Subject: [PATCH 2/5] Use min(start_date) for the task-reschedules start_date route ti_run now returns the earliest reschedule start date via min(start_date). Use the same definition in GET /task-reschedules/{ti_id}/start_date instead of picking the first row by id, so the two cannot drift apart. Co-Authored-By: Claude --- .../execution_api/routes/task_reschedules.py | 11 +++----- .../versions/head/test_task_instances.py | 26 ++++++++++++++++++- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py index c98d9417f509b..02b7810681bdb 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py @@ -20,7 +20,7 @@ from uuid import UUID from fastapi import APIRouter, Security, status -from sqlalchemy import select +from sqlalchemy import func, select from airflow.api_fastapi.common.db.common import AsyncSessionDep from airflow.api_fastapi.common.types import UtcDateTime @@ -42,11 +42,6 @@ @router.get("/{task_instance_id}/start_date") async def get_start_date(task_instance_id: UUID, session: AsyncSessionDep) -> UtcDateTime | None: """Get the first reschedule date if found, None if no records exist.""" - start_date = await session.scalar( - select(TaskReschedule.start_date) - .where(TaskReschedule.ti_id == task_instance_id) - .order_by(TaskReschedule.id.asc()) - .limit(1) + return await session.scalar( + select(func.min(TaskReschedule.start_date)).where(TaskReschedule.ti_id == task_instance_id) ) - - return start_date 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 983465d2296f2..40b7f029a50dc 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 @@ -21,7 +21,7 @@ import time from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext -from datetime import datetime +from datetime import datetime, timedelta from threading import Barrier, Event, local from types import SimpleNamespace from typing import TYPE_CHECKING @@ -5308,6 +5308,30 @@ def test_get_start_date(self, client, session, create_task_instance): assert response.status_code == 200 assert response.json() == "2024-01-01T00:00:00Z" + def test_get_start_date_returns_earliest_not_first_inserted(self, client, session, create_task_instance): + ti = create_task_instance( + task_id="test_get_start_date_returns_earliest_not_first_inserted", + state=State.RUNNING, + start_date=timezone.datetime(2024, 1, 1), + session=session, + ) + # Insert the later reschedule first so that ordering by id and by start_date disagree + for start_date in (timezone.datetime(2024, 1, 2), timezone.datetime(2024, 1, 1)): + session.add( + TaskReschedule( + ti_id=ti.id, + start_date=start_date, + end_date=start_date + timedelta(hours=1), + reschedule_date=start_date + timedelta(hours=2), + ) + ) + session.flush() + session.commit() + + response = client.get(f"/execution/task-reschedules/{ti.id}/start_date") + assert response.status_code == 200 + assert response.json() == "2024-01-01T00:00:00Z" + def test_get_start_date_not_found(self, client): ti_id = "0182e924-0f1e-77e6-ab50-e977118bc139" response = client.get(f"/execution/task-reschedules/{ti_id}/start_date") From 21a507e085a1955de0b746d3be7b148fe1c1eaef Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Tue, 6 Oct 2026 00:23:03 +0200 Subject: [PATCH 3/5] Sync the Go SDK supervisor schema and models with the new run-context field The Task SDK supervisor schema gained first_task_reschedule_start_date, and the Go SDK keeps a vendored copy of that schema plus models generated from it, so the static checks failed until both were refreshed. Generated-by: Claude Opus 5 --- go-sdk/pkg/execution/genmodels/models.gen.go | 142 ++++++++++--------- go-sdk/schema/supervisor-schema.json | 13 ++ 2 files changed, 86 insertions(+), 69 deletions(-) diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index d7d779b15cf5f..6e53729efe936 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1684,6 +1684,10 @@ type TIRunContext struct { // DagRun corresponds to the JSON schema field "dag_run". DagRun DagRun `msgpack:"dag_run"` + // FirstTaskRescheduleStartDate corresponds to the JSON schema field + // "first_task_reschedule_start_date". + FirstTaskRescheduleStartDate interface{} `msgpack:"first_task_reschedule_start_date,omitempty"` + // Loop corresponds to the JSON schema field "loop". Loop *LoopContext `msgpack:"loop,omitempty"` @@ -1716,51 +1720,19 @@ type TIRunContext struct { XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } -type TaskArgBinding interface{} - -type TaskBreadcrumbsResult struct { - // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". - Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} - -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` - - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` - - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` - - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` +type TriggerKwargs map[string]interface{} - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` +type Warnings []interface{} - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` +type TaskArgBinding interface{} - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` - // VersionData corresponds to the JSON schema field "version_data". - VersionData *VersionData `msgpack:"version_data,omitempty"` + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` } type TaskIds []string @@ -1804,24 +1776,69 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } -type TaskInstanceState string - const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRemoved TaskInstanceState = "removed" const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateSuccess TaskInstanceState = "success" const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateQueued TaskInstanceState = "queued" const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateRemoved TaskInstanceState = "removed" + +type TaskInstanceState string type TaskOutlets []AssetProfile +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} + +type TaskBreadcrumbsResult struct { + // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". + Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` +} + +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` + + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` + + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` + + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` + + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` + + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` +} + // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { // StartDate corresponds to the JSON schema field "start_date". @@ -1831,6 +1848,12 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateSkipped TaskStateState = "skipped" +const TaskStateStateRemoved TaskStateState = "removed" + // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1854,12 +1877,6 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateRemoved TaskStateState = "removed" -const TaskStateStateSkipped TaskStateState = "skipped" - // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1880,19 +1897,6 @@ type TaskStatesResult struct { Type string `msgpack:"type,omitempty"` } -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - -type Warnings []interface{} - -type TriggerKwargs map[string]interface{} - type TriggerDagRun struct { // Conf corresponds to the JSON schema field "conf". Conf *Conf `msgpack:"conf,omitempty"` diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 7e09a736de793..77888c5adfdea 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -5226,6 +5226,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" From c8527f483f2d54adafe4e147e8d2ccb72b61fa85 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Tue, 6 Oct 2026 10:50:19 +0530 Subject: [PATCH 4/5] Add real-bundle downgrade test for the first reschedule start date field AddFirstTaskRescheduleStartDateToSupervisorTIRunContext is didnt_exist-only, so it is only visible through SchemaVersionMigrator.downgrade. Pin it the same way the arg_bindings and retry_reason changes are pinned. Co-Authored-By: Claude --- .../execution_time/schema/test_migrator.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index a5ac1de22d286..802764446df49 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -28,6 +28,7 @@ from __future__ import annotations +import datetime import subprocess import sys from typing import Literal @@ -580,3 +581,75 @@ def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): body = {"type": "TaskState", "state": "failed", "end_date": None, "rendered_map_index": None} out = real_migrator.upgrade(body, TaskState, "2026-06-16") assert out["retry_reason"] is None + + +class TestRealBundleFirstTaskRescheduleStartDate: + """ + Drive the *real* supervisor bundle through the ``first_task_reschedule_start_date`` migration. + + ``AddFirstTaskRescheduleStartDateToSupervisorTIRunContext`` is ``didnt_exist``-only, so the + downgrade re-validation is the only place it is visible: dropping it from the bundle, or + registering it under the wrong version, leaves the field on the wire for a pinned runtime. + """ + + FIRST_RESCHEDULE = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + + @pytest.fixture + def startup_details(self): + import uuid + + from airflow.sdk.api.datamodels._generated import ( + BundleInfo, + DagRun, + DagRunState, + DagRunType, + TaskInstance, + TIRunContext, + ) + from airflow.sdk.execution_time.comms import StartupDetails + + now = datetime.datetime.now(datetime.timezone.utc) + return StartupDetails( + ti=TaskInstance( + id=uuid.uuid4(), + task_id="wait", + dag_id="d", + run_id="r", + try_number=1, + dag_version_id=uuid.uuid4(), + ), + dag_rel_path="d.py", + bundle_info=BundleInfo(name="b", version=None), + start_date=now, + ti_context=TIRunContext( + dag_run=DagRun( + dag_id="d", + run_id="r", + logical_date=now, + data_interval_start=None, + data_interval_end=None, + start_date=now, + end_date=None, + run_type=DagRunType.MANUAL, + state=DagRunState.RUNNING, + run_after=now, + consumed_asset_events=[], + partition_key=None, + ), + max_tries=1, + first_task_reschedule_start_date=self.FIRST_RESCHEDULE, + ), + sentry_integration="", + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_field_for_previous_version(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() + assert "first_task_reschedule_start_date" not in out["ti_context"] + + def test_head_version_keeps_field(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-10-30") + assert out.ti_context.first_task_reschedule_start_date == self.FIRST_RESCHEDULE From 4d2464c5465e0b1e4f4328c0fb1b846278c16ff4 Mon Sep 17 00:00:00 2001 From: Hemkumar Chheda Date: Sat, 10 Oct 2026 23:35:02 +0530 Subject: [PATCH 5/5] Address review nits and sync the TS SDK vendored supervisor schema Drop the API date from the older-server comment, move the migrator test imports to module level, and add first_task_reschedule_start_date to the TS SDK's vendored supervisor schema that #74415 introduced. Co-Authored-By: Claude --- .../airflow/sdk/execution_time/task_runner.py | 3 +- .../execution_time/schema/test_migrator.py | 29 +++++++++---------- ts-sdk/schema/supervisor-schema.json | 13 +++++++++ 3 files changed, 28 insertions(+), 17 deletions(-) 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 88485e8c0f639..b90ec01e21e99 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -731,8 +731,7 @@ def get_first_reschedule_date(self, context: Context) -> AwareDatetime | None: return None if from_server := self._ti_context_from_server: - # Servers older than the 2026-10-30 API version omit this, so fall through to - # asking the supervisor below. + # Older servers don't send this field, so fall through to asking the supervisor below. if (first_reschedule_date := from_server.first_task_reschedule_start_date) is not None: return first_reschedule_date diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 802764446df49..e8c4a16aff248 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -31,6 +31,7 @@ import datetime import subprocess import sys +import uuid from typing import Literal from uuid import UUID @@ -51,13 +52,23 @@ ) from airflow.sdk import TaskInstanceState -from airflow.sdk.api.datamodels._generated import LoopContext, PreviousTIResponse +from airflow.sdk.api.datamodels._generated import ( + BundleInfo, + DagRun, + DagRunState, + DagRunType, + LoopContext, + PreviousTIResponse, + TaskInstance, + TIRunContext, +) from airflow.sdk.execution_time.comms import ( GetXCom, GetXComCount, GetXComSequenceItem, GetXComSequenceSlice, PreviousTIResult, + StartupDetails, TaskState, ) from airflow.sdk.execution_time.schema import ( @@ -592,23 +603,11 @@ class TestRealBundleFirstTaskRescheduleStartDate: registering it under the wrong version, leaves the field on the wire for a pinned runtime. """ - FIRST_RESCHEDULE = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + FIRST_RESCHEDULE = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC) @pytest.fixture def startup_details(self): - import uuid - - from airflow.sdk.api.datamodels._generated import ( - BundleInfo, - DagRun, - DagRunState, - DagRunType, - TaskInstance, - TIRunContext, - ) - from airflow.sdk.execution_time.comms import StartupDetails - - now = datetime.datetime.now(datetime.timezone.utc) + now = datetime.datetime.now(datetime.UTC) return StartupDetails( ti=TaskInstance( id=uuid.uuid4(), diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index 7e09a736de793..77888c5adfdea 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -5226,6 +5226,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"