#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""The entrypoint for the actual task execution process."""

from __future__ import annotations

import contextvars
import functools
import inspect
import os
import sys
import time
from collections.abc import Callable, Iterable, Iterator, Mapping
from contextlib import ExitStack, contextmanager, suppress
from datetime import datetime, timedelta, timezone
from itertools import product
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any, Literal
from urllib.parse import quote

import attrs
import lazy_object_proxy
import structlog
from opentelemetry import trace
from opentelemetry.trace import INVALID_SPAN, Status, StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter
from structlog.contextvars import bind_contextvars

from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock
from airflow.dag_processing.bundles.manager import DagBundlesManager
from airflow.sdk._shared.observability.metrics import stats
from airflow.sdk._shared.observability.traces import get_task_span_detail_level
from airflow.sdk._shared.template_rendering import truncate_rendered_value
from airflow.sdk.api.client import get_hostname, getuser
from airflow.sdk.api.datamodels._generated import (
    AssetProfile,
    DagRun,
    PreviousTIResponse,
    TaskInstance,
    TaskInstanceState,
    TIRunContext,
)
from airflow.sdk.bases.operator import BaseOperator, ExecutorSafeguard
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.configuration import conf
from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
from airflow.sdk.definitions.asset import (
    Asset,
    AssetAlias,
    AssetNameRef,
    AssetUniqueKey,
    AssetUriRef,
)
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.param import process_params
from airflow.sdk.exceptions import (
    AirflowException,
    AirflowInactiveAssetInInletOrOutletException,
    AirflowRescheduleException,
    AirflowRuntimeError,
    AirflowTaskTimeout,
    ErrorType,
    TaskAwaitingInput,
    TaskDeferred,
)
from airflow.sdk.execution_time.callback_runner import create_executable_runner
from airflow.sdk.execution_time.comms import (
    AssetEventDagRunReferenceResult,
    AwaitInputTask,
    CommsDecoder,
    DagResult,
    DagRunStateResult,
    DeferTask,
    DRCount,
    ErrorResponse,
    GetDag,
    GetDagRunState,
    GetDRCount,
    GetPreviousDagRun,
    GetPreviousTI,
    GetTaskBreadcrumbs,
    GetTaskRescheduleStartDate,
    GetTaskStates,
    GetTICount,
    InactiveAssetsResult,
    PreviousDagRunResult,
    PreviousTIResult,
    RescheduleTask,
    ResendLoggingFD,
    RetryTask,
    SentFDs,
    SetRenderedFields,
    SetRenderedMapIndex,
    SkipDownstreamTasks,
    StartupDetails,
    SucceedTask,
    TaskBreadcrumbsResult,
    TaskRescheduleStartDate,
    TaskState,
    TaskStatesResult,
    TICount,
    ToSupervisor,
    ToTask,
    TriggerDagRun,
    ValidateInletsAndOutlets,
)
from airflow.sdk.execution_time.context import (
    AssetStateStoreAccessors,
    ConnectionAccessor,
    InletEventsAccessors,
    MacrosAccessor,
    OutletEventAccessors,
    TaskStateStoreAccessor,
    TriggeringAssetEventsAccessor,
    VariableAccessor,
    _get_worker_state_store_backend,
    context_get_outlet_events,
    context_to_airflow_vars,
    get_previous_dagrun_success,
    set_current_context,
)
from airflow.sdk.execution_time.sentry import Sentry
from airflow.sdk.execution_time.xcom import XCom
from airflow.sdk.listener import get_listener_manager
from airflow.sdk.observability.metrics import stats_utils
from airflow.sdk.serde import allow_class, iter_pydantic_models
from airflow.sdk.state import TaskScope
from airflow.sdk.timezone import coerce_datetime

if TYPE_CHECKING:
    import jinja2
    from pendulum.datetime import DateTime
    from structlog.typing import FilteringBoundLogger as Logger

    from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
    from airflow.sdk.definitions.context import Context
    from airflow.sdk.definitions.retry_policy import RetryDecision
    from airflow.sdk.exceptions import DagRunTriggerException
    from airflow.sdk.types import OutletEventAccessorsProtocol

log = structlog.get_logger("task")

tracer = trace.get_tracer(__name__)


class detail_span:
    """Context manager and decorator that creates a child span when detail level > 1."""

    def __init__(self, *args, **kwargs):
        self._args = args
        self._kwargs = kwargs
        self._ctx = None

    def _make_ctx(self):
        parent_span = trace.get_current_span()
        config_level = get_task_span_detail_level(span=parent_span)
        if config_level > 1:
            return tracer.start_as_current_span(*self._args, **self._kwargs)
        return trace.INVALID_SPAN

    def __enter__(self):
        self._ctx = self._make_ctx()
        return self._ctx.__enter__()

    def __exit__(self, *exc_info):
        return self._ctx.__exit__(*exc_info)

    def __call__(self, f):
        @functools.wraps(f)
        def wrapper(*inner_args, **inner_kwargs):
            with self._make_ctx():
                return f(*inner_args, **inner_kwargs)

        wrapper.__signature__ = inspect.signature(f)
        return wrapper


@contextmanager
def _make_task_span(msg: StartupDetails):
    parent_context = (
        TraceContextTextMapPropagator().extract(msg.ti.context_carrier) if msg.ti.context_carrier else None
    )
    ti = msg.ti
    span_name = f"worker.{ti.task_id}"
    if ti.map_index is not None and ti.map_index >= 0:
        span_name += f"[{ti.map_index}]"
    with tracer.start_as_current_span(span_name, context=parent_context) as span:
        span.set_attributes(
            {
                "airflow.dag_id": ti.dag_id,
                "airflow.task_id": ti.task_id,
                "airflow.dag_run.run_id": ti.run_id,
                "airflow.task_instance.try_number": ti.try_number,
                "airflow.task_instance.map_index": ti.map_index if ti.map_index is not None else -1,
            }
        )
        yield span


class TaskRunnerMarker:
    """Marker for listener hooks, to properly detect from which component they are called."""


