This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3af85e214f4 Avoid a supervisor round trip before poke for rescheduled
sensors (#68012)
3af85e214f4 is described below
commit 3af85e214f489748fd39dab9bbbfb99073ffc3b5
Author: Hemkumar Chheda <[email protected]>
AuthorDate: Mon Oct 12 18:12:56 2026 +0530
Avoid a supervisor round trip before poke for rescheduled sensors (#68012)
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
* 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.
* 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
* 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.
---
.../execution_api/datamodels/taskinstance.py | 3 +
.../execution_api/routes/task_instances.py | 14 +-
.../execution_api/routes/task_reschedules.py | 11 +-
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_10_30.py | 10 ++
.../versions/head/test_task_instances.py | 80 +++++++++++-
.../versions/v2026_10_30/test_task_instances.py | 45 ++++++-
go-sdk/pkg/execution/genmodels/models.gen.go | 142 +++++++++++----------
go-sdk/schema/supervisor-schema.json | 13 ++
.../src/airflow/sdk/api/datamodels/_generated.py | 3 +
.../airflow/sdk/execution_time/schema/schema.json | 13 ++
.../sdk/execution_time/schema/versions/__init__.py | 2 +
.../execution_time/schema/versions/v2026_10_30.py | 16 +++
.../src/airflow/sdk/execution_time/task_runner.py | 5 +
task-sdk/tests/conftest.py | 6 +
.../execution_time/schema/test_migrator.py | 74 ++++++++++-
.../task_sdk/execution_time/test_task_runner.py | 18 ++-
ts-sdk/schema/supervisor-schema.json | 13 ++
ts-sdk/src/generated/supervisor.ts | 2 +
19 files changed, 384 insertions(+), 88 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 2d158476ada..81000c67d30 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 cb42639c984..3f462882ed9 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/routes/task_reschedules.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_reschedules.py
index c98d9417f50..02b7810681b 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 __future__ import annotations
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 = APIRouter(
@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/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index f23972db004..aa4d1028d2f 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 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
AddArgBindingsToTIRunContext,
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
+ AddFirstTaskRescheduleStartDateField,
AddLoopContext,
AddMultiTeamToTIRunContext,
AddStoppedTaskReport,
@@ -81,6 +82,7 @@ bundle = VersionBundle(
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 a75c7591982..ec4b6ddf125 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 584b112f213..40b7f029a50 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 re
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
@@ -877,6 +877,60 @@ class TestTIRunState:
)
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
):
@@ -5254,6 +5308,30 @@ class TestGetRescheduleStartDate:
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")
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 55280830e14..3b4fcd392aa 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.taskreschedule import TaskReschedule
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 @@ RUN_PATCH_BODY = {
@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/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index d7d779b15cf..6e53729efe9 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 7e09a736de7..77888c5adfd 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"
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index d6e9e45f0d2..b5343588148 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 7e09a736de7..77888c5adfd 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 647e5dc9193..0f35e9f0e6f 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 baffc8c9d0f..d1efcf0b8c3 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 652225b6673..b90ec01e21e 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,11 @@ class RuntimeTaskInstance(TaskInstance):
# 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:
+ # 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
+
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 a7d14ef1b54..1fee940436c 100644
--- a/task-sdk/tests/conftest.py
+++ b/task-sdk/tests/conftest.py
@@ -239,6 +239,7 @@ class MakeTIContextCallable(Protocol):
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 @@ class MakeTIContextDictCallable(Protocol):
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_ti_context() -> MakeTIContextCallable:
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_ti_context() -> MakeTIContextCallable:
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_ti_context_dict(make_ti_context:
MakeTIContextCallable) -> MakeTIContex
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_ti_context_dict(make_ti_context:
MakeTIContextCallable) -> MakeTIContex
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/schema/test_migrator.py
b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py
index a5ac1de22d2..e8c4a16aff2 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,8 +28,10 @@ schema diverges from head.
from __future__ import annotations
+import datetime
import subprocess
import sys
+import uuid
from typing import Literal
from uuid import UUID
@@ -50,13 +52,23 @@ from task_sdk.execution_time.schema._mock_version_bundle
import (
)
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 (
@@ -580,3 +592,63 @@ class TestRealBundleRetryReason:
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.UTC)
+
+ @pytest.fixture
+ def startup_details(self):
+ now = datetime.datetime.now(datetime.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
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 8d4ade1a340..2cd63cce594 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 @@ class TestRuntimeTaskInstance:
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 @@ class TestRuntimeTaskInstance:
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/schema/supervisor-schema.json
b/ts-sdk/schema/supervisor-schema.json
index 7e09a736de7..77888c5adfd 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"
diff --git a/ts-sdk/src/generated/supervisor.ts
b/ts-sdk/src/generated/supervisor.ts
index 31196f868a5..7faf8c48770 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;