# TODO: Move this entire class into a separate file:
#  `airflow/sdk/execution_time/task_instance.py`
#   or `airflow/sdk/execution_time/runtime_ti.py`
class RuntimeTaskInstance(TaskInstance):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    task: BaseOperator
    bundle_instance: BaseDagBundle
    _cached_template_context: Context | None = None
    """The Task Instance context. This is used to cache get_template_context."""

    _terminal_state_send_failed: bool = False
    """True when the supervisor IPC send for a non-success terminal state raised; signals main() to sys.exit(1) after finalize() so the supervisor doesn't misclassify the run as SUCCESS via exit code 0."""

    _ti_context_from_server: Annotated[TIRunContext | None, Field(repr=False)] = None
    """The Task Instance context from the API server, if any."""

    max_tries: int = 0
    """The maximum number of retries for the task."""

    start_date: AwareDatetime
    """Start date of the task instance."""

    end_date: AwareDatetime | None = None

    state: TaskInstanceState | None = None

    is_mapped: bool | None = None
    """True if the original task was mapped."""

    rendered_map_index: str | None = None

    sentry_integration: str = ""

    @property
    def stats_tags(self) -> dict[str, str]:
        """Metric tags for this task instance, including team_name when available."""
        tags: dict[str, str] = {"dag_id": self.dag_id, "task_id": self.task_id}
        if self._ti_context_from_server and self._ti_context_from_server.dag_run.team_name:
            tags["team_name"] = self._ti_context_from_server.dag_run.team_name
        return tags

    def __rich_repr__(self):
        yield "id", self.id
        yield "task_id", self.task_id
        yield "dag_id", self.dag_id
        yield "run_id", self.run_id
        yield "max_tries", self.max_tries
        yield "task", type(self.task)
        yield "start_date", self.start_date

    __rich_repr__.angular = True  # type: ignore[attr-defined]

    @detail_span("get_template_context")
    def get_template_context(self) -> Context:
        # TODO: Move this to `airflow.sdk.execution_time.context`
        #   once we port the entire context logic from airflow/utils/context.py ?
        from airflow.sdk.plugins_manager import integrate_macros_plugins

        integrate_macros_plugins()

        dag_run_conf: dict[str, Any] | None = None
        if from_server := self._ti_context_from_server:
            dag_run_conf = from_server.dag_run.conf or dag_run_conf

        validated_params = process_params(self.task.dag, self.task, dag_run_conf, suppress_exception=False)

        # Cache the context object, which ensures that all calls to get_template_context
        # are operating on the same context object.
        if self._cached_template_context is None:
            self._cached_template_context = {
                # From the Task Execution interface
                "dag": self.task.dag,
                "inlets": self.task.inlets,
                "map_index_template": self.task.map_index_template,
                "outlets": self.task.outlets,
                "run_id": self.run_id,
                "task": self.task,
                "task_instance": self,
                "ti": self,
                "outlet_events": OutletEventAccessors(),
                "inlet_events": InletEventsAccessors(self.task.inlets),
                "macros": MacrosAccessor(),
                "params": validated_params,
                # TODO: Make this go through Public API longer term.
                # "test_mode": task_instance.test_mode,
                "var": {
                    "json": VariableAccessor(deserialize_json=True),
                    "value": VariableAccessor(deserialize_json=False),
                },
                "conn": ConnectionAccessor(),
                "task_state_store": TaskStateStoreAccessor(
                    ti_id=self.id,
                    scope=TaskScope(
                        dag_id=self.dag_id,
                        run_id=self.run_id,
                        task_id=self.task_id,
                        map_index=self.map_index if self.map_index is not None else -1,
                    ),
                ),
            }
            _asset_types = (Asset, AssetNameRef, AssetUriRef, AssetAlias)
            if any(isinstance(i, _asset_types) for i in self.task.inlets + self.task.outlets):
                self._cached_template_context["asset_state_store"] = AssetStateStoreAccessors(
                    self.task.inlets, self.task.outlets
                )
                # AssetAlias inlets are resolved to their concrete assets at context build time via
                # GetAssetsByAlias comms. If an alias maps to no active assets, it doesn't contribute to
                # asset_state. AssetAlias outlets are skipped downstream in context.py
        if TYPE_CHECKING:
            assert self._cached_template_context is not None
        if from_server:
            dag_run = from_server.dag_run
            context_from_server: Context = {
                # TODO: Assess if we need to pass these through timezone.coerce_datetime
                "dag_run": dag_run,  # type: ignore[typeddict-item]  # Removable after #46522
                "partition_key": dag_run.partition_key,
                "partition_date": coerce_datetime(dag_run.partition_date),
                "triggering_asset_events": TriggeringAssetEventsAccessor.build(
                    AssetEventDagRunReferenceResult.from_asset_event_dag_run_reference(event)
                    for event in dag_run.consumed_asset_events
                ),
                "task_instance_key_str": f"{self.task.dag_id}__{self.task.task_id}__{dag_run.run_id}",
                "task_reschedule_count": from_server.task_reschedule_count or 0,
                "prev_start_date_success": lazy_object_proxy.Proxy(
                    lambda: coerce_datetime(get_previous_dagrun_success(self.id).start_date)
                ),
                "prev_end_date_success": lazy_object_proxy.Proxy(
                    lambda: coerce_datetime(get_previous_dagrun_success(self.id).end_date)
                ),
            }
            self._cached_template_context.update(context_from_server)

            if logical_date := coerce_datetime(dag_run.logical_date):
                if TYPE_CHECKING:
                    assert isinstance(logical_date, DateTime)
                ds = logical_date.strftime("%Y-%m-%d")
                ds_nodash = ds.replace("-", "")
                ts = logical_date.isoformat()
                ts_nodash = logical_date.strftime("%Y%m%dT%H%M%S")
                ts_nodash_with_tz = ts.replace("-", "").replace(":", "")
                # logical_date and data_interval either coexist or be None together
                self._cached_template_context.update(
                    {
                        # keys that depend on logical_date
                        "logical_date": logical_date,
                        "ds": ds,
                        "ds_nodash": ds_nodash,
                        "task_instance_key_str": f"{self.task.dag_id}__{self.task.task_id}__{ds_nodash}",
                        "ts": ts,
                        "ts_nodash": ts_nodash,
                        "ts_nodash_with_tz": ts_nodash_with_tz,
                        # keys that depend on data_interval
                        "data_interval_end": coerce_datetime(dag_run.data_interval_end),
                        "data_interval_start": coerce_datetime(dag_run.data_interval_start),
                        "prev_data_interval_start_success": lazy_object_proxy.Proxy(
                            lambda: coerce_datetime(get_previous_dagrun_success(self.id).data_interval_start)
                        ),
                        "prev_data_interval_end_success": lazy_object_proxy.Proxy(
                            lambda: coerce_datetime(get_previous_dagrun_success(self.id).data_interval_end)
                        ),
                    }
                )

            # Backward compatibility: old servers may still send upstream_map_indexes
            upstream_map_indexes = getattr(from_server, "upstream_map_indexes", None)
            if upstream_map_indexes is not None:
                setattr(self, "_upstream_map_indexes", upstream_map_indexes)

        return self._cached_template_context

    @detail_span("render_templates")
    def render_templates(
        self, context: Context | None = None, jinja_env: jinja2.Environment | None = None
    ) -> BaseOperator:
        """
        Render templates in the operator fields.

        If the task was originally mapped, this may replace ``self.task`` with
        the unmapped, fully rendered BaseOperator. The original ``self.task``
        before replacement is returned.
        """
        if not context:
            context = self.get_template_context()
        original_task = self.task

        if TYPE_CHECKING:
            assert context

        ti = context["ti"]

        if TYPE_CHECKING:
            assert original_task
            assert self.task
            assert ti.task

        # If self.task is mapped, this call replaces self.task to point to the
        # unmapped BaseOperator created by this function! This is because the
        # MappedOperator is useless for template rendering, and we need to be
        # able to access the unmapped task instead.
        self.task.render_template_fields(context, jinja_env)
        self.is_mapped = original_task.is_mapped
        return original_task

    def _normalize_xcom_pull_params(
        self,
        task_ids: str | Iterable[str] | None,
        dag_id: str | None,
        map_indexes: int | Iterable[int] | None | ArgNotSet,
        run_id: str | None,
    ) -> tuple[str, str, list[str], bool, bool, Iterable[int | None] | ArgNotSet]:
        """
        Normalize and validate the parameters shared by ``xcom_pull`` and ``axcom_pull``.

        :returns: A tuple of ``(dag_id, run_id, task_ids, single_task_requested,
            single_map_index_requested, map_indexes_iterable)``.  When
            ``map_indexes_iterable`` is ``NOTSET`` the caller should fetch *all*
            map-indexes (i.e. use ``get_all`` / ``aget_all``); otherwise it is an
            iterable of explicit map indexes to fetch one-by-one.
        """
        if dag_id is None:
            dag_id = self.dag_id
        if run_id is None:
            run_id = self.run_id

        single_task_requested = isinstance(task_ids, (str, type(None)))
        single_map_index_requested = isinstance(map_indexes, (int, type(None)))

        if task_ids is None:
            task_ids = [self.task_id]
        elif isinstance(task_ids, str):
            task_ids = [task_ids]
        else:
            task_ids = list(task_ids)

        if not is_arg_set(map_indexes):
            map_indexes_iterable: Iterable[int | None] | ArgNotSet = NOTSET
        elif isinstance(map_indexes, int) or map_indexes is None:
            map_indexes_iterable = [map_indexes]
        elif isinstance(map_indexes, Iterable):
            map_indexes_iterable = map_indexes
        else:
            raise TypeError(
                f"Invalid type for map_indexes: expected int, iterable of ints, or None, got {type(map_indexes)}"
            )

        return (
            dag_id,
            run_id,
            task_ids,
            single_task_requested,
            single_map_index_requested,
            map_indexes_iterable,
        )

    def xcom_pull(
        self,
        task_ids: str | Iterable[str] | None = None,
        dag_id: str | None = None,
        key: str = BaseXCom.XCOM_RETURN_KEY,
        include_prior_dates: bool = False,
        *,
        map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET,
        default: Any = None,
        run_id: str | None = None,
    ) -> Any:
        """
        Pull XComs either from the API server (BaseXCom) or from the custom XCOM backend if configured.

        The pull can be filtered optionally by certain criterion.

        :param key: A key for the XCom. If provided, only XComs with matching
            keys will be returned. The default key is ``'return_value'``, also
            available as constant ``XCOM_RETURN_KEY``. This key is automatically
            given to XComs returned by tasks (as opposed to being pushed
            manually).
        :param task_ids: Only XComs from tasks with matching ids will be
            pulled. If *None* (default), the task_id of the calling task is used.
        :param dag_id: If provided, only pulls XComs from this Dag. If *None*
            (default), the Dag of the calling task is used.
        :param map_indexes: If provided, only pull XComs with matching indexes.
            If *None* (default), this is inferred from the task(s) being pulled
            (see below for details).
        :param include_prior_dates: If False, only XComs from the current
            logical_date are returned. If *True*, XComs from previous dates
            are returned as well.
        :param run_id: If provided, only pulls XComs from a DagRun w/a matching run_id.
            If *None* (default), the run_id of the calling task is used.

        When pulling one single task (``task_id`` is *None* or a str) without
        specifying ``map_indexes``, the return value is a single XCom entry
        (map_indexes is set to map_index of the calling task instance).

        When pulling task is mapped the specified ``map_index`` is used, so by default
        pulling on mapped task will result in no matching XComs if the task instance
        of the method call is not mapped. Otherwise, the map_index of the calling task
        instance is used. Setting ``map_indexes`` to *None* will pull XCom as it would
        from a non mapped task.

        In either case, ``default`` (*None* if not specified) is returned if no
        matching XComs are found.

        When pulling multiple tasks (i.e. either ``task_id`` or ``map_index`` is
        a non-str iterable), a list of matching XComs is returned. Elements in
        the list is ordered by item ordering in ``task_id`` and ``map_index``.
        """
        dag_id, run_id, task_ids, single_task_requested, single_map_index_requested, map_indexes_iterable = (
            self._normalize_xcom_pull_params(task_ids, dag_id, map_indexes, run_id)
        )

        xcoms: list[Any] = []

        if not is_arg_set(map_indexes_iterable):
            # map_indexes was not specified — fetch all map indexes for each task
            for t_id in task_ids:
                values = XCom.get_all(
                    run_id=run_id,
                    key=key,
                    task_id=t_id,
                    dag_id=dag_id,
                    include_prior_dates=include_prior_dates,
                )
                xcoms.append(None) if values is None else xcoms.extend(values)
            # For a single task pulling from an unmapped task, return a single value
            if single_task_requested and len(xcoms) == 1:
                return xcoms[0]
            return xcoms

        for t_id, m_idx in product(task_ids, map_indexes_iterable):
            value = XCom.get_one(
                run_id=run_id,
                key=key,
                task_id=t_id,
                dag_id=dag_id,
                map_index=m_idx,
                include_prior_dates=include_prior_dates,
            )
            xcoms.append(default if value is None else value)

        if single_task_requested and single_map_index_requested:
            return xcoms[0]
        return xcoms

    async def axcom_pull(
        self,
        task_ids: str | Iterable[str] | None = None,
        dag_id: str | None = None,
        key: str = BaseXCom.XCOM_RETURN_KEY,
        include_prior_dates: bool = False,
        *,
        map_indexes: int | Iterable[int] | None | ArgNotSet = NOTSET,
        default: Any = None,
        run_id: str | None = None,
    ) -> Any:
        """Async version of :meth:`xcom_pull`; see that method for full documentation."""
        dag_id, run_id, task_ids, single_task_requested, single_map_index_requested, map_indexes_iterable = (
            self._normalize_xcom_pull_params(task_ids, dag_id, map_indexes, run_id)
        )

        xcoms: list[Any] = []

        if not is_arg_set(map_indexes_iterable):
            # map_indexes was not specified — fetch all map indexes for each task
            for t_id in task_ids:
                values = await XCom.aget_all(
                    run_id=run_id,
                    key=key,
                    task_id=t_id,
                    dag_id=dag_id,
                    include_prior_dates=include_prior_dates,
                )
                xcoms.append(None) if values is None else xcoms.extend(values)
            # For a single task pulling from an unmapped task, return a single value
            if single_task_requested and len(xcoms) == 1:
                return xcoms[0]
            return xcoms

        for t_id, m_idx in product(task_ids, map_indexes_iterable):
            value = await XCom.aget_one(
                run_id=run_id,
                key=key,
                task_id=t_id,
                dag_id=dag_id,
                map_index=m_idx,
                include_prior_dates=include_prior_dates,
            )
            xcoms.append(default if value is None else value)

        if single_task_requested and single_map_index_requested:
            return xcoms[0]

        return xcoms

    def xcom_push(self, key: str, value: Any):
        """
        Make an XCom available for tasks to pull.

        :param key: Key to store the value under.
        :param value: Value to store. Only be JSON-serializable values may be used.
        """
        _xcom_push(self, key, value)

    async def axcom_push(self, key: str, value: Any):
        """
        Make an XCom available for tasks to pull asynchronously.

        :param key: Key to store the value under.
        :param value: Value to store. Only be JSON-serializable values may be used.
        """
        await _axcom_push(self, key, value)

    def get_relevant_upstream_map_indexes(
        self, upstream: BaseOperator, ti_count: int | None, session: Any
    ) -> int | range | None:
        """
        Compute the relevant upstream map indexes for XCom resolution.

        :param upstream: The upstream operator
        :param ti_count: The total count of task instances for this task's expansion
        :param session: Not used (kept for API compatibility)
        :return: None (use entire value), int (single index), or range (subset of indexes)
        """
        from airflow.sdk.execution_time.task_mapping import get_relevant_map_indexes, get_ti_count_for_task

        map_index = self.map_index
        if map_index is None or map_index < 0:
            return None

        # If ti_count not provided, we need to query it
        if ti_count is None:
            ti_count = get_ti_count_for_task(self.task_id, self.dag_id, self.run_id)

        if not ti_count:
            return None

        return get_relevant_map_indexes(
            task=self.task,
            run_id=self.run_id,
            map_index=map_index,
            ti_count=ti_count,
            relative=upstream,
            dag_id=self.dag_id,
        )

    def get_first_reschedule_date(self, context: Context) -> AwareDatetime | None:
        """Get the first reschedule date for the task instance if found, none otherwise."""
        if context.get("task_reschedule_count", 0) == 0:
            # If the task has not been rescheduled, there is no need to ask the supervisor
            return None

        max_tries: int = self.max_tries
        retries: int = self.task.retries or 0
        first_try_number = max_tries - retries + 1

        log.debug("Requesting first reschedule date from supervisor")

        response = SUPERVISOR_COMMS.send(
            msg=GetTaskRescheduleStartDate(ti_id=self.id, try_number=first_try_number)
        )

        if TYPE_CHECKING:
            assert isinstance(response, TaskRescheduleStartDate)

        return response.start_date

    def get_previous_dagrun(self, state: str | None = None) -> DagRun | None:
        """Return the previous Dag run before the given logical date, optionally filtered by state."""
        context = self.get_template_context()
        dag_run = context.get("dag_run")

        log.debug("Getting previous Dag run", dag_run=dag_run)

        if dag_run is None:
            return None

        if dag_run.logical_date is None:
            return None

        response = SUPERVISOR_COMMS.send(
            msg=GetPreviousDagRun(dag_id=self.dag_id, logical_date=dag_run.logical_date, state=state)
        )

        if TYPE_CHECKING:
            assert isinstance(response, PreviousDagRunResult)

        return response.dag_run

    def get_previous_ti(
        self,
        state: TaskInstanceState | None = None,
        logical_date: AwareDatetime | None = None,
        map_index: int = -1,
    ) -> PreviousTIResponse | None:
        """
        Return the previous task instance matching the given criteria.

        :param state: Filter by TaskInstance state
        :param logical_date: Filter by logical date (returns TI before this date)
        :param map_index: Filter by map_index (defaults to -1 for non-mapped tasks)
        :return: Previous task instance or None if not found
        """
        context = self.get_template_context()
        dag_run = context.get("dag_run")

        log.debug("Getting previous task instance", task_id=self.task_id, state=state)

        # Use current dag run's logical_date if not provided
        effective_logical_date = logical_date
        if effective_logical_date is None and dag_run and dag_run.logical_date:
            effective_logical_date = dag_run.logical_date

        response = SUPERVISOR_COMMS.send(
            msg=GetPreviousTI(
                dag_id=self.dag_id,
                task_id=self.task_id,
                logical_date=effective_logical_date,
                map_index=map_index,
                state=state,
            )
        )

        if TYPE_CHECKING:
            assert isinstance(response, PreviousTIResult)

        return response.task_instance

    @staticmethod
    def get_ti_count(
        dag_id: str,
        map_index: int | None = None,
        task_ids: list[str] | None = None,
        task_group_id: str | None = None,
        logical_dates: list[datetime] | None = None,
        run_ids: list[str] | None = None,
        states: list[str] | None = None,
    ) -> int:
        """Return the number of task instances matching the given criteria."""
        response = SUPERVISOR_COMMS.send(
            GetTICount(
                dag_id=dag_id,
                map_index=map_index,
                task_ids=task_ids,
                task_group_id=task_group_id,
                logical_dates=logical_dates,
                run_ids=run_ids,
                states=states,
            ),
        )

        if TYPE_CHECKING:
            assert isinstance(response, TICount)

        return response.count

    @staticmethod
    def get_task_states(
        dag_id: str,
        map_index: int | None = None,
        task_ids: list[str] | None = None,
        task_group_id: str | None = None,
        logical_dates: list[datetime] | None = None,
        run_ids: list[str] | None = None,
    ) -> dict[str, Any]:
        """Return the task states matching the given criteria."""
        response = SUPERVISOR_COMMS.send(
            GetTaskStates(
                dag_id=dag_id,
                map_index=map_index,
                task_ids=task_ids,
                task_group_id=task_group_id,
                logical_dates=logical_dates,
                run_ids=run_ids,
            ),
        )

        if TYPE_CHECKING:
            assert isinstance(response, TaskStatesResult)

        return response.task_states

    @staticmethod
    def get_task_breadcrumbs(dag_id: str, run_id: str) -> Iterable[dict[str, Any]]:
        """Return task breadcrumbs for the given dag run."""
        response = SUPERVISOR_COMMS.send(GetTaskBreadcrumbs(dag_id=dag_id, run_id=run_id))
        if TYPE_CHECKING:
            assert isinstance(response, TaskBreadcrumbsResult)
        return response.breadcrumbs

    @staticmethod
    def get_dr_count(
        dag_id: str,
        logical_dates: list[datetime] | None = None,
        run_ids: list[str] | None = None,
        states: list[str] | None = None,
    ) -> int:
        """Return the number of Dag runs matching the given criteria."""
        response = SUPERVISOR_COMMS.send(
            GetDRCount(
                dag_id=dag_id,
                logical_dates=logical_dates,
                run_ids=run_ids,
                states=states,
            ),
        )

        if TYPE_CHECKING:
            assert isinstance(response, DRCount)

        return response.count

    @staticmethod
    def get_dagrun_state(dag_id: str, run_id: str) -> str:
        """Return the state of the Dag run with the given Run ID."""
        response = SUPERVISOR_COMMS.send(msg=GetDagRunState(dag_id=dag_id, run_id=run_id))

        if TYPE_CHECKING:
            assert isinstance(response, DagRunStateResult)

        return response.state

    @staticmethod
    def get_dag(dag_id: str) -> DagResult:
        """Return the DAG with the given dag_id."""
        response = SUPERVISOR_COMMS.send(msg=GetDag(dag_id=dag_id))

        if TYPE_CHECKING:
            assert isinstance(response, DagResult)

        return response

    @property
    def log_url(self) -> str:
        run_id = quote(self.run_id)
        base_url = conf.get("api", "base_url", fallback="http://localhost:8080/")
        map_index_value = self.map_index
        map_index = (
            f"/mapped/{map_index_value}" if map_index_value is not None and map_index_value >= 0 else ""
        )
        try_number_value = self.try_number
        try_number = (
            f"?try_number={try_number_value}" if try_number_value is not None and try_number_value > 0 else ""
        )
        _log_uri = f"{base_url.rstrip('/')}/dags/{self.dag_id}/runs/{run_id}/tasks/{self.task_id}{map_index}{try_number}"
        return _log_uri

    @property
    def mark_success_url(self) -> str:
        """URL to mark TI success."""
        return self.log_url


def _xcom_push(
    ti: RuntimeTaskInstance,
    key: str,
    value: Any,
    *,
    mapped_length: int | None = None,
) -> None:
    """Push a XCom through XCom.set, which pushes to XCom Backend if configured."""
    # Private function, as we don't want to expose the ability to manually set `mapped_length` to SDK
    # consumers

    XCom.set(
        key=key,
        value=value,
        dag_id=ti.dag_id,
        task_id=ti.task_id,
        run_id=ti.run_id,
        map_index=ti.map_index,
        dag_result=ti.task.returns_dag_result,
        _mapped_length=mapped_length,
    )


async def _axcom_push(
    ti: RuntimeTaskInstance,
    key: str,
    value: Any,
    *,
    mapped_length: int | None = None,
) -> None:
    """Push a XCom through XCom.aset, which pushes to XCom Backend if configured."""
    # Private function, as we don't want to expose the ability to manually set `mapped_length` to SDK
    # consumers

    await XCom.aset(
        key=key,
        value=value,
        dag_id=ti.dag_id,
        task_id=ti.task_id,
        run_id=ti.run_id,
        map_index=ti.map_index,
        dag_result=ti.task.returns_dag_result,
        _mapped_length=mapped_length,
    )


def _xcom_push_to_db(ti: RuntimeTaskInstance, key: str, value: Any) -> None:
    """Push a XCom directly to metadata DB, bypassing custom xcom_backend."""
    XCom._set_xcom_in_db(
        key=key,
        value=value,
        dag_id=ti.dag_id,
        task_id=ti.task_id,
        run_id=ti.run_id,
        map_index=ti.map_index,
    )


def _maybe_reschedule_startup_failure(
    *,
    ti_context: TIRunContext,
    log: Logger,
) -> None:
    """
    Attempt to reschedule the task when a startup failure occurs.

    This does not count as a retry. If the reschedule limit is exceeded, this function
    returns and the caller should fail the task.
    """
    missing_dag_retries = conf.getint("workers", "missing_dag_retries", fallback=3)
    missing_dag_retry_delay = conf.getint("workers", "missing_dag_retry_delay", fallback=60)

    reschedule_count = int(getattr(ti_context, "task_reschedule_count", 0) or 0)
    if missing_dag_retries > 0 and reschedule_count < missing_dag_retries:
        raise AirflowRescheduleException(
            reschedule_date=datetime.now(tz=timezone.utc) + timedelta(seconds=missing_dag_retry_delay)
        )

    log.error(
        "Startup reschedule limit exceeded",
        reschedule_count=reschedule_count,
        max_reschedules=missing_dag_retries,
    )


def _register_deserialization_allowed_classes(dag, log: Logger) -> None:
    """
    Register every operator-declared XCom model class in the deserialization allow-list.

    Runs once per task-run startup, walking the whole DAG, so a consumer task can
    deserialize a producer's structured output even though only the producer
    constructs the value. Reads the declared classes off real operators
    (``getattr``) and off not-yet-expanded mapped operators (``partial_kwargs``),
    and works regardless of how the DAG was loaded -- a fresh parse, or a cache
    that reconstructs operators without running ``__init__``.

    Failures to register a single class are logged and skipped rather than failing
    task startup; a genuinely undeserializable declaration surfaces later as the
    normal allow-list ImportError at consume time.
    """
    for op in dag.tasks:
        if isinstance(op, MappedOperator):
            fields = getattr(op.operator_class, "deserialization_allowed_class_fields", ())
            values = [op.partial_kwargs.get(field) for field in fields]
        else:
            fields = getattr(op, "deserialization_allowed_class_fields", ())
            values = [getattr(op, field, None) for field in fields]
        for value in values:
            if value is None:
                continue
            for model_cls in iter_pydantic_models(value):
                try:
                    allow_class(model_cls)
                except ValueError as exc:
                    log.warning(
                        "Skipping XCom deserialization registration for a model class",
                        task_id=op.task_id,
                        error=str(exc),
                    )


@detail_span("parse")
def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
    # TODO: Task-SDK:
    # Using BundleDagBag here is about 98% wrong, but it'll do for now
    from airflow.dag_processing.dagbag import BundleDagBag

    bundle_info = what.bundle_info
    bundle_prepare_start = time.monotonic()
    bundle_instance = DagBundlesManager().get_bundle(
        name=bundle_info.name,
        version=bundle_info.version,
    )
    bundle_instance.initialize()
    _verify_bundle_access(bundle_instance, log)
    bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000)

    dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path))
    dag_file_parse_start = time.monotonic()
    bag = BundleDagBag(
        dag_folder=dag_absolute_path,
        safe_mode=False,
        load_op_links=False,
        bundle_path=bundle_instance.path,
        bundle_name=bundle_info.name,
    )
    dag_file_parse_ms = int((time.monotonic() - dag_file_parse_start) * 1000)
    if TYPE_CHECKING:
        assert what.ti.dag_id

    try:
        dag = bag.dags[what.ti.dag_id]
    except KeyError:
        log.error(
            "Dag not found during start up", dag_id=what.ti.dag_id, bundle=bundle_info, path=what.dag_rel_path
        )
        _maybe_reschedule_startup_failure(ti_context=what.ti_context, log=log)
        sys.exit(1)

    # install_loader()

    try:
        task = dag.task_dict[what.ti.task_id]
    except KeyError:
        log.error(
            "Task not found in Dag during start up",
            dag_id=dag.dag_id,
            task_id=what.ti.task_id,
            bundle=bundle_info,
            path=what.dag_rel_path,
        )
        _maybe_reschedule_startup_failure(ti_context=what.ti_context, log=log)
        sys.exit(1)

    if not isinstance(task, (BaseOperator, MappedOperator)):
        raise TypeError(
            f"task is of the wrong type, got {type(task)}, wanted {BaseOperator} or {MappedOperator}"
        )

    # Register operator-declared XCom model classes (e.g. ``output_type``) for
    # the whole DAG so this task can deserialize structured output produced by
    # any other task -- including mapped producers and DAGs loaded from a cache
    # that bypasses operator ``__init__``.
    _register_deserialization_allowed_classes(dag, log)

    # Surface the post-RUNNING startup breakdown so support engineers and DAG authors can
    # attribute apparent slow startup to bundle prep (Airflow-side, e.g. git fetch) vs.
    # DAG file parse (user code). Emitted before return so it lands in the task log.
    # Under `run_as_user` impersonation, parse() runs once pre-sudo and again post-sudo,
    # so this event fires twice for those tasks -- the two entries show impersonation overhead.
    log.info(
        "Worker startup parse complete",
        bundle_name=bundle_info.name,
        bundle_version=bundle_info.version,
        dag_file=what.dag_rel_path,
        dag_id=what.ti.dag_id,
        bundle_prepare_ms=bundle_prepare_ms,
        dag_file_parse_ms=dag_file_parse_ms,
    )
    return RuntimeTaskInstance.model_construct(
        **what.ti.model_dump(exclude_unset=True),
        task=task,
        bundle_instance=bundle_instance,
        _ti_context_from_server=what.ti_context,
        max_tries=what.ti_context.max_tries,
        start_date=what.start_date,
        state=TaskInstanceState.RUNNING,
        sentry_integration=what.sentry_integration,
    )


# This global variable will be used by Connection/Variable/XCom classes, or other parts of the task's execution,
# to send requests back to the supervisor process.
#
# Why it needs to be a global:
# - Many parts of Airflow's codebase (e.g., connections, variables, and XComs) may rely on making dynamic requests
#   to the parent process during task execution.
# - These calls occur in various locations and cannot easily pass the `CommsDecoder` instance through the
#   deeply nested execution stack.
# - By defining `SUPERVISOR_COMMS` as a global, it ensures that this communication mechanism is readily
#   accessible wherever needed during task execution without modifying every layer of the call stack.
SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor]


# State machine!
# 1. Start up (receive details from supervisor)
# 2. Execution (run task code, possibly send requests)
# 3. Shutdown and report status


@detail_span("_verify_bundle_access")
def _verify_bundle_access(bundle_instance: BaseDagBundle, log: Logger) -> None:
    """
    Verify bundle is accessible by the current user.

    This is called after user impersonation (if any) to ensure the bundle
    is actually accessible. Uses os.access() which works with any permission
    scheme (standard Unix permissions, ACLs, SELinux, etc.).

    :param bundle_instance: The bundle instance to check
    :param log: Logger instance
    :raises AirflowException: if bundle is not accessible
    """
    from getpass import getuser

    from airflow.sdk.exceptions import AirflowException

    bundle_path = bundle_instance.path

    if not bundle_path.exists():
        # Already handled by initialize() with a warning
        return

    # Check read permission (and execute for directories to list contents)
    access_mode = os.R_OK
    if bundle_path.is_dir():
        access_mode |= os.X_OK

    if not os.access(bundle_path, access_mode):
        raise AirflowException(
            f"Bundle '{bundle_instance.name}' path '{bundle_path}' is not accessible "
            f"by user '{getuser()}'. When using run_as_user, ensure bundle directories "
            f"are readable by the impersonated user. "
            f"See: https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html"
        )


def get_startup_details() -> StartupDetails:
    # The parent sends us a StartupDetails message un-prompted. After this, every single message is only sent
    # in response to us sending a request.

    if os.environ.get("_AIRFLOW__REEXECUTED_PROCESS") == "1" and (
        msgjson := os.environ.get("_AIRFLOW__STARTUP_MSG")
    ):
        # Clear any Kerberos replace cache if there is one, so new process can't reuse it.
        os.environ.pop("KRB5CCNAME", None)
        # entrypoint of re-exec process

        msg: StartupDetails = TypeAdapter(StartupDetails).validate_json(msgjson)
        reinit_supervisor_comms()

        # We delay this message until _after_ we've got the logging re-configured, otherwise it will show up
        # on stdout
        log.debug("Using serialized startup message from environment", msg=msg)
    else:
        # normal entry point
        msg = SUPERVISOR_COMMS._get_response()  # type: ignore[assignment]

        if not isinstance(msg, StartupDetails):
            raise RuntimeError(f"Unhandled startup message {type(msg)} {msg}")
    return msg


@detail_span("startup")
def startup(msg: StartupDetails) -> tuple[RuntimeTaskInstance, Context, Logger]:
    # setproctitle causes issue on Mac OS: https://github.com/benoitc/gunicorn/issues/3021
    os_type = sys.platform
    if os_type == "darwin":
        log.debug("Mac OS detected, skipping setproctitle")
    else:
        from setproctitle import setproctitle

        setproctitle(f"airflow worker -- {msg.ti.id}")

    # Bind TI identifiers so every subsequent log line from this worker process
    # carries ti_id, enabling single-grep lifecycle reconstruction across components.
    bind_contextvars(
        ti_id=str(msg.ti.id),
        dag_id=msg.ti.dag_id,
        task_id=msg.ti.task_id,
        run_id=msg.ti.run_id,
        try_number=msg.ti.try_number,
        map_index=msg.ti.map_index,
    )

    try:
        get_listener_manager().hook.on_starting(component=TaskRunnerMarker())
    except Exception:
        log.exception("error calling listener")

    with _airflow_parsing_context_manager(dag_id=msg.ti.dag_id, task_id=msg.ti.task_id):
        ti = parse(msg, log)
    log.debug("Dag file parsed", file=msg.dag_rel_path)

    run_as_user = getattr(ti.task, "run_as_user", None) or conf.get(
        "core", "default_impersonation", fallback=None
    )

    if os.environ.get("_AIRFLOW__REEXECUTED_PROCESS") != "1" and run_as_user and run_as_user != getuser():
        # enters here for re-exec process
        os.environ["_AIRFLOW__REEXECUTED_PROCESS"] = "1"
        # store startup message in environment for re-exec process
        os.environ["_AIRFLOW__STARTUP_MSG"] = msg.model_dump_json()
        os.set_inheritable(SUPERVISOR_COMMS.socket.fileno(), True)

        # Import main directly from the module instead of re-executing the file.
        # This ensures that when other parts modules import
        # airflow.sdk.execution_time.task_runner, they get the same module instance
        # with the properly initialized SUPERVISOR_COMMS global variable.
        # If we re-executed the module with `python -m`, it would load as __main__ and future
        # imports would get a fresh copy without the initialized globals.
        rexec_python_code = "from airflow.sdk.execution_time.task_runner import main; main()"
        cmd = ["sudo", "-E", "-H", "-u", run_as_user, sys.executable, "-c", rexec_python_code]
        log.info(
            "Running command",
            command=cmd,
        )
        os.execvp("sudo", cmd)

        # ideally, we should never reach here, but if we do, we should return None, None, None
        return None, None, None

    return ti, ti.get_template_context(), log


def _serialize_template_field(
    template_field: Any, name: str
) -> str | dict | list | int | float | bool | None:
    """
    Return a serializable representation of the templated field.

    The walk has two responsibilities:

    1. **Make the template_field JSON-encodable** — every container is rebuilt
       with primitive leaves (str/int/float/bool/None), tuples and sets are
       flattened to lists, and unsupported objects fall through to ``str()``
       so ``json.dumps`` never raises on the result.
    2. **Keep the output deterministic across parses** — callables are replaced
       with their qualified name (never the default ``<function ... at 0x...>``
       repr), dicts are key-sorted, and (frozen)sets are sorted by element so
       the same input always produces the same string.

    Uses the SDK secrets masker to redact secrets in the serialized output.
    """
    import inspect

    from airflow.sdk._shared.module_loading import qualname
    from airflow.sdk._shared.secrets_masker import redact

    def normalize_dict_key(key) -> str:
        """Normalize a dict key to a serialized string type."""
        # Serialized template_field keys must all be strings, not a mix of types, so that
        # downstream json.dumps(..., sort_keys=True) does not raise on mixed-type keys.
        return str(serialize_object(key))

    def serialize_object(obj):
        """Recursively rewrite ``obj`` into a JSON-encodable, hash-stable structure."""
        if obj is None or isinstance(obj, (str, int, float, bool)):
            return obj

        if isinstance(obj, dict):
            # Serialize keys/values first so each key is a string and the output is hash-stable,
            # then sort by the serialized key to prevent hash inconsistencies when dict ordering varies.
            serialized_pairs = [(normalize_dict_key(k), serialize_object(v)) for k, v in obj.items()]
            return dict(sorted(serialized_pairs, key=lambda kv: kv[0]))

        if isinstance(obj, (list, tuple)):
            return [serialize_object(item) for item in obj]

        if isinstance(obj, (set, frozenset)):
            # JSON has no set type → flatten to a list with deterministic ordering
            # so hash randomization on element types cannot shift cross-process iteration order.
            serialized_set = [serialize_object(e) for e in obj]
            return sorted(serialized_set, key=lambda x: (type(x).__name__, str(x)))

        # Use inspect.getattr_static to bypass any custom __getattr__ / metaclass magic
        if callable(inspect.getattr_static(obj, "serialize", None)):
            return serialize_object(obj.serialize())

        # Kubernetes client objects (V1Pod, V1Container, ...) expose their content via to_dict().
        # Scope the branch to the kubernetes namespace so unrelated user classes that happen to
        # define a to_dict() method fall through to str() instead of being treated as K8s payloads.
        if getattr(type(obj), "__module__", "").startswith(
            ("kubernetes.", "kubernetes_asyncio.")
        ) and callable(inspect.getattr_static(obj, "to_dict", None)):
            return serialize_object(obj.to_dict())

        if callable(obj):
            # Use qualified name; default repr embeds memory addresses, which would change the DAG hash on every parse
            return f"<callable {qualname(obj, True)}>"

        # A custom __str__ or __repr__ is treated as an intentional textual representation
        # supplied by the author and used as-is.
        if type(obj).__str__ is not object.__str__ or type(obj).__repr__ is not object.__repr__:
            return str(obj)

        # Otherwise fall back to a qualname marker. The default object repr is
        # `<ClassName object at 0x...>`, which embeds a memory address that flips per process
        # and would break DAG hash stability — use the class qualname instead.
        return f"<{qualname(type(obj), True)} object>"

    max_length = conf.getint("core", "max_templated_field_length")

    serialized = serialize_object(template_field)

    if len(str(serialized)) > max_length:
        # Redact while still structured to preserve nested-key context (so values under
        # documented sensitive keys such as `password`, `token`, `secret`, `api_key`
        # are masked recursively); only stringify the redacted result for truncation.
        rendered = redact(serialized, name)
        return truncate_rendered_value(str(rendered), max_length)

    return serialized


@detail_span("_serialize_rendered_fields")
def _serialize_rendered_fields(task: AbstractOperator) -> dict[str, JsonValue]:
    from airflow.sdk._shared.secrets_masker import redact

    rendered_fields: dict[str, JsonValue] = {}
    for field in task.template_fields:
        value = getattr(task, field)
        serialized = _serialize_template_field(value, field)
        # Redact secrets in the task process itself before sending to API server
        # This ensures that the secrets those are registered via mask_secret() on workers / dag processor are properly masked
        # on the UI.
        redacted = redact(serialized, field)
        rendered_fields[field] = TypeAdapter(JsonValue).validate_python(redacted)

    renderers = getattr(task, "template_fields_renderers", {})
    for renderer_path in renderers:
        if "." not in renderer_path:
            continue

        base_field, _, remainder = renderer_path.partition(".")
        if base_field not in task.template_fields:
            continue

        base_value = getattr(task, base_field, None)
        if base_value is None:
            continue

        current = base_value
        for key in remainder.split("."):
            if isinstance(current, dict):
                current = current.get(key)
            else:
                current = getattr(current, key, None)
            if current is None:
                break

        if current is not None:
            serialized = _serialize_template_field(current, renderer_path)
            redacted = redact(serialized, renderer_path)
            rendered_fields[renderer_path] = TypeAdapter(JsonValue).validate_python(redacted)

    return rendered_fields


def _build_asset_profiles(lineage_objects: list) -> Iterator[AssetProfile]:
    # Lineage can have other types of objects besides assets, so we need to process them a bit.
    for obj in lineage_objects or ():
        if isinstance(obj, Asset):
            yield AssetProfile(name=obj.name, uri=obj.uri, type=Asset.__name__)
        elif isinstance(obj, AssetNameRef):
            yield AssetProfile(name=obj.name, type=AssetNameRef.__name__)
        elif isinstance(obj, AssetUriRef):
            yield AssetProfile(uri=obj.uri, type=AssetUriRef.__name__)
        elif isinstance(obj, AssetAlias):
            yield AssetProfile(name=obj.name, type=AssetAlias.__name__)


def _serialize_outlet_events(events: OutletEventAccessorsProtocol) -> Iterator[dict[str, JsonValue]]:
    if TYPE_CHECKING:
        assert isinstance(events, OutletEventAccessors)
    # We just collect everything the user recorded in the accessors.
    # Further filtering will be done in the API server.
    for key, accessor in events._dict.items():
        if isinstance(key, AssetUniqueKey):
            if accessor.partition_keys:
                for partition_key in accessor.partition_keys:
                    yield {
                        "dest_asset_key": attrs.asdict(key),
                        "extra": accessor.extra,
                        "partition_key": partition_key,
                    }
            else:
                yield {"dest_asset_key": attrs.asdict(key), "extra": accessor.extra}
        for alias_event in accessor.asset_alias_events:
            yield attrs.asdict(alias_event)


@detail_span("_prepare")
def _prepare(ti: RuntimeTaskInstance, log: Logger, context: Context) -> ToSupervisor | None:
    ti.hostname = get_hostname()
    ti.task = ti.task.prepare_for_execution()
    # Since context is now cached, and calling `ti.get_template_context` will return the same dict, we want to
    # update the value of the task that is sent from there
    context["task"] = ti.task

    jinja_env = ti.task.dag.get_template_env()
    ti.render_templates(context=context, jinja_env=jinja_env)

    if rendered_fields := _serialize_rendered_fields(ti.task):
        # so that we do not call the API unnecessarily
        SUPERVISOR_COMMS.send(msg=SetRenderedFields(rendered_fields=rendered_fields))

    # Try to render map_index_template early with available context (will be re-rendered after execution)
    # This provides a partial label during task execution for templates using pre-execution context
    # If rendering fails here, we suppress the error since it will be re-rendered after execution
    try:
        if rendered_map_index := _render_map_index(context, ti=ti, log=log):
            ti.rendered_map_index = rendered_map_index
            log.debug("Sending early rendered map index", length=len(rendered_map_index))
            SUPERVISOR_COMMS.send(msg=SetRenderedMapIndex(rendered_map_index=rendered_map_index))
    except Exception:
        log.debug(
            "Early rendering of map_index_template failed, will retry after task execution", exc_info=True
        )

    _validate_task_inlets_and_outlets(ti=ti, log=log)

    try:
        # TODO: Call pre execute etc.
        get_listener_manager().hook.on_task_instance_running(
            previous_state=TaskInstanceState.QUEUED, task_instance=ti
        )
    except Exception:
        log.exception("error calling listener")

    # No error, carry on and execute the task
    return None


@detail_span("_validate_task_inlets_and_outlets")
def _validate_task_inlets_and_outlets(*, ti: RuntimeTaskInstance, log: Logger) -> None:
    if not ti.task.inlets and not ti.task.outlets:
        return

    inactive_assets_resp = SUPERVISOR_COMMS.send(msg=ValidateInletsAndOutlets(ti_id=ti.id))
    if TYPE_CHECKING:
        assert isinstance(inactive_assets_resp, InactiveAssetsResult)
    if inactive_assets := inactive_assets_resp.inactive_assets:
        raise AirflowInactiveAssetInInletOrOutletException(
            inactive_asset_keys=[
                AssetUniqueKey.from_profile(asset_profile) for asset_profile in inactive_assets
            ]
        )


def _defer_task(
    defer: TaskDeferred, ti: RuntimeTaskInstance, log: Logger
) -> tuple[ToSupervisor, TaskInstanceState]:
    log.info("Pausing task as DEFERRED. ", dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id)
    classpath, trigger_kwargs = defer.trigger.serialize()
    queue: str | None = None
    # Currently, only task-associated BaseTrigger instances may have a non-None queue,
    # and only when triggerer.queues_enabled conf is True.
    if conf.getboolean("triggerer", "queues_enabled", fallback=False) and getattr(
        defer.trigger, "supports_triggerer_queue", True
    ):
        queue = ti.task.queue

    from airflow.sdk.serde import serialize as serde_serialize

    trigger_kwargs = serde_serialize(trigger_kwargs)
    next_kwargs = serde_serialize(defer.kwargs or {})

    if TYPE_CHECKING:
        assert isinstance(next_kwargs, dict)
        assert isinstance(trigger_kwargs, dict)

    msg = DeferTask(
        classpath=classpath,
        trigger_kwargs=trigger_kwargs,
        trigger_timeout=defer.timeout,
        queue=queue,
        next_method=defer.method_name,
        next_kwargs=next_kwargs,
    )
    state = TaskInstanceState.DEFERRED

    return msg, state


def _await_input_task(
    awaiting: TaskAwaitingInput, ti: RuntimeTaskInstance, log: Logger
) -> tuple[ToSupervisor, TaskInstanceState]:
    """Build the message that parks a task in AWAITING_INPUT (HITL), with no trigger."""
    log.info("Pausing task as AWAITING_INPUT.", dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id)

    from airflow.sdk.serde import serialize as serde_serialize

    next_kwargs = serde_serialize(awaiting.kwargs or {})

    if TYPE_CHECKING:
        assert isinstance(next_kwargs, dict)

    msg = AwaitInputTask(
        timeout=awaiting.timeout,
        next_method=awaiting.method_name,
        next_kwargs=next_kwargs,
    )
    state = TaskInstanceState.AWAITING_INPUT

    return msg, state


@Sentry.enrich_errors
@detail_span("run")
def run(
    ti: RuntimeTaskInstance,
    context: Context,
    log: Logger,
) -> tuple[TaskInstanceState, ToSupervisor | None, BaseException | None]:
    """Run the task in this process."""
    import signal

    from airflow.sdk.exceptions import (
        AirflowFailException,
        AirflowRescheduleException,
        AirflowSensorTimeout,
        AirflowSkipException,
        AirflowTaskTerminated,
        DagRunTriggerException,
        DownstreamTasksSkipped,
        TaskAwaitingInput,
        TaskDeferred,
    )

    if TYPE_CHECKING:
        assert ti.task is not None
        assert isinstance(ti.task, BaseOperator)

    parent_pid = os.getpid()

    def _on_term(signum, frame):
        pid = os.getpid()
        if pid != parent_pid:
            return

        ti.task.on_kill()

    signal.signal(signal.SIGTERM, _on_term)

    msg: ToSupervisor | None = None
    state: TaskInstanceState | None = None
    error: BaseException | None = None

    stats_tags = ti.stats_tags
    stats.incr("ti.start", tags=stats_tags)

    try:
        # First, clear the xcom data sent from server
        if ti._ti_context_from_server and (keys_to_delete := ti._ti_context_from_server.xcom_keys_to_clear):
            for x in keys_to_delete:
                log.debug("Clearing XCom with key", key=x)
                XCom.delete(
                    key=x,
                    dag_id=ti.dag_id,
                    task_id=ti.task_id,
                    run_id=ti.run_id,
                    map_index=ti.map_index,
                )

        with set_current_context(context):
            # This is the earliest that we can render templates -- as if it excepts for any reason we need to
            # catch it and handle it like a normal task failure
            if early_exit := _prepare(ti, log, context):
                msg = early_exit
                ti.state = state = TaskInstanceState.FAILED
                return state, msg, error

            try:
                result = _execute_task(context=context, ti=ti, log=log)
                log.info("::group::Post Execute")
            except Exception:
                import jinja2

                # If the task failed, swallow rendering error so it doesn't mask the main error.
                with suppress(jinja2.TemplateSyntaxError, jinja2.UndefinedError):
                    previous_rendered_map_index = ti.rendered_map_index
                    ti.rendered_map_index = _render_map_index(context, ti=ti, log=log)
                    # Send update only if value changed (e.g., user set context variables during execution)
                    if ti.rendered_map_index and ti.rendered_map_index != previous_rendered_map_index:
                        SUPERVISOR_COMMS.send(
                            msg=SetRenderedMapIndex(rendered_map_index=ti.rendered_map_index)
                        )
                raise
            else:  # If the task succeeded, render normally to let rendering error bubble up.
                previous_rendered_map_index = ti.rendered_map_index
                ti.rendered_map_index = _render_map_index(context, ti=ti, log=log)
                # Send update only if value changed (e.g., user set context variables during execution)
                if ti.rendered_map_index and ti.rendered_map_index != previous_rendered_map_index:
                    SUPERVISOR_COMMS.send(msg=SetRenderedMapIndex(rendered_map_index=ti.rendered_map_index))

        _push_xcom_if_needed(result, ti, log)

        msg, state = _handle_current_task_success(context, ti)
    except DownstreamTasksSkipped as skip:
        log.info("::group::Post Execute")
        log.info("Skipping downstream tasks.")
        tasks_to_skip = skip.tasks if isinstance(skip.tasks, list) else [skip.tasks]
        SUPERVISOR_COMMS.send(msg=SkipDownstreamTasks(tasks=tasks_to_skip))
        msg, state = _handle_current_task_success(context, ti)
    except DagRunTriggerException as drte:
        log.info("::group::Post Execute")
        msg, state = _handle_trigger_dag_run(drte, context, ti, log)
    except TaskDeferred as defer:
        log.info("::group::Post Execute")
        msg, state = _defer_task(defer, ti, log)
    except TaskAwaitingInput as awaiting:
        log.info("::group::Post Execute")
        msg, state = _await_input_task(awaiting, ti, log)
    except AirflowSkipException as e:
        log.info("::group::Post Execute")
        if e.args:
            log.info("Skipping task.", reason=e.args[0])
        msg = TaskState(
            state=TaskInstanceState.SKIPPED,
            end_date=datetime.now(tz=timezone.utc),
            rendered_map_index=ti.rendered_map_index,
        )
        state = TaskInstanceState.SKIPPED
    except AirflowRescheduleException as reschedule:
        log.info("::group::Post Execute")
        log.info("Rescheduling task, marking task as UP_FOR_RESCHEDULE")
        msg = RescheduleTask(
            reschedule_date=reschedule.reschedule_date, end_date=datetime.now(tz=timezone.utc)
        )
        state = TaskInstanceState.UP_FOR_RESCHEDULE
    except (AirflowFailException, AirflowSensorTimeout) as e:
        # If AirflowFailException is raised, task should not retry.
        # If a sensor in reschedule mode reaches timeout, task should not retry.
        log.exception("Task failed with exception")
        log.info("::group::Post Execute")
        ti.end_date = datetime.now(tz=timezone.utc)
        msg = TaskState(
            state=TaskInstanceState.FAILED,
            end_date=ti.end_date,
            rendered_map_index=ti.rendered_map_index,
        )
        state = TaskInstanceState.FAILED
        error = e
    except (AirflowTaskTimeout, AirflowException, AirflowRuntimeError) as e:
        # We should allow retries if the task has defined it.
        log.exception("Task failed with exception")
        log.info("::group::Post Execute")
        msg, state = _handle_current_task_failed(ti, e, log, context)
        error = e
    except AirflowTaskTerminated as e:
        # External state updates are already handled with `ti_heartbeat` and will be
        # updated already be another UI API. So, these exceptions should ideally never be thrown.
        # If these are thrown, we should mark the TI state as failed.
        log.exception("Task failed with exception")
        log.info("::group::Post Execute")
        ti.end_date = datetime.now(tz=timezone.utc)
        msg = TaskState(
            state=TaskInstanceState.FAILED,
            end_date=ti.end_date,
            rendered_map_index=ti.rendered_map_index,
        )
        state = TaskInstanceState.FAILED
        error = e
    except SystemExit as e:
        # SystemExit needs to be retried if they are eligible.
        log.error("Task exited", exit_code=e.code)
        log.info("::group::Post Execute")
        msg, state = _handle_current_task_failed(ti, e, log, context)
        error = e
    except BaseException as e:
        log.exception("Task failed with exception")
        log.info("::group::Post Execute")
        msg, state = _handle_current_task_failed(ti, e, log, context)
        error = e
    finally:
        # `state` may still be unset if an exception handler above raised before
        # binding it
        if state is not None:
            stats.incr("ti.finish", tags={**stats_tags, "state": state.value})

        if msg:
            # If the supervisor rejects the terminal-state report
            # (e.g. the server already moved the TI to a terminal state and
            # returns 409 - for example due to issue fixed by #65594),
            # we still want to run finalize() so listeners observe the task state.
            try:
                SUPERVISOR_COMMS.send(msg=msg)
            except Exception:
                log.exception(
                    "Failed to report terminal task state to supervisor",
                    state=state.value if state is not None else None,
                )
                # Fail closed for FAILED / UP_FOR_RETRY: when the supervisor
                # never receives the terminal-state message, exiting 0 would
                # let the supervisor's final_state property default to
                # SUCCESS (exit_code == 0 with no _terminal_state set),
                # turning a real failure into a silent data-quality bug for
                # every downstream task. We signal main() to sys.exit(1)
                # AFTER finalize() runs, so on_failure_callback /
                # on_retry_callback / listener hooks / email_on_failure /
                # email_on_retry still fire. sys.exit(1) directly here would
                # raise SystemExit, which is BaseException, not Exception —
                # main()'s `except Exception:` would not catch it and
                # finalize() at the call site would be skipped.
                #
                # SKIPPED / UP_FOR_RESCHEDULE / DEFERRED are intentionally
                # not fail-closed: supervisor's final_state would misclassify
                # them too, but exiting non-zero would map them to FAILED,
                # which is strictly worse than the default. Those need a
                # separate fix in supervisor's final_state.
                if state in (TaskInstanceState.FAILED, TaskInstanceState.UP_FOR_RETRY):
                    ti._terminal_state_send_failed = True

    # Return the message to make unit tests easier too
    ti.state = state
    return state, msg, error


def _handle_current_task_success(
    context: Context,
    ti: RuntimeTaskInstance,
) -> tuple[SucceedTask, TaskInstanceState]:
    end_date = datetime.now(tz=timezone.utc)
    ti.end_date = end_date

    # Record operator and task instance success metrics
    operator = ti.task.__class__.__name__
    stats_tags = ti.stats_tags

    stats.incr("operator_successes", tags={**stats_tags, "operator_name": operator})
    stats.incr("ti_successes", tags=stats_tags)

    task_outlets = list(_build_asset_profiles(ti.task.outlets))
    outlet_events = list(_serialize_outlet_events(context["outlet_events"]))

    if conf.getboolean("state_store", "clear_on_success") and _get_worker_state_store_backend() is not None:
        log.info(
            "Clearing task state from custom backend as clear_on_success is enabled. The database references will be cleared by the API server."
        )
        context["task_state_store"]._clear_backend_only()

    msg = SucceedTask(
        end_date=end_date,
        task_outlets=task_outlets,
        outlet_events=outlet_events,
        rendered_map_index=ti.rendered_map_index,
    )
    return msg, TaskInstanceState.SUCCESS


def _evaluate_retry_policy(
    ti: RuntimeTaskInstance,
    exception: BaseException,
    log: Logger,
    context: Context | None = None,
) -> RetryDecision | None:
    """
    Evaluate the task's retry policy if one is configured.

    Returns ``None`` when no policy is configured so the caller falls through
    to the standard retry logic.
    """
    policy = getattr(ti.task, "retry_policy", None)
    if policy is None:
        return None
    try:
        max_tries = ti._ti_context_from_server.max_tries if ti._ti_context_from_server else 0
        decision = policy.evaluate(
            exception=exception,
            try_number=ti.try_number,
            max_tries=max_tries,
            context=context,
        )
        if decision.reason:
            log.info("Retry policy decision", action=decision.action.value, reason=decision.reason)
        return decision
    except Exception:
        log.exception("Retry policy evaluation failed, falling back to default behaviour")
        return None


def _handle_current_task_failed(
    ti: RuntimeTaskInstance,
    exception: BaseException,
    log: Logger,
    context: Context | None = None,
) -> tuple[RetryTask | TaskState, TaskInstanceState]:
    """
    Handle a failed task: evaluate the retry policy (if any) and decide the next state.

    This is the entry point every failure path routes through. When the policy
    returns FAIL the task is marked as failed immediately, bypassing the normal
    retry-count check.  When it returns RETRY with a custom delay, that delay is
    forwarded in the ``RetryTask`` message.  For DEFAULT (or when no policy is
    configured), the standard ``_finalize_task_failure`` logic runs.
    """
    from airflow.sdk.definitions.retry_policy import RetryAction

    decision = _evaluate_retry_policy(ti, exception, log, context)
    if decision is not None and decision.action == RetryAction.FAIL:
        ti.end_date = datetime.now(tz=timezone.utc)
        return (
            TaskState(
                state=TaskInstanceState.FAILED,
                end_date=ti.end_date,
                rendered_map_index=ti.rendered_map_index,
            ),
            TaskInstanceState.FAILED,
        )
    if decision is not None and decision.action == RetryAction.RETRY:
        return _finalize_task_failure(
            ti, retry_delay_override=decision.retry_delay, retry_reason=decision.reason
        )
    return _finalize_task_failure(ti)


def _finalize_task_failure(
    ti: RuntimeTaskInstance,
    retry_delay_override: timedelta | None = None,
    retry_reason: str | None = None,
) -> tuple[RetryTask, TaskInstanceState] | tuple[TaskState, TaskInstanceState]:
    """
    Record failure metrics and build the standard retry-or-fail outcome.

    Returns an ``UP_FOR_RETRY`` ``RetryTask`` when the server marked the task
    retry-eligible (optionally carrying a policy-supplied delay/reason), else a
    ``FAILED`` ``TaskState``. This is the default path; retry-policy overrides
    are decided in :func:`_handle_current_task_failed` before this is called.
    """
    end_date = datetime.now(tz=timezone.utc)
    ti.end_date = end_date

    # Record operator and task instance failed metrics
    operator = ti.task.__class__.__name__
    stats_tags = ti.stats_tags

    stats.incr("operator_failures", tags={**stats_tags, "operator_name": operator})
    stats.incr("ti_failures", tags=stats_tags)

    if ti._ti_context_from_server and ti._ti_context_from_server.should_retry:
        retry_kwargs: dict[str, Any] = {"end_date": end_date}
        if retry_delay_override is not None:
            retry_kwargs["retry_delay_seconds"] = retry_delay_override.total_seconds()
        if retry_reason is not None:
            retry_kwargs["retry_reason"] = retry_reason[:500]
        return RetryTask(**retry_kwargs), TaskInstanceState.UP_FOR_RETRY
    return (
        TaskState(
            state=TaskInstanceState.FAILED, end_date=end_date, rendered_map_index=ti.rendered_map_index
        ),
        TaskInstanceState.FAILED,
    )


def _handle_trigger_dag_run(
    drte: DagRunTriggerException, context: Context, ti: RuntimeTaskInstance, log: Logger
) -> tuple[ToSupervisor, TaskInstanceState]:
    """Handle exception from TriggerDagRunOperator."""
    log.info("Triggering Dag Run.", trigger_dag_id=drte.trigger_dag_id)
    comms_msg = SUPERVISOR_COMMS.send(
        TriggerDagRun(
            dag_id=drte.trigger_dag_id,
            run_id=drte.dag_run_id,
            logical_date=drte.logical_date,
            run_after=drte.run_after,
            conf=drte.conf,
            reset_dag_run=drte.reset_dag_run,
            note=drte.note,
        ),
    )

    if isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAGRUN_ALREADY_EXISTS:
        if drte.skip_when_already_exists:
            log.info(
                "Dag Run already exists, skipping task as skip_when_already_exists is set to True.",
                dag_id=drte.trigger_dag_id,
            )
            msg = TaskState(
                state=TaskInstanceState.SKIPPED,
                end_date=datetime.now(tz=timezone.utc),
                rendered_map_index=ti.rendered_map_index,
            )
            state = TaskInstanceState.SKIPPED
        else:
            log.error("Dag Run already exists, marking task as failed.", dag_id=drte.trigger_dag_id)
            msg = TaskState(
                state=TaskInstanceState.FAILED,
                end_date=datetime.now(tz=timezone.utc),
                rendered_map_index=ti.rendered_map_index,
            )
            state = TaskInstanceState.FAILED

        return msg, state

    log.info("Dag Run triggered successfully.", trigger_dag_id=drte.trigger_dag_id)

    # Store the run id from the dag run (either created or found above) to
    # be used when creating the extra link on the webserver.
    ti.xcom_push(key="trigger_run_id", value=drte.dag_run_id)

    if drte.wait_for_completion:
        if drte.deferrable:
            from airflow.providers.standard.triggers.external_task import DagStateTrigger

            defer = TaskDeferred(
                trigger=DagStateTrigger(
                    dag_id=drte.trigger_dag_id,
                    states=drte.allowed_states + drte.failed_states,  # type: ignore[arg-type]
                    # Don't filter by execution_dates when run_ids is provided.
                    # run_id uniquely identifies a DAG run, and when reset_dag_run=True,
                    # drte.logical_date might be a newly calculated value that doesn't match
                    # the persisted logical_date in the database, causing the trigger to never find the run.
                    execution_dates=None,
                    run_ids=[drte.dag_run_id],
                    poll_interval=drte.poke_interval,
                ),
                method_name="execute_complete",
            )
            return _defer_task(defer, ti, log)
        while True:
            log.info(
                "Waiting for dag run to complete execution in allowed state.",
                dag_id=drte.trigger_dag_id,
                run_id=drte.dag_run_id,
                allowed_state=drte.allowed_states,
            )
            time.sleep(drte.poke_interval)

            comms_msg = SUPERVISOR_COMMS.send(
                GetDagRunState(dag_id=drte.trigger_dag_id, run_id=drte.dag_run_id)
            )
            if TYPE_CHECKING:
                assert isinstance(comms_msg, DagRunStateResult)
            if comms_msg.state in drte.failed_states:
                log.error(
                    "DagRun finished with failed state.", dag_id=drte.trigger_dag_id, state=comms_msg.state
                )
                # Mirror the deferrable path (DagStateTrigger -> execute_complete raises
                # AirflowException), which flows through run()'s exception handler and
                # therefore honours a configured retry_policy. Synthesize the same
                # exception here so non-deferrable waits evaluate the policy too,
                # falling back to the standard retry-count check when none is set.
                return _handle_current_task_failed(
                    ti,
                    AirflowException(f"{drte.trigger_dag_id} failed with failed state {comms_msg.state}"),
                    log,
                    context,
                )
            if comms_msg.state in drte.allowed_states:
                log.info(
                    "DagRun finished with allowed state.", dag_id=drte.trigger_dag_id, state=comms_msg.state
                )
                break
            log.debug(
                "DagRun not yet in allowed or failed state.",
                dag_id=drte.trigger_dag_id,
                state=comms_msg.state,
            )
    else:
        # Fire-and-forget mode: wait_for_completion=False
        if drte.deferrable:
            log.info(
                "Ignoring deferrable=True because wait_for_completion=False. "
                "Task will complete immediately without waiting for the triggered DAG run.",
                trigger_dag_id=drte.trigger_dag_id,
            )

    return _handle_current_task_success(context, ti)


def _run_task_state_change_callbacks(
    task: BaseOperator,
    kind: Literal[
        "on_execute_callback",
        "on_failure_callback",
        "on_success_callback",
        "on_retry_callback",
        "on_skipped_callback",
    ],
    context: Context,
    log: Logger,
) -> None:
    callback: Callable[[Context], None]
    for i, callback in enumerate(getattr(task, kind)):
        try:
            create_executable_runner(callback, context_get_outlet_events(context), logger=log).run(context)
        except Exception:
            log.exception("Failed to run task callback", kind=kind, index=i, callback=callback)


def _send_error_email_notification(
    task: BaseOperator | MappedOperator,
    ti: RuntimeTaskInstance,
    context: Context,
    error: BaseException | str | None,
    log: Logger,
) -> None:
    """Send email notification for task errors using SmtpNotifier."""
    try:
        from airflow.providers.smtp.notifications.smtp import SmtpNotifier
    except ImportError:
        log.error(
            "Failed to send task failure or retry email notification: "
            "`apache-airflow-providers-smtp` is not installed. "
            "Install this provider to enable email notifications."
        )
        return

    if not task.email:
        return

    subject_template_file = conf.get("email", "subject_template", fallback=None)

    # Read the template file if configured
    if subject_template_file and Path(subject_template_file).exists():
        subject = Path(subject_template_file).read_text()
    else:
        # Fallback to default
        subject = "Airflow alert: {{ti}}"

    html_content_template_file = conf.get("email", "html_content_template", fallback=None)

    # Read the template file if configured
    if html_content_template_file and Path(html_content_template_file).exists():
        html_content = Path(html_content_template_file).read_text()
    else:
        # Fallback to default
        # For reporting purposes, we report based on 1-indexed,
        # not 0-indexed lists (i.e. Try 1 instead of Try 0 for the first attempt).
        html_content = (
            "Try {{try_number}} out of {{max_tries + 1}}<br>"
            "Exception:<br>{{exception_html}}<br>"
            'Log: <a href="{{ti.log_url}}">Link</a><br>'
            "Host: {{ti.hostname}}<br>"
            'Mark success: <a href="{{ti.mark_success_url}}">Link</a><br>'
        )

    # Add exception_html to context for template rendering
    import html

    exception_html = html.escape(str(error)).replace("\n", "<br>")
    additional_context = {
        "exception": error,
        "exception_html": exception_html,
        "try_number": ti.try_number,
        "max_tries": ti.max_tries,
    }
    email_context = {**context, **additional_context}
    to_emails = task.email
    if not to_emails:
        return

    try:
        notifier = SmtpNotifier(
            to=to_emails,
            subject=subject,
            html_content=html_content,
            from_email=conf.get("email", "from_email", fallback="airflow@airflow"),
        )
        notifier(email_context)
    except Exception:
        log.exception("Failed to send email notification")


@detail_span("_execute_task")
def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger):
    """Execute Task (optionally with a Timeout) and push Xcom results."""
    task = ti.task
    execute = task.execute

    if ti._ti_context_from_server and (next_method := ti._ti_context_from_server.next_method):
        from airflow.sdk.serde import deserialize

        next_kwargs_data = ti._ti_context_from_server.next_kwargs or {}
        try:
            if TYPE_CHECKING:
                assert isinstance(next_kwargs_data, dict)
            kwargs = deserialize(next_kwargs_data)
        except (ImportError, KeyError, AttributeError, TypeError):
            from airflow.serialization.serialized_objects import BaseSerialization

            kwargs = BaseSerialization.deserialize(next_kwargs_data)

        if TYPE_CHECKING:
            assert isinstance(kwargs, dict)
        execute = functools.partial(task.resume_execution, next_method=next_method, next_kwargs=kwargs)

    ctx = contextvars.copy_context()
    # Populate the context var so ExecutorSafeguard doesn't complain
    ctx.run(ExecutorSafeguard.tracker.set, task)

    # Export context in os.environ to make it available for operators to use.
    airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True)
    os.environ.update(airflow_context_vars)

    outlet_events = context_get_outlet_events(context)

    if (pre_execute_hook := task._pre_execute_hook) is not None:
        create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context)
    if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not BaseOperator.pre_execute:
        create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context)

    _run_task_state_change_callbacks(task, "on_execute_callback", context, log)

    log.info("::endgroup::")

    if task.execution_timeout:
        from airflow.sdk.execution_time.timeout import timeout

        # TODO: handle timeout in case of deferral
        timeout_seconds = task.execution_timeout.total_seconds()
        try:
            # It's possible we're already timed out, so fast-fail if true
            if timeout_seconds <= 0:
                raise AirflowTaskTimeout()
            # Run task in timeout wrapper
            with timeout(timeout_seconds):
                result = ctx.run(execute, context=context)
        except AirflowTaskTimeout:
            task.on_kill()
            raise
    else:
        result = ctx.run(execute, context=context)

    if (post_execute_hook := task._post_execute_hook) is not None:
        create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result)
    if getattr(post_execute_hook := task.post_execute, "__func__", None) is not BaseOperator.post_execute:
        create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context)

    return result


def _render_map_index(context: Context, ti: RuntimeTaskInstance, log: Logger) -> str | None:
    """Render named map index if the Dag author defined map_index_template at the task level."""
    if (template := context.get("map_index_template")) is None:
        return None
    log.debug("Rendering map_index_template", template_length=len(template))
    jinja_env = ti.task.dag.get_template_env()
    rendered_map_index = jinja_env.from_string(template).render(context)
    log.debug("Map index rendered", length=len(rendered_map_index))
    return rendered_map_index


def _push_xcom_if_needed(result: Any, ti: RuntimeTaskInstance, log: Logger):
    """Push XCom values when task has ``do_xcom_push`` set to ``True`` and the task returns a result."""
    if ti.task.do_xcom_push:
        xcom_value = result
    else:
        xcom_value = None

    has_mapped_dep = next(ti.task.iter_mapped_dependants(), None) is not None
    if xcom_value is None:
        if not ti.is_mapped and has_mapped_dep:
            # Uhoh, a downstream mapped task depends on us to push something to map over
            from airflow.sdk.exceptions import XComForMappingNotPushed

            raise XComForMappingNotPushed()
        return

    mapped_length: int | None = None
    if not ti.is_mapped and has_mapped_dep:
        from airflow.sdk.definitions.mappedoperator import is_mappable_value
        from airflow.sdk.exceptions import UnmappableXComTypePushed

        if not is_mappable_value(xcom_value):
            raise UnmappableXComTypePushed(xcom_value)
        mapped_length = len(xcom_value)

    log.info("Pushing xcom", ti=ti)

    # If the task has multiple outputs, push each output as a separate XCom.
    if ti.task.multiple_outputs:
        if not isinstance(xcom_value, Mapping):
            raise TypeError(
                f"Returned output was type {type(xcom_value)} expected dictionary for multiple_outputs"
            )
        for key in xcom_value.keys():
            if not isinstance(key, str):
                raise TypeError(
                    "Returned dictionary keys must be strings when using "
                    f"multiple_outputs, found {key} ({type(key)}) instead"
                )
        for k, v in result.items():
            ti.xcom_push(k, v)

    _xcom_push(ti, BaseXCom.XCOM_RETURN_KEY, result, mapped_length=mapped_length)


@detail_span("finalize")
def finalize(
    ti: RuntimeTaskInstance,
    state: TaskInstanceState,
    context: Context,
    log: Logger,
    error: BaseException | None = None,
):
    # Record task duration metrics for all terminal states
    if ti.start_date and ti.end_date:
        duration_ms = (ti.end_date - ti.start_date).total_seconds() * 1000
        stats.timing("task.duration", duration_ms, tags=ti.stats_tags)

    task = ti.task
    # Pushing xcom for each operator extra links defined on the operator only.
    for oe in task.operator_extra_links:
        try:
            link, xcom_key = oe.get_link(operator=task, ti_key=ti), oe.xcom_key  # type: ignore[arg-type]
            log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key)
            _xcom_push_to_db(ti, key=xcom_key, value=link)
        except Exception:
            log.exception(
                "Failed to push an xcom for task operator extra link",
                link_name=oe.name,
                xcom_key=oe.xcom_key,
                ti=ti,
            )

    if getattr(ti.task, "overwrite_rtif_after_execution", False):
        log.debug("Overwriting Rendered template fields.")
        if ti.task.template_fields:
            try:
                SUPERVISOR_COMMS.send(SetRenderedFields(rendered_fields=_serialize_rendered_fields(ti.task)))
            except Exception:
                log.exception("Failed to set rendered fields during finalization", ti=ti, task=ti.task)

    log.debug("Running finalizers", ti=ti)
    if state == TaskInstanceState.SUCCESS:
        _run_task_state_change_callbacks(task, "on_success_callback", context, log)
        try:
            get_listener_manager().hook.on_task_instance_success(
                previous_state=TaskInstanceState.RUNNING, task_instance=ti
            )
        except Exception:
            log.exception("error calling listener")
    elif state == TaskInstanceState.SKIPPED:
        _run_task_state_change_callbacks(task, "on_skipped_callback", context, log)
        try:
            get_listener_manager().hook.on_task_instance_skipped(
                previous_state=TaskInstanceState.RUNNING, task_instance=ti
            )
        except Exception:
            log.exception("error calling listener")
    elif state == TaskInstanceState.UP_FOR_RETRY:
        _run_task_state_change_callbacks(task, "on_retry_callback", context, log)
        try:
            get_listener_manager().hook.on_task_instance_failed(
                previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
            )
        except Exception:
            log.exception("error calling listener")
        if error and task.email_on_retry and task.email:
            _send_error_email_notification(task, ti, context, error, log)
    elif state == TaskInstanceState.FAILED:
        _run_task_state_change_callbacks(task, "on_failure_callback", context, log)
        try:
            get_listener_manager().hook.on_task_instance_failed(
                previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
            )
        except Exception:
            log.exception("error calling listener")
        if error and task.email_on_failure and task.email:
            _send_error_email_notification(task, ti, context, error, log)

    try:
        get_listener_manager().hook.before_stopping(component=TaskRunnerMarker())
    except Exception:
        log.exception("error calling listener")

    log.info("::endgroup::")


@contextmanager
def flush_spans():
    try:
        yield
    finally:
        provider = trace.get_tracer_provider()
        if hasattr(provider, "force_flush"):
            from airflow.sdk.configuration import conf

            timeout_millis = conf.getint("traces", "task_runner_flush_timeout_milliseconds", fallback=30000)
            provider.force_flush(timeout_millis=timeout_millis)


@flush_spans()
def main():
    log = structlog.get_logger(logger_name="task")

    global SUPERVISOR_COMMS
    SUPERVISOR_COMMS = CommsDecoder[ToTask, ToSupervisor](log=log)

    stats.initialize(
        factory=stats_utils.get_stats_factory(),
        export_legacy_names=conf.getboolean("metrics", "legacy_names_on"),
    )

    stack = ExitStack()
    span = INVALID_SPAN
    with stack:
        try:
            try:
                log.info("::group::Pre Execute")
                startup_details = get_startup_details()

                # On macOS fork+exec path, the structured log channel wasn't
                # inherited (exec replaces the address space). Request it from
                # the supervisor using the existing ResendLoggingFD mechanism.
                # Must happen after get_startup_details() so we don't read the
                # startup message as a ResendLoggingFD response.
                if os.environ.pop("_AIRFLOW_FORK_EXEC", None) == "1":
                    reinit_supervisor_comms()
                span_ctx_mgr = _make_task_span(msg=startup_details)
                span = stack.enter_context(span_ctx_mgr)
                ti, context, log = startup(msg=startup_details)
            except AirflowRescheduleException as reschedule:
                log.warning("Rescheduling task during startup, marking task as UP_FOR_RESCHEDULE")
                SUPERVISOR_COMMS.send(
                    msg=RescheduleTask(
                        reschedule_date=reschedule.reschedule_date,
                        end_date=datetime.now(tz=timezone.utc),
                    )
                )
                span.record_exception(reschedule)
                span.set_status(
                    Status(StatusCode.ERROR, description=f"Exception: {type(reschedule).__name__}")
                )
                sys.exit(0)
            with BundleVersionLock(
                bundle_name=ti.bundle_instance.name,
                bundle_version=ti.bundle_instance.version,
            ):
                state, _, error = run(ti, context, log)
                context["exception"] = error
                finalize(ti, state, context, log, error)
                # If run() couldn't deliver a FAILED / UP_FOR_RETRY terminal
                # state to the supervisor, fail closed now — finalize() has
                # already run, so callbacks and listeners observed the state.
                if getattr(ti, "_terminal_state_send_failed", False):
                    sys.exit(1)
        except KeyboardInterrupt as e:
            log.exception("Ctrl-c hit")
            span.record_exception(e)
            span.set_status(Status(StatusCode.ERROR, description=f"Exception: {type(e).__name__}"))
            sys.exit(2)
        except Exception as e:
            log.exception("Top level error")
            span.record_exception(e)
            span.set_status(Status(StatusCode.ERROR, description=f"Exception: {type(e).__name__}"))
            sys.exit(1)
        finally:
            # Ensure the request socket is closed on the child side in all circumstances
            # before the process fully terminates.
            if SUPERVISOR_COMMS and SUPERVISOR_COMMS.socket:
                with suppress(Exception):
                    SUPERVISOR_COMMS.socket.close()


def reinit_supervisor_comms() -> None:
    """
    Re-initialize supervisor comms and logging channel in subprocess.

    This is not needed for most cases, but is used when either we re-launch the process via sudo for
    run_as_user, or from inside the python code in a virtualenv (et al.) operator to re-connect so those tasks
    can continue to access variables etc.
    """
    import socket

    if "SUPERVISOR_COMMS" not in globals():
        global SUPERVISOR_COMMS

        fd = int(os.environ.get("__AIRFLOW_SUPERVISOR_FD", "0"))

        SUPERVISOR_COMMS = CommsDecoder[ToTask, ToSupervisor](log=log, socket=socket.socket(fileno=fd))

    logs = SUPERVISOR_COMMS.send(ResendLoggingFD())
    if isinstance(logs, SentFDs):
        from airflow.sdk.log import configure_logging

        log_io = os.fdopen(logs.fds[0], "wb", buffering=0)
        configure_logging(json_output=True, output=log_io, sending_to_supervisor=True)
    else:
        print("Unable to re-configure logging after sudo, we didn't get an FD", file=sys.stderr)


if __name__ == "__main__":
    main()
