# 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.

from __future__ import annotations

import contextlib
import contextvars
import functools
import json
import os
import textwrap
import time
from collections.abc import Iterable
from datetime import datetime, timedelta, timezone as dt_timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any
from unittest import mock
from unittest.mock import call, patch

import pandas as pd
import pytest
import structlog
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import BaseModel
from task_sdk import FAKE_BUNDLE
from uuid6 import uuid7

from airflow._shared.observability.traces import (
    OverrideableRandomIdGenerator,
    new_dagrun_trace_carrier,
    new_task_run_carrier,
)
from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span
from airflow.listeners import hookimpl
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.providers.standard.triggers.temporal import DateTimeTrigger
from airflow.sdk import (
    DAG,
    BaseOperator,
    BaseOperatorLink,
    Connection,
    dag as dag_decorator,
    get_current_context,
    task as task_decorator,
    timezone,
)
from airflow.sdk._shared.observability.metrics.base_stats_logger import StatsLogger
from airflow.sdk._shared.state import AssetScope, TaskScope
from airflow.sdk.api.datamodels._generated import (
    AssetProfile,
    AssetResponse,
    DagRun,
    DagRunState,
    PreviousTIResponse,
    TaskInstance,
    TaskInstanceState,
    TIRunContext,
)
from airflow.sdk.bases.operator import ExecutorSafeguard
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.definitions._internal.types import NOTSET, SET_DURING_EXECUTION, is_arg_set
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetUniqueKey, AssetUriRef, Dataset, Model
from airflow.sdk.definitions.param import DagParam
from airflow.sdk.definitions.retry_policy import (
    ExceptionRetryPolicy,
    RetryAction,
    RetryDecision,
    RetryPolicy,
    RetryRule,
)
from airflow.sdk.exceptions import (
    AirflowException,
    AirflowFailException,
    AirflowRescheduleException,
    AirflowRuntimeError,
    AirflowSensorTimeout,
    AirflowSkipException,
    AirflowTaskTerminated,
    AirflowTaskTimeout,
    DownstreamTasksSkipped,
    ErrorType,
    TaskDeferred,
)
from airflow.sdk.execution_time import task_runner
from airflow.sdk.execution_time.comms import (
    AssetEventResult,
    AssetEventsResult,
    AssetResult,
    AssetsByAliasResult,
    BundleInfo,
    ClearAssetStateStoreByName,
    ClearTaskStateStore,
    ConnectionResult,
    DagResult,
    DagRunStateResult,
    DeferTask,
    DeleteAssetStateStoreByName,
    DeleteTaskStateStore,
    DRCount,
    ErrorResponse,
    GetAssetByUri,
    GetAssetsByAlias,
    GetAssetStateStoreByName,
    GetAssetStateStoreByUri,
    GetConnection,
    GetDag,
    GetDagRunState,
    GetDRCount,
    GetPreviousDagRun,
    GetPreviousTI,
    GetTaskStates,
    GetTaskStateStore,
    GetTICount,
    GetVariable,
    GetXCom,
    GetXComSequenceSlice,
    InactiveAssetsResult,
    MaskSecret,
    OKResponse,
    PreviousDagRunResult,
    PreviousTIResult,
    PrevSuccessfulDagRunResult,
    RescheduleTask,
    RetryTask,
    SetAssetStateStoreByName,
    SetAssetStateStoreByUri,
    SetRenderedFields,
    SetTaskStateStore,
    SetXCom,
    SkipDownstreamTasks,
    StartupDetails,
    SucceedTask,
    TaskRescheduleStartDate,
    TaskState,
    TaskStatesResult,
    TICount,
    TriggerDagRun,
    ValidateInletsAndOutlets,
    VariableResult,
    XComResult,
    XComSequenceSliceResult,
)
from airflow.sdk.execution_time.context import (
    ConnectionAccessor,
    InletEventsAccessors,
    MacrosAccessor,
    OutletEventAccessors,
    TaskStateStoreAccessor,
    TriggeringAssetEventsAccessor,
    VariableAccessor,
    _wrap_external_ref,
)
from airflow.sdk.execution_time.task_runner import (
    RuntimeTaskInstance,
    TaskRunnerMarker,
    _defer_task,
    _execute_task,
    _make_task_span,
    _push_xcom_if_needed,
    _register_deserialization_allowed_classes,
    _run_execute_callable,
    _serialize_outlet_events,
    _xcom_push,
    detail_span,
    finalize,
    get_startup_details,
    parse,
    run,
    startup,
)
from airflow.sdk.execution_time.xcom import XCom
from airflow.sdk.serde import deserialize
from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent
from airflow.triggers.callback import CallbackTrigger
from airflow.triggers.testing import SuccessTrigger

from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.mock_operators import AirflowLink

if TYPE_CHECKING:
    from kgb import SpyAgency
import time_machine


def get_inline_dag(dag_id: str, task: BaseOperator) -> DAG:
    """Creates an inline dag and returns it based on dag_id and task."""
    dag = DAG(dag_id=dag_id, start_date=timezone.datetime(2024, 12, 3))
    task.dag = dag

    return dag


class CustomOperator(BaseOperator):
    def execute(self, context):
        task_id = context["task_instance"].task_id
        print(f"Hello World {task_id}!")


def test_parse(test_dags_dir: Path, make_ti_context):
    """Test that checks parsing of a basic dag with an un-mocked parse."""
    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="a",
            dag_id="super_basic",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="super_basic.py",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    with patch.dict(
        os.environ,
        {
            "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(
                [
                    {
                        "name": "my-bundle",
                        "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
                        "kwargs": {"path": str(test_dags_dir), "refresh_interval": 1},
                    }
                ]
            ),
        },
    ):
        ti = parse(what, mock.Mock())

    assert ti.task
    assert ti.task.dag


@mock.patch("airflow.dag_processing.dagbag.BundleDagBag")
def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context):
    """Test that checks that the BundleDagBag is constructed as expected during parsing"""
    mock_bag_instance = mock.Mock()
    mock_dagbag.return_value = mock_bag_instance
    mock_dag = mock.Mock(spec=DAG)
    mock_task = mock.Mock(spec=BaseOperator)
    # The worker walks dag.tasks to register declared deserialization classes;
    # give the mock an iterable tasks list and an empty declaration so the walk
    # is a no-op for this BundleDagBag-construction test.
    mock_task.deserialization_allowed_class_fields = ()

    mock_bag_instance.dags = {"super_basic": mock_dag}
    mock_dag.task_dict = {"a": mock_task}
    mock_dag.tasks = [mock_task]

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="a",
            dag_id="super_basic",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="super_basic.py",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    with patch.dict(
        os.environ,
        {
            "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(
                [
                    {
                        "name": "my-bundle",
                        "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
                        "kwargs": {"path": str(test_dags_dir), "refresh_interval": 1},
                    }
                ]
            ),
        },
    ):
        parse(what, mock.Mock())

    mock_dagbag.assert_called_once_with(
        dag_folder=mock.ANY,
        safe_mode=False,
        load_op_links=False,
        bundle_path=test_dags_dir,
        bundle_name="my-bundle",
    )


@pytest.mark.parametrize(
    ("dag_id", "task_id", "expected_error"),
    (
        pytest.param(
            "madeup_dag_id",
            "a",
            mock.call(mock.ANY, dag_id="madeup_dag_id", path="super_basic.py"),
            id="dag-not-found",
        ),
        pytest.param(
            "super_basic",
            "no-such-task",
            mock.call(mock.ANY, task_id="no-such-task", dag_id="super_basic", path="super_basic.py"),
            id="task-not-found",
        ),
    ),
)
def test_parse_not_found(test_dags_dir: Path, make_ti_context, dag_id, task_id, expected_error):
    """Check for nice error messages on dag not found."""
    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id=task_id,
            dag_id=dag_id,
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="super_basic.py",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    log = mock.Mock()

    with (
        patch.dict(
            os.environ,
            {
                "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(
                    [
                        {
                            "name": "my-bundle",
                            "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
                            "kwargs": {"path": str(test_dags_dir), "refresh_interval": 1},
                        }
                    ]
                ),
            },
        ),
        pytest.raises(AirflowRescheduleException),
    ):
        parse(what, log)

    expected_error.kwargs["bundle"] = what.bundle_info
    log.error.assert_has_calls([expected_error])


def test_parse_not_found_does_not_reschedule_when_max_attempts_reached(test_dags_dir: Path, make_ti_context):
    """
    If the startup reschedule attempt limit is reached, parsing failures should not be rescheduled
    and should surface as a hard failure (SystemExit in the task runner process).
    """
    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="a",
            dag_id="madeup_dag_id",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="super_basic.py",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(task_reschedule_count=3),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    log = mock.Mock()

    with (
        conf_vars(
            {
                ("workers", "missing_dag_retries"): "3",
                ("workers", "missing_dag_retry_delay"): "60",
            }
        ),
        patch.dict(
            os.environ,
            {
                "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(
                    [
                        {
                            "name": "my-bundle",
                            "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
                            "kwargs": {"path": str(test_dags_dir), "refresh_interval": 1},
                        }
                    ]
                ),
            },
        ),
        pytest.raises(SystemExit),
    ):
        parse(what, log)


@mock.patch("builtins.exit", side_effect=lambda code: (_ for _ in ()).throw(SystemExit(code)))
@mock.patch("airflow.sdk.execution_time.task_runner.startup")
@mock.patch("airflow.sdk.execution_time.task_runner.get_startup_details")
@mock.patch("airflow.sdk.execution_time.task_runner.CommsDecoder")
def test_main_sends_reschedule_task_when_startup_reschedules(
    mock_comms_decoder_cls, mock_get_startup_details, mock_startup, mock_exit, time_machine, make_ti_context
):
    """
    If startup raises AirflowRescheduleException, the task runner should report a RescheduleTask
    message to the supervisor and exit cleanly (code 0).
    """
    ts = datetime(2025, 1, 1, tzinfo=dt_timezone.utc)
    reschedule_date = ts + timedelta(seconds=60)

    mock_comms_instance = mock.Mock()
    mock_comms_instance.socket = None
    mock_comms_decoder_cls.__getitem__.return_value.return_value = mock_comms_instance
    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="my_task",
            dag_id="test_dag",
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
            context_carrier={},
        ),
        dag_rel_path="",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )
    mock_get_startup_details.return_value = what
    mock_startup.side_effect = AirflowRescheduleException(reschedule_date=reschedule_date)

    # Move time
    time_machine.move_to(ts, tick=False)

    # Run & assert
    with pytest.raises(SystemExit) as exc:
        task_runner.main()

    assert exc.value.code == 0
    assert mock_comms_instance.mock_calls == [
        call.send(msg=RescheduleTask(reschedule_date=reschedule_date, end_date=ts))
    ]


def test_run_swallows_supervisor_terminal_send_failure(create_runtime_ti, mock_supervisor_comms):
    """
    When the supervisor rejects the terminal-state send (e.g. the server
    already moved the TI to a terminal state and returns 409), `run()` must
    not let that exception escape — the task body succeeded, so `run()` should
    still return (SUCCESS, msg, None) so `main()` can notify listeners with
    the correct local terminal state.
    """
    from airflow.sdk.execution_time.task_runner import run

    class TrivialOperator(BaseOperator):
        def execute(self, context):
            return "ok"

    task = TrivialOperator(task_id="trivial")
    runtime_ti = create_runtime_ti(task=task)
    context = runtime_ti.get_template_context()
    log = mock.MagicMock()

    # Let every send succeed except the terminal-state send — simulate the
    # server rejecting the SucceedTask with a 409-wrapped AirflowRuntimeError.
    def send_side_effect(msg=None, **kwargs):
        if isinstance(msg, SucceedTask):
            raise RuntimeError("409 invalid_state: previous_state=success")
        return mock.DEFAULT

    mock_supervisor_comms.send.side_effect = send_side_effect

    # Must not raise.
    state, msg, error = run(runtime_ti, context, log)

    assert state == TaskInstanceState.SUCCESS
    assert isinstance(msg, SucceedTask)
    assert error is None


def test_run_signals_fail_closed_when_failure_terminal_send_fails(create_runtime_ti, mock_supervisor_comms):
    """
    When the task FAILS and the terminal-state send to the supervisor fails too
    (e.g. broken Unix socket / supervisor crashed / IPC channel dead), `run()`
    must signal to main() that the process should exit non-zero — otherwise
    the supervisor's `final_state` property defaults exit_code-0-with-no-
    terminal-state to SUCCESS, turning a transient IPC blip into a silent
    data-quality bug downstream.

    The signal is deferred to main() (via `_terminal_state_send_failed` on the
    ti) so finalize() still runs first — on_failure_callback / listener hooks /
    email_on_failure must observe the FAILED state before the process exits.
    """

    class FailingOperator(BaseOperator):
        def execute(self, context):
            raise RuntimeError("task body failed")

    task = FailingOperator(task_id="failing")
    runtime_ti = create_runtime_ti(task=task)
    context = runtime_ti.get_template_context()
    log = mock.MagicMock()

    # Let the terminal-state send raise an IPC-level failure.
    def send_side_effect(msg=None, **kwargs):
        if isinstance(msg, TaskState):
            raise BrokenPipeError("supervisor IPC broken")
        return mock.DEFAULT

    mock_supervisor_comms.send.side_effect = send_side_effect

    # run() must not raise — fail-closed is signalled via the ti attribute
    # so main() can sys.exit(1) only after finalize() has run.
    state, _, _ = run(runtime_ti, context, log)

    assert state == TaskInstanceState.FAILED
    assert runtime_ti._terminal_state_send_failed is True


@pytest.mark.parametrize(
    ("state_when_send_fails", "should_fail_closed"),
    [
        (TaskInstanceState.SUCCESS, False),
        (TaskInstanceState.SKIPPED, False),
    ],
)
def test_run_does_not_signal_fail_closed_for_non_failed_states(
    create_runtime_ti, mock_supervisor_comms, state_when_send_fails, should_fail_closed
):
    """
    Only FAILED / UP_FOR_RETRY are fail-closed. SUCCESS is exempt (the existing
    409-rejection softening). SKIPPED is also exempt: supervisor's final_state
    misclassifies it either way, and exiting non-zero would map it to FAILED,
    which is strictly worse than the default mapping.
    """

    class Op(BaseOperator):
        def execute(self, context):
            if state_when_send_fails == TaskInstanceState.SKIPPED:
                raise AirflowSkipException("skip")
            return "ok"

    task = Op(task_id="op")
    runtime_ti = create_runtime_ti(task=task)
    context = runtime_ti.get_template_context()
    log = mock.MagicMock()

    def send_side_effect(msg=None, **kwargs):
        if isinstance(msg, (TaskState, SucceedTask)):
            raise BrokenPipeError("supervisor IPC broken")
        return mock.DEFAULT

    mock_supervisor_comms.send.side_effect = send_side_effect

    state, _, _ = run(runtime_ti, context, log)

    assert state == state_when_send_fails
    assert getattr(runtime_ti, "_terminal_state_send_failed", False) is should_fail_closed


def test_task_span_is_child_of_dag_run_span(make_ti_context):
    """Full trace hierarchy: dag_run → task_run.my_task (API server) → worker.my_task (task runner)."""
    # Single provider shared by all spans so contexts are compatible.
    in_mem_exporter = InMemorySpanExporter()
    provider = TracerProvider(id_generator=OverrideableRandomIdGenerator())
    provider.add_span_processor(SimpleSpanProcessor(in_mem_exporter))

    # Step 1: create the dag run span and capture its carrier.
    dag_run_tracer = provider.get_tracer("dag_run")
    with dag_run_tracer.start_as_current_span("dag_run.test_dag") as dag_run_span:
        dag_run_carrier: dict[str, str] = {}
        TraceContextTextMapPropagator().inject(dag_run_carrier)
        dag_run_span_ctx = dag_run_span.get_span_context()

    # Step 2: derive the parent task span carrier (child of dag run), as the scheduler does.
    ti_model_tracer = provider.get_tracer("airflow.models.taskinstance")
    with mock.patch("airflow.models.taskinstance.tracer", ti_model_tracer):
        ti_carrier = new_task_run_carrier(dag_run_carrier)

    # Extract the parent task span context (the stable span ID stored in ti_carrier).
    parent_task_span_ctx = trace.get_current_span(
        context=TraceContextTextMapPropagator().extract(ti_carrier)
    ).get_span_context()

    # Step 3: build StartupDetails with ti.context_carrier = ti_carrier.
    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="my_task",
            dag_id="test_dag",
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
            context_carrier=ti_carrier,
        ),
        dag_rel_path="",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    # Step 4: emit the worker span (task runner side).
    task_runner_tracer = provider.get_tracer("airflow.sdk.execution_time.task_runner")
    with mock.patch("airflow.sdk.execution_time.task_runner.tracer", task_runner_tracer):
        with _make_task_span(what):
            pass

    # Step 5: emit the parent task span (API server side, as happens on task completion).
    mock_ti = mock.MagicMock()
    mock_ti.dag_id = "test_dag"
    mock_ti.task_id = "my_task"
    mock_ti.run_id = "test_run"
    mock_ti.try_number = 1
    mock_ti.map_index = -1
    mock_ti.queued_dttm = None
    mock_ti.start_date = timezone.utcnow()
    mock_ti.dag_run.context_carrier = dag_run_carrier
    mock_ti.context_carrier = ti_carrier
    api_tracer = provider.get_tracer("airflow.api_fastapi.execution_api.routes.task_instances")
    with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.tracer", api_tracer):
        _emit_task_span(mock_ti, TaskInstanceState.SUCCESS)

    finished = in_mem_exporter.get_finished_spans()

    # task_run.my_task: emitted by API server, child of dag run, span_id == parent_task_span_ctx.span_id.
    task_spans = [s for s in finished if s.name == "task_run.my_task"]
    assert len(task_spans) == 1
    task_span = task_spans[0]
    assert task_span.parent is not None
    assert task_span.parent.span_id == dag_run_span_ctx.span_id
    assert task_span.context.span_id == parent_task_span_ctx.span_id

    # worker.my_task: created by task runner, child of the parent task span.
    worker_spans = [s for s in finished if s.name == "worker.my_task"]
    assert len(worker_spans) == 1
    assert worker_spans[0].parent is not None
    assert worker_spans[0].parent.span_id == parent_task_span_ctx.span_id

    # All spans share the same trace ID.
    assert {s.context.trace_id for s in finished} == {dag_run_span_ctx.trace_id}


def test_task_span_no_parent_when_no_context_carrier(make_ti_context):
    """When context_carrier is absent, the task span should be a root span (no parent)."""
    from opentelemetry.sdk.trace import TracerProvider
    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
    from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

    in_mem_exporter = InMemorySpanExporter()
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(in_mem_exporter))

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="standalone_task",
            dag_id="test_dag",
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
            context_carrier=None,
        ),
        dag_rel_path="",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    task_tracer = provider.get_tracer("airflow.sdk.execution_time.task_runner")
    with mock.patch("airflow.sdk.execution_time.task_runner.tracer", task_tracer):
        with _make_task_span(what):
            pass

    finished = in_mem_exporter.get_finished_spans()
    assert len(finished) == 1
    assert finished[0].parent is None


def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context):
    """Check that the bundle path is added to sys.path, so Dags can import shared modules."""
    tmp_path.joinpath("util.py").write_text("NAME = 'dag_name'")

    dag1_path = tmp_path.joinpath("path_test.py")
    dag1_code = """
    from util import NAME
    from airflow.sdk import DAG
    from airflow.sdk.bases.operator import BaseOperator
    with DAG(NAME):
        BaseOperator(task_id="a")
    """
    dag1_path.write_text(textwrap.dedent(dag1_code))

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="a",
            dag_id="dag_name",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="path_test.py",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    with patch.dict(
        os.environ,
        {
            "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(
                [
                    {
                        "name": "my-bundle",
                        "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
                        "kwargs": {"path": str(tmp_path), "refresh_interval": 1},
                    }
                ]
            ),
        },
    ):
        ti = parse(what, mock.Mock())

    assert ti.task.dag.dag_id == "dag_name"


def test_verify_bundle_access_raises_when_not_accessible(tmp_path: Path, make_ti_context):
    """Test that _verify_bundle_access raises AirflowException when bundle path is not accessible."""
    from airflow.sdk.execution_time.task_runner import _verify_bundle_access

    # Create a directory that exists
    bundle_path = tmp_path / "test_bundle"
    bundle_path.mkdir()

    # Create a mock bundle instance
    mock_bundle = mock.Mock()
    mock_bundle.path = bundle_path
    mock_bundle.name = "test-bundle"

    # Mock os.access to simulate permission denied (avoids root user issues in CI)
    with patch("airflow.sdk.execution_time.task_runner.os.access", return_value=False):
        with pytest.raises(AirflowException) as exc_info:
            _verify_bundle_access(mock_bundle, mock.Mock())

        assert "not accessible" in str(exc_info.value)
        assert "test-bundle" in str(exc_info.value)


def test_verify_bundle_access_succeeds_when_readable(tmp_path: Path, make_ti_context):
    """Test that _verify_bundle_access succeeds when bundle path is accessible."""
    from airflow.sdk.execution_time.task_runner import _verify_bundle_access

    # Create a directory with read permissions
    bundle_path = tmp_path / "accessible_bundle"
    bundle_path.mkdir()

    mock_bundle = mock.Mock()
    mock_bundle.path = bundle_path
    mock_bundle.name = "test-bundle"

    # Should not raise
    _verify_bundle_access(mock_bundle, mock.Mock())


def test_verify_bundle_access_skips_nonexistent_path(tmp_path: Path):
    """Test that _verify_bundle_access does nothing when bundle path doesn't exist."""
    from airflow.sdk.execution_time.task_runner import _verify_bundle_access

    mock_bundle = mock.Mock()
    mock_bundle.path = tmp_path / "nonexistent"
    mock_bundle.name = "test-bundle"

    # Should not raise - nonexistent paths are handled by initialize()
    _verify_bundle_access(mock_bundle, mock.Mock())


@pytest.mark.parametrize("use_queues", [False, True])
def test_run_deferred_basic(time_machine, create_runtime_ti, mock_supervisor_comms, use_queues: bool):
    """Test that a task can transition to a deferred state."""
    from airflow.providers.standard.sensors.date_time import DateTimeSensorAsync

    task_queue = "fake_q"
    deferred_queue = task_queue if use_queues else None
    # Use the time machine to set the current time
    instant = timezone.datetime(2024, 11, 22)
    with conf_vars({("triggerer", "queues_enabled"): str(use_queues)}):
        task = DateTimeSensorAsync(
            task_id="async",
            target_time=str(instant + timedelta(seconds=3)),
            poke_interval=60,
            timeout=600,
            queue=task_queue,
        )
        time_machine.move_to(instant, tick=False)

        # Expected DeferTask, it is constructed by _defer_task from exception and is sent to supervisor
        expected_defer_task = DeferTask(
            state="deferred",
            classpath="airflow.providers.standard.triggers.temporal.DateTimeTrigger",
            trigger_kwargs={
                "moment": {
                    "__classname__": "pendulum.datetime.DateTime",
                    "__version__": 2,
                    "__data__": {
                        "timestamp": 1732233603.0,
                        "tz": {
                            "__classname__": "builtins.tuple",
                            "__version__": 1,
                            "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                        },
                    },
                },
                "end_from_trigger": False,
            },
            trigger_timeout=None,
            queue=deferred_queue,
            next_method="execute_complete",
            next_kwargs={},
            rendered_map_index=None,
            type="DeferTask",
        )

        # Run the task
        ti = create_runtime_ti(dag_id="basic_deferred_run", task=task)
        run(ti, context=ti.get_template_context(), log=mock.MagicMock())

        assert ti.state == TaskInstanceState.DEFERRED

        # send will only be called when the TaskDeferred exception is raised
        mock_supervisor_comms.send.assert_any_call(expected_defer_task)


class FakeEventTrigger(BaseEventTrigger):
    """Fake event trigger class for testing"""

    def serialize(self) -> tuple[str, dict[str, Any]]:
        return ("tests.task_sdk.execution_time.test_task_runner.FakeEventTrigger", {})

    async def run(self):
        yield TriggerEvent(True)


@conf_vars({("triggerer", "queues_enabled"): "True"})
@pytest.mark.parametrize(
    ("mock_trigger", "expected_trigger_queue"),
    [
        (SuccessTrigger(kwargs={}), "task_q_test"),
        (FakeEventTrigger(kwarg={}), None),
        (
            CallbackTrigger(
                callback_path="classpath.test_callback",
                callback_kwargs={"message": "test_msg", "context": {"dag_run": "test"}},
            ),
            None,
        ),
    ],
)
def test_defer_task_queue_assignment(
    create_runtime_ti, mock_trigger: BaseTrigger, expected_trigger_queue: str | None
) -> None:
    """Ensure `_defer_task` will only pass along origin task queue information to the trigger message when expected."""
    from airflow.providers.standard.operators.empty import EmptyOperator

    mock_task_queue = "task_q_test"
    task = EmptyOperator(task_id="empty_trig_queue_test", queue=mock_task_queue)
    runtime_ti = create_runtime_ti(dag_id="deferred_run", task=task)
    actual_msg, actual_state = _defer_task(
        defer=TaskDeferred(trigger=mock_trigger, method_name="foo"), ti=runtime_ti, log=mock.MagicMock()
    )
    assert isinstance(actual_msg, DeferTask)
    assert actual_state == TaskInstanceState.DEFERRED
    actual_queue = actual_msg.queue
    assert actual_queue == expected_trigger_queue, (
        f"Expected DeferTask's queue value to be {mock_task_queue}, but got {actual_queue}"
    )


@pytest.mark.parametrize(
    ("should_retry", "expected_state"),
    [
        (True, TaskInstanceState.UP_FOR_RETRY),
        (False, TaskInstanceState.FAILED),
    ],
)
def test_defer_with_unserializable_kwargs_honours_retries_and_callbacks(
    should_retry, expected_state, create_runtime_ti, mock_supervisor_comms
):
    """
    A task that defers with a non-serializable ``next_kwargs`` value must fail like any
    other task, rather than taking the whole run down.

    ``_defer_task`` runs ``serde_serialize`` on the deferral kwargs, which raises
    ``TypeError`` for anything it has no serializer for (a file handle, a client object,
    a lambda). That raise happens inside ``run()``'s ``except TaskDeferred`` handler, so
    before the fix it escaped ``run()`` without evaluating retries or running callbacks,
    skipping the retry decision and callbacks even without involving the API server at
    all.
    """
    callbacks_run = []

    class _DeferWithBadKwargs(BaseOperator):
        def execute(self, context):
            raise TaskDeferred(
                trigger=DateTimeTrigger(moment=timezone.datetime(2024, 11, 22)),
                method_name="next",
                # A live handle is the realistic version of this mistake.
                kwargs={"client": object()},
            )

    task = _DeferWithBadKwargs(
        task_id="defer_bad_kwargs",
        on_failure_callback=lambda context: callbacks_run.append("failure"),
        on_retry_callback=lambda context: callbacks_run.append("retry"),
    )
    ti = create_runtime_ti(
        dag_id="test_defer_with_unserializable_kwargs",
        run_id="test_run",
        task=task,
        should_retry=should_retry,
    )

    log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
    context = ti.get_template_context()

    state, _, error = run(ti, context, log)

    assert state == expected_state
    assert isinstance(error, TypeError)

    context["exception"] = error
    finalize(ti, state, context, log, error)

    assert callbacks_run == ["retry" if should_retry else "failure"]


def test_handler_failure_keeps_non_retryable_exceptions_non_retryable(
    create_runtime_ti, mock_supervisor_comms
):
    """
    A handler that raises a non-retryable exception must still fail without retrying.

    `AirflowFailException` means "do not retry" wherever it is raised. Routing handler
    failures through `_handle_current_task_failed` would consult the retry count instead
    and hand back UP_FOR_RETRY, so `_handle_handler_failure` keeps the main chain's
    classification for these types.
    """

    class _FailingTrigger(DateTimeTrigger):
        def serialize(self):
            raise AirflowFailException("trigger cannot be serialized, do not retry")

    class _DeferWithFailingTrigger(BaseOperator):
        def execute(self, context):
            raise TaskDeferred(
                trigger=_FailingTrigger(moment=timezone.datetime(2024, 11, 22)),
                method_name="next",
            )

    task = _DeferWithFailingTrigger(task_id="defer_fail_exc")
    # Retries are available; the exception type must still win.
    ti = create_runtime_ti(
        dag_id="test_handler_failure_non_retryable",
        run_id="test_run",
        task=task,
        should_retry=True,
    )

    log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)

    state, msg, error = run(ti, ti.get_template_context(), log)

    assert state == TaskInstanceState.FAILED
    assert isinstance(msg, TaskState)
    assert isinstance(error, AirflowFailException)


def test_handler_failure_lets_keyboard_interrupt_propagate(create_runtime_ti, mock_supervisor_comms):
    """
    A ``KeyboardInterrupt`` inside a handler must reach ``main()``, not become a task failure.

    The supervisor's default termination signal is SIGINT, so swallowing it here would
    convert an operator-initiated kill into an ordinary retry.
    """

    class _InterruptingTrigger(DateTimeTrigger):
        def serialize(self):
            raise KeyboardInterrupt

    class _DeferWithInterrupt(BaseOperator):
        def execute(self, context):
            raise TaskDeferred(
                trigger=_InterruptingTrigger(moment=timezone.datetime(2024, 11, 22)),
                method_name="next",
            )

    task = _DeferWithInterrupt(task_id="defer_interrupt")
    ti = create_runtime_ti(dag_id="test_handler_interrupt", run_id="test_run", task=task)

    log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)

    with pytest.raises(KeyboardInterrupt):
        run(ti, ti.get_template_context(), log)


def test_handler_failure_fails_closed_when_failure_path_also_fails(create_runtime_ti, mock_supervisor_comms):
    """
    If the failure path itself raises, `run()` must still return a terminal state.

    A `retry_policy` is user code and `RetryDecision` is an unvalidated dataclass, so a
    policy handing back `retry_delay=30` (seconds, rather than a `timedelta`) reaches
    `_finalize_task_failure`, which calls `.total_seconds()` on it. That `AttributeError`
    is raised outside `_evaluate_retry_policy`'s own error handling, so it propagates out
    of the failure path. Re-entering that path with its own exception would fail the same
    way and escape `run()`, losing the callbacks this handling exists to preserve.
    """

    class _BadDelayPolicy(RetryPolicy):
        def evaluate(self, exception, try_number, max_tries, context=None):
            # Seconds as an int, not a timedelta.
            return RetryDecision(action=RetryAction.RETRY, retry_delay=30)

    class _AlwaysFails(BaseOperator):
        def execute(self, context):
            raise RuntimeError("boom")

    task = _AlwaysFails(task_id="bad_delay_policy", retry_policy=_BadDelayPolicy())
    ti = create_runtime_ti(
        dag_id="test_handler_double_fault",
        run_id="test_run",
        task=task,
        should_retry=True,
    )

    log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)

    state, msg, error = run(ti, ti.get_template_context(), log)

    assert state == TaskInstanceState.FAILED
    assert isinstance(msg, TaskState)
    assert msg.state == TaskInstanceState.FAILED
    # The terminal state must actually reach the supervisor, not merely be returned.
    assert any(call.kwargs.get("msg") is msg for call in mock_supervisor_comms.send.call_args_list)
    # `error` is the exception that actually ended the run, so a callback sees the broken
    # policy rather than a misleadingly clean task error. The task's own failure stays
    # reachable on the implicit exception chain.
    assert isinstance(error, AttributeError)
    assert isinstance(error.__context__, RuntimeError)


def test_handler_failure_counts_the_failure_once(create_runtime_ti, mock_supervisor_comms):
    """
    One failure is one increment, even when the failure path runs twice.

    The first pass through `_finalize_task_failure` records the counters before the broken
    retry delay makes it raise, so counting again on the second pass would report two
    failures for a single task run.
    """

    class _BadDelayPolicy(RetryPolicy):
        def evaluate(self, exception, try_number, max_tries, context=None):
            return RetryDecision(action=RetryAction.RETRY, retry_delay=30)

    class _AlwaysFails(BaseOperator):
        def execute(self, context):
            raise RuntimeError("boom")

    task = _AlwaysFails(task_id="count_once", retry_policy=_BadDelayPolicy())
    ti = create_runtime_ti(
        dag_id="test_handler_failure_counts_once",
        run_id="test_run",
        task=task,
        should_retry=True,
    )

    log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
    stats_backend = mock.MagicMock(spec=StatsLogger)

    with mock.patch("airflow.sdk.execution_time.task_runner.stats", stats_backend):
        run(ti, ti.get_template_context(), log)

    counted = [call.args[0] for call in stats_backend.incr.call_args_list if call.args]
    assert counted.count("ti_failures") == 1
    assert counted.count("operator_failures") == 1


def test_run_downstream_skipped(mocked_parse, create_runtime_ti, mock_supervisor_comms, listener_manager):
    listener = TestTaskRunnerCallsListeners.CustomListener()
    listener_manager(listener)

    class CustomOperator(BaseOperator):
        def execute(self, context):
            raise DownstreamTasksSkipped(tasks=["task1", "task2"])

    task = CustomOperator(
        task_id="test_task_runner_calls_listeners_skipped", do_xcom_push=True, multiple_outputs=True
    )
    ti = create_runtime_ti(task=task)

    context = ti.get_template_context()
    log = mock.MagicMock()
    run(ti, context=context, log=log)
    finalize(ti, context=context, log=mock.MagicMock(), state=TaskInstanceState.SUCCESS)

    assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.SUCCESS]
    log.info.assert_called_with("Skipping downstream tasks.")
    mock_supervisor_comms.send.assert_any_call(
        SkipDownstreamTasks(tasks=["task1", "task2"], type="SkipDownstreamTasks")
    )


def test_run_emits_endgroup_before_execute(create_runtime_ti, mock_supervisor_comms):
    """::endgroup:: is emitted to close the pre-execute log group just before execute() runs."""
    call_order: list[str] = []

    class TrackingOperator(BaseOperator):
        def execute(self, context):
            call_order.append("execute")

    task = TrackingOperator(task_id="tracking_task")
    ti = create_runtime_ti(task=task)
    log = mock.MagicMock(spec=["info", "debug", "warning", "error", "exception", "bind"])

    def tracking_info(msg, *args, **kwargs):
        if msg == "::endgroup::":
            call_order.append("::endgroup::")

    log.info.side_effect = tracking_info

    run(ti, context=ti.get_template_context(), log=log)

    assert "::endgroup::" in call_order, "::endgroup:: was never emitted"
    assert "execute" in call_order, "execute() was never called"
    assert call_order.index("::endgroup::") < call_order.index("execute")


def test_run_emits_post_execute_group_before_xcom_push(create_runtime_ti, mock_supervisor_comms):
    """::group::Post Execute is opened before 'Pushing xcom' so the push appears inside the group."""
    call_order: list[str] = []

    class ReturningOperator(BaseOperator):
        def execute(self, context):
            return "some_value"

    task = ReturningOperator(task_id="returning_task")
    ti = create_runtime_ti(task=task)
    log = mock.MagicMock(spec=["info", "debug", "warning", "error", "exception", "bind"])

    def tracking_info(msg, *args, **kwargs):
        if msg in ("::group::Post Execute", "Pushing xcom"):
            call_order.append(msg)

    log.info.side_effect = tracking_info

    run(ti, context=ti.get_template_context(), log=log)

    assert "::group::Post Execute" in call_order
    assert "Pushing xcom" in call_order
    assert call_order.index("::group::Post Execute") < call_order.index("Pushing xcom")


def test_finalize_emits_endgroup(create_runtime_ti, mock_supervisor_comms):
    """finalize() closes the post-execute log group but does not open it."""
    task = BaseOperator(task_id="some_task")
    ti = create_runtime_ti(task=task)
    log = mock.MagicMock()

    finalize(ti, state=TaskInstanceState.SUCCESS, context=ti.get_template_context(), log=log)

    info_calls = log.info.call_args_list
    assert call("::group::Post Execute") not in info_calls
    assert info_calls[-1] == call("::endgroup::")


def test_resume_from_deferred(time_machine, create_runtime_ti, mock_supervisor_comms, spy_agency: SpyAgency):
    from airflow.providers.standard.sensors.date_time import DateTimeSensorAsync

    instant_str = "2024-09-30T12:00:00Z"
    instant = timezone.parse(instant_str)
    task = DateTimeSensorAsync(
        task_id="async",
        target_time=instant + timedelta(seconds=3),
        poke_interval=60,
        timeout=600,
    )

    ti = create_runtime_ti(dag_id="basic_deferred_run", task=task)
    ti._ti_context_from_server.next_method = "execute_complete"
    ti._ti_context_from_server.next_kwargs = {
        "__type": "dict",
        "__var": {"event": {"__type": "datetime", "__var": 1727697600.0}},
    }

    spy = spy_agency.spy_on(task.execute_complete)
    state, msg, err = run(ti, context=ti.get_template_context(), log=mock.MagicMock())
    assert err is None
    assert state == TaskInstanceState.SUCCESS
    assert ti.state == TaskInstanceState.SUCCESS

    spy_agency.assert_spy_called_with(spy, mock.ANY, event=instant)


def test_run_basic_skipped(time_machine, create_runtime_ti, mock_supervisor_comms):
    """Test running a basic task that marks itself skipped."""

    task = PythonOperator(
        task_id="skip",
        python_callable=lambda: (_ for _ in ()).throw(
            AirflowSkipException("This task is being skipped intentionally."),
        ),
    )

    ti = create_runtime_ti(dag_id="basic_skipped", task=task)

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.SKIPPED

    mock_supervisor_comms.send.assert_called_with(
        TaskState(state=TaskInstanceState.SKIPPED, end_date=instant)
    )


def test_run_raises_base_exception(time_machine, create_runtime_ti, mock_supervisor_comms):
    """Test running a basic task that raises a base exception which should send fail_with_retry state."""

    task = PythonOperator(
        task_id="zero_division_error",
        python_callable=lambda: 1 / 0,
    )

    ti = create_runtime_ti(dag_id="basic_dag_base_exception", task=task)

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.FAILED

    mock_supervisor_comms.send.assert_called_with(
        msg=TaskState(
            state=TaskInstanceState.FAILED,
            end_date=instant,
        ),
    )


def test_run_raises_system_exit(time_machine, create_runtime_ti, mock_supervisor_comms):
    """Test running a basic task that exits with SystemExit exception."""

    task = PythonOperator(
        task_id="system_exit_task",
        python_callable=lambda: exit(10),
    )

    ti = create_runtime_ti(task=task, dag_id="basic_dag_system_exit")

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    log = mock.MagicMock()
    run(ti, context=ti.get_template_context(), log=log)

    assert ti.state == TaskInstanceState.FAILED

    mock_supervisor_comms.send.assert_called_with(TaskState(state=TaskInstanceState.FAILED, end_date=instant))

    log.exception.assert_not_called()
    log.error.assert_called_with(mock.ANY, exit_code=10)


def test_run_raises_airflow_exception(time_machine, create_runtime_ti, mock_supervisor_comms):
    """Test running a basic task that exits with AirflowException."""

    task = PythonOperator(
        task_id="af_exception_task",
        python_callable=lambda: (_ for _ in ()).throw(
            AirflowException("Oops! I am failing with AirflowException!"),
        ),
    )

    ti = create_runtime_ti(task=task, dag_id="basic_dag_af_exception")

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.FAILED

    mock_supervisor_comms.send.assert_called_with(TaskState(state=TaskInstanceState.FAILED, end_date=instant))


def test_run_task_timeout(time_machine, create_runtime_ti, mock_supervisor_comms):
    """Test running a basic task that times out."""
    from time import sleep

    task = PythonOperator(
        task_id="sleep",
        execution_timeout=timedelta(milliseconds=10),
        python_callable=lambda: sleep(2),
    )

    ti = create_runtime_ti(task=task, dag_id="basic_dag_time_out")

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.FAILED

    # this state can only be reached if the try block passed down the exception to handler of AirflowTaskTimeout
    mock_supervisor_comms.send.assert_called_with(TaskState(state=TaskInstanceState.FAILED, end_date=instant))


def test_execution_timeout(create_runtime_ti):
    def sleep_and_catch_other_exceptions():
        with contextlib.suppress(Exception):
            # Catching Exception should NOT catch AirflowTaskTimeout
            time.sleep(5)

    op = PythonOperator(
        task_id="test_timeout",
        execution_timeout=timedelta(seconds=1),
        python_callable=sleep_and_catch_other_exceptions,
    )

    ti = create_runtime_ti(task=op, dag_id="dag_execution_timeout")

    with pytest.raises(AirflowTaskTimeout):
        _execute_task(context=ti.get_template_context(), ti=ti, log=mock.MagicMock())


def test_basic_templated_dag(mocked_parse, make_ti_context, mock_supervisor_comms, spy_agency):
    """Test running a Dag with templated task."""
    from airflow.providers.standard.operators.bash import BashOperator

    task = BashOperator(
        task_id="templated_task",
        bash_command="echo 'Logical date is {{ logical_date }}'",
    )

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="templated_task",
            dag_id="basic_templated_dag",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        bundle_info=FAKE_BUNDLE,
        dag_rel_path="",
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )
    ti = mocked_parse(what, "basic_templated_dag", task)

    # Ensure that task is locked for execution
    spy_agency.spy_on(task.prepare_for_execution)
    assert not task._lock_for_execution

    run(ti, context=ti.get_template_context(), log=mock.Mock())

    spy_agency.assert_spy_called(task.prepare_for_execution)
    assert ti.task._lock_for_execution
    assert ti.task is not task, "ti.task should be a copy of the original task"
    assert ti.task is ti.get_template_context()["task"], "task in context should be updated too"
    assert ti.state == TaskInstanceState.SUCCESS

    mock_supervisor_comms.send.assert_any_call(
        msg=SetRenderedFields(
            rendered_fields={
                "bash_command": "echo 'Logical date is 2024-12-01 01:00:00+00:00'",
                "cwd": None,
                "env": None,
            }
        ),
    )


@pytest.mark.parametrize(
    ("task_params", "expected_rendered_fields"),
    [
        pytest.param(
            {"op_args": [], "op_kwargs": {}, "templates_dict": None},
            {"op_args": [], "op_kwargs": {}, "templates_dict": None},
            id="no_templates",
        ),
        pytest.param(
            {
                "op_args": ["arg1", "arg2", 1, 2, 3.75, {"key": "value"}],
                "op_kwargs": {"key1": "value1", "key2": 99.0, "key3": {"nested_key": "nested_value"}},
            },
            {
                "op_args": ["arg1", "arg2", 1, 2, 3.75, {"key": "value"}],
                "op_kwargs": {"key1": "value1", "key2": 99.0, "key3": {"nested_key": "nested_value"}},
            },
            id="mixed_types",
        ),
        pytest.param(
            {"my_tup": (1, 2), "my_set": {1, 2, 3}},
            {"my_tup": [1, 2], "my_set": [1, 2, 3]},
            id="tuples_and_sets",
        ),
        pytest.param(
            {"op_args": [("a", "b", "c")], "op_kwargs": {}, "templates_dict": None},
            {"op_args": [["a", "b", "c"]], "op_kwargs": {}, "templates_dict": None},
            id="nested_tuples_within_lists",
        ),
        pytest.param(
            {
                "op_args": [
                    [
                        ("t0.task_id", "t1.task_id", "branch one"),
                        ("t0.task_id", "t2.task_id", "branch two"),
                        ("t0.task_id", "t3.task_id", "branch three"),
                    ]
                ],
                "op_kwargs": {},
                "templates_dict": None,
            },
            {
                "op_args": [
                    [
                        ["t0.task_id", "t1.task_id", "branch one"],
                        ["t0.task_id", "t2.task_id", "branch two"],
                        ["t0.task_id", "t3.task_id", "branch three"],
                    ]
                ],
                "op_kwargs": {},
                "templates_dict": None,
            },
            id="nested_tuples_within_lists_higher_nesting",
        ),
    ],
)
def test_startup_and_run_dag_with_rtif(
    mocked_parse, task_params, expected_rendered_fields, make_ti_context, time_machine, mock_supervisor_comms
):
    """Test startup of a Dag with various rendered templated fields."""

    class CustomOperator(BaseOperator):
        template_fields = tuple(task_params.keys())

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            for key, value in task_params.items():
                setattr(self, key, value)

        def execute(self, context):
            for key in self.template_fields:
                print(key, getattr(self, key))

    task = CustomOperator(task_id="templated_task")
    instant = timezone.datetime(2024, 12, 3, 10, 0)

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="templated_task",
            dag_id="basic_dag",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="",
        bundle_info=FAKE_BUNDLE,
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )
    mocked_parse(what, "basic_dag", task)

    time_machine.move_to(instant, tick=False)

    mock_supervisor_comms._get_response.return_value = what

    run(*startup(get_startup_details()))
    expected_calls = [
        mock.call.send(SetRenderedFields(rendered_fields=expected_rendered_fields)),
        mock.call.send(
            msg=SucceedTask(
                end_date=instant,
                state=TaskInstanceState.SUCCESS,
                task_outlets=[],
                outlet_events=[],
            ),
        ),
    ]
    mock_supervisor_comms.assert_has_calls(expected_calls)


@patch("os.execvp")
@patch("os.set_inheritable")
def test_task_run_with_user_impersonation(
    mock_set_inheritable, mock_execvp, mocked_parse, make_ti_context, time_machine, mock_supervisor_comms
):
    class CustomOperator(BaseOperator):
        def execute(self, context):
            print("Hi from CustomOperator!")

    task = CustomOperator(task_id="impersonation_task", run_as_user="airflowuser")
    instant = timezone.datetime(2024, 12, 3, 10, 0)

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="impersonation_task",
            dag_id="basic_dag",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="",
        bundle_info=FAKE_BUNDLE,
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    mocked_parse(what, "basic_dag", task)
    time_machine.move_to(instant, tick=False)

    mock_supervisor_comms._get_response.return_value = what
    mock_supervisor_comms.socket.fileno.return_value = 42

    with mock.patch.dict(os.environ, {}, clear=True):
        startup(get_startup_details())

        assert os.environ["_AIRFLOW__REEXECUTED_PROCESS"] == "1"
        assert "_AIRFLOW__STARTUP_MSG" in os.environ

        mock_set_inheritable.assert_called_once_with(42, True)
        actual_cmd = mock_execvp.call_args.args[1]

        assert actual_cmd[:5] == ["sudo", "-E", "-H", "-u", "airflowuser"]
        assert "python" in actual_cmd[5]
        assert actual_cmd[6] == "-c"
        assert actual_cmd[7] == "from airflow.sdk.execution_time.task_runner import main; main()"


@patch("airflow.sdk.execution_time.task_runner.getuser")
def test_task_run_with_user_impersonation_default_user(
    mock_get_user, mocked_parse, make_ti_context, time_machine, mock_supervisor_comms
):
    class CustomOperator(BaseOperator):
        def execute(self, context):
            print("Hi from CustomOperator!")

    task = CustomOperator(task_id="impersonation_task", run_as_user="default_user")
    instant = timezone.datetime(2024, 12, 3, 10, 0)

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="impersonation_task",
            dag_id="basic_dag",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="",
        bundle_info=FAKE_BUNDLE,
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    mocked_parse(what, "basic_dag", task)
    time_machine.move_to(instant, tick=False)

    mock_supervisor_comms._get_response.return_value = what
    mock_supervisor_comms.socket.fileno.return_value = 42
    mock_get_user.return_value = "default_user"

    with mock.patch.dict(os.environ, {}, clear=True):
        startup(get_startup_details())

        assert "_AIRFLOW__REEXECUTED_PROCESS" not in os.environ
        assert "_AIRFLOW__STARTUP_MSG" not in os.environ


def test_task_run_with_user_impersonation_remove_krb5ccname_on_reexecuted_process(
    mocked_parse, make_ti_context, time_machine, mock_supervisor_comms
):
    class CustomOperator(BaseOperator):
        def execute(self, context):
            print("Hi from CustomOperator!")

    task = CustomOperator(task_id="impersonation_task", run_as_user="airflowuser")
    instant = timezone.datetime(2024, 12, 3, 10, 0)

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id="impersonation_task",
            dag_id="basic_dag",
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="",
        bundle_info=FAKE_BUNDLE,
        ti_context=make_ti_context(),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    mocked_parse(what, "basic_dag", task)
    time_machine.move_to(instant, tick=False)

    mock_supervisor_comms._get_response.return_value = what

    mock_os_env = {
        "KRB5CCNAME": "/tmp/airflow_krb5_ccache",
        "_AIRFLOW__REEXECUTED_PROCESS": "1",
        "_AIRFLOW__STARTUP_MSG": what.model_dump_json(),
    }
    with mock.patch.dict("os.environ", mock_os_env, clear=True):
        startup(get_startup_details())

        assert os.environ["_AIRFLOW__REEXECUTED_PROCESS"] == "1"
        assert "_AIRFLOW__STARTUP_MSG" in os.environ
        assert "KRB5CCNAME" not in os.environ


@pytest.mark.parametrize(
    ("command", "rendered_command"),
    [
        ("{{ task.task_id }}", "templated_task"),
        ("{{ run_id }}", "c"),
        ("{{ logical_date }}", "2024-12-01 01:00:00+00:00"),
    ],
)
@pytest.mark.usefixtures("mock_supervisor_comms")
def test_startup_and_run_dag_with_templated_fields(
    command, rendered_command, create_runtime_ti, time_machine
):
    """Test startup of a Dag with various templated fields."""
    from airflow.providers.standard.operators.bash import BashOperator

    task = BashOperator(task_id="templated_task", bash_command=command)

    ti = create_runtime_ti(
        task=task, dag_id="basic_dag", logical_date="2024-12-01 01:00:00+00:00", run_id="c"
    )

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)
    run(ti, context=ti.get_template_context(), log=mock.MagicMock())
    assert ti.task.bash_command == rendered_command


def test_get_context_in_task(create_runtime_ti, time_machine, mock_supervisor_comms):
    """Test that the `get_current_context` & `set_current_context` work correctly."""

    class MyContextAssertOperator(BaseOperator):
        def execute(self, context):
            # Ensure the context returned by get_current_context is the same as the
            # context passed to the operator
            assert context == get_current_context()

    task = MyContextAssertOperator(task_id="assert_context")

    ti = create_runtime_ti(task=task)

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.SUCCESS

    # Ensure the task is Successful
    mock_supervisor_comms.send.assert_called_once_with(
        msg=SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
    )


@pytest.mark.parametrize(
    ("dag_id", "task_id", "fail_with_exception"),
    [
        pytest.param(
            "basic_failed", "fail-exception", AirflowFailException("Oops. Failing by AirflowFailException!")
        ),
        pytest.param(
            "basic_failed2",
            "sensor-timeout-exception",
            AirflowSensorTimeout("Oops. Failing by AirflowSensorTimeout!"),
        ),
        pytest.param(
            "basic_failed3",
            "task-terminated-exception",
            AirflowTaskTerminated("Oops. Failing by AirflowTaskTerminated!"),
        ),
    ],
)
def test_run_basic_failed(
    time_machine, create_runtime_ti, dag_id, task_id, fail_with_exception, mock_supervisor_comms
):
    """Test running a basic task that marks itself as failed by raising exception."""

    class CustomOperator(BaseOperator):
        def __init__(self, e, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.e = e

        def execute(self, context):
            print(f"raising exception {self.e}")
            raise self.e

    task = CustomOperator(task_id=task_id, e=fail_with_exception)

    ti = create_runtime_ti(task=task, dag_id=dag_id)

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.FAILED

    mock_supervisor_comms.send.assert_called_once_with(
        msg=TaskState(state=TaskInstanceState.FAILED, end_date=instant)
    )


@pytest.mark.parametrize("retries", [0, 1])
def test_airflow_fail_exception_does_not_retry(
    time_machine, create_runtime_ti, mock_supervisor_comms, retries
):
    """Test that AirflowFailException does not cause task to retry."""

    def fail():
        raise AirflowFailException("hopeless")

    task = PythonOperator(
        task_id="test_raise_airflow_fail_exception",
        python_callable=fail,
        retries=retries,
    )

    ti = create_runtime_ti(task=task, dag_id=f"test_airflow_fail_exception_no_retry_{retries}")

    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    assert ti.state == TaskInstanceState.FAILED
    mock_supervisor_comms.send.assert_called_with(
        msg=TaskState(state=TaskInstanceState.FAILED, end_date=instant)
    )


def test_dag_parsing_context(make_ti_context, mock_supervisor_comms, monkeypatch, test_dags_dir):
    """
    Test that the Dag parsing context is correctly set during the startup process.

    This test verifies that the Dag and task IDs are correctly set in the parsing context
    when a Dag is started up.
    """
    dag_id = "dag_parsing_context_test"
    task_id = "conditional_task"

    what = StartupDetails(
        ti=TaskInstance(
            id=uuid7(),
            task_id=task_id,
            dag_id=dag_id,
            run_id="c",
            try_number=1,
            dag_version_id=uuid7(),
            queue="default",
        ),
        dag_rel_path="dag_parsing_context.py",
        bundle_info=BundleInfo(name="my-bundle", version=None),
        ti_context=make_ti_context(dag_id=dag_id, run_id="c"),
        start_date=timezone.utcnow(),
        sentry_integration="",
    )

    mock_supervisor_comms._get_response.return_value = what

    # Set the environment variable for Dag bundles
    # We use the Dag defined in `task_sdk/tests/dags/dag_parsing_context.py` for this test!
    dag_bundle_val = json.dumps(
        [
            {
                "name": "my-bundle",
                "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle",
                "kwargs": {"path": str(test_dags_dir), "refresh_interval": 1},
            }
        ]
    )

    monkeypatch.setenv("AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST", dag_bundle_val)
    ti, _, _ = startup(get_startup_details())

    # Presence of `conditional_task` below means Dag ID is properly set in the parsing context!
    # Check the dag file for the actual logic!
    assert ti.task.dag.task_dict.keys() == {"visible_task", "conditional_task"}


@pytest.mark.parametrize(
    ("task_outlets", "expected_msg"),
    [
        pytest.param(
            [Asset(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
            SucceedTask(
                state="success",
                end_date=timezone.datetime(2024, 12, 3, 10, 0),
                task_outlets=[
                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", type="Asset")
                ],
                outlet_events=[],
            ),
            id="asset",
        ),
        pytest.param(
            [Dataset(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
            SucceedTask(
                state="success",
                end_date=timezone.datetime(2024, 12, 3, 10, 0),
                task_outlets=[
                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", type="Asset")
                ],
                outlet_events=[],
            ),
            id="dataset",
        ),
        pytest.param(
            [Model(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
            SucceedTask(
                state="success",
                end_date=timezone.datetime(2024, 12, 3, 10, 0),
                task_outlets=[
                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", type="Asset")
                ],
                outlet_events=[],
            ),
            id="model",
        ),
        pytest.param(
            [Asset.ref(name="s3://bucket/my-task")],
            SucceedTask(
                state="success",
                end_date=timezone.datetime(2024, 12, 3, 10, 0),
                task_outlets=[AssetProfile(name="s3://bucket/my-task", type="AssetNameRef")],
                outlet_events=[],
            ),
            id="name-ref",
        ),
        pytest.param(
            [Asset.ref(uri="s3://bucket/my-task")],
            SucceedTask(
                state="success",
                end_date=timezone.datetime(2024, 12, 3, 10, 0),
                task_outlets=[AssetProfile(uri="s3://bucket/my-task", type="AssetUriRef")],
                outlet_events=[],
            ),
            id="uri-ref",
        ),
        pytest.param(
            [AssetAlias(name="example-alias", group="asset")],
            SucceedTask(
                state="success",
                end_date=timezone.datetime(2024, 12, 3, 10, 0),
                task_outlets=[AssetProfile(name="example-alias", type="AssetAlias")],
                outlet_events=[],
            ),
            id="asset-alias",
        ),
    ],
)
def test_run_with_asset_outlets(
    time_machine, create_runtime_ti, mock_supervisor_comms, task_outlets, expected_msg
):
    """Test running a basic task that contains asset outlets."""
    from airflow.providers.standard.operators.bash import BashOperator

    task = BashOperator(
        outlets=task_outlets,
        task_id="asset-outlet-task",
        bash_command="echo 'hi'",
    )

    ti = create_runtime_ti(task=task, dag_id="dag_with_asset_outlet_task")
    instant = timezone.datetime(2024, 12, 3, 10, 0)
    time_machine.move_to(instant, tick=False)

    with mock.patch(
        "airflow.sdk.execution_time.task_runner._validate_task_inlets_and_outlets"
    ) as validate_mock:
        run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    validate_mock.assert_called_once()

    mock_supervisor_comms.send.assert_any_call(expected_msg)


def test_run_with_asset_inlets(create_runtime_ti, mock_supervisor_comms):
    """Test running a basic task that contains asset inlets."""
    asset_event_resp = AssetEventResult(
        id=1,
        created_dagruns=[],
        timestamp=timezone.utcnow(),
        asset=AssetResponse(name="test", uri="test", group="asset"),
    )
    events_result = AssetEventsResult(asset_events=[asset_event_resp])
    mock_supervisor_comms.send.return_value = events_result

    from airflow.providers.standard.operators.bash import BashOperator

    task = BashOperator(
        inlets=[Asset(name="test", uri="test://uri"), AssetAlias(name="alias-name")],
        task_id="asset-outlet-task",
        bash_command="echo 0",
    )

    ti = create_runtime_ti(task=task, dag_id="dag_with_asset_outlet_task")
    run(ti, context=ti.get_template_context(), log=mock.MagicMock())
    inlet_events = ti.get_template_context()["inlet_events"]

    # access the asset events of Asset(name="test", uri="test://uri")
    assert list(inlet_events[0]) == [asset_event_resp]
    assert list(inlet_events[-2]) == [asset_event_resp]
    assert list(inlet_events[Asset(name="test", uri="test://uri")]) == [asset_event_resp]

    # access the asset events of AssetAlias(name="alias-name")
    assert list(inlet_events[1]) == [asset_event_resp]
    assert list(inlet_events[-1]) == [asset_event_resp]
    assert list(inlet_events[AssetAlias(name="alias-name")]) == [asset_event_resp]

    # access with invalid index
    with pytest.raises(IndexError):
        inlet_events[2]

    with pytest.raises(IndexError):
        inlet_events[-3]

    with pytest.raises(KeyError):
        inlet_events[Asset(name="no such asset in inlets")]


@mock.patch("airflow.sdk.execution_time.task_runner.context_to_airflow_vars")
@mock.patch.dict(os.environ, {}, clear=True)
def test_execute_task_exports_env_vars(
    mock_context_to_airflow_vars, create_runtime_ti, mock_supervisor_comms
):
    """Test that _execute_task exports airflow context to environment variables."""

    def test_function():
        return "test function"

    task = PythonOperator(
        task_id="test_task",
        python_callable=test_function,
    )

    ti = create_runtime_ti(task=task, dag_id="dag_with_env_vars")

    mock_env_vars = {"AIRFLOW_CTX_DAG_ID": "test_dag_env_vars", "AIRFLOW_CTX_TASK_ID": "test_env_task"}
    mock_context_to_airflow_vars.return_value = mock_env_vars
    run(ti, ti.get_template_context(), log=mock.MagicMock())

    assert os.environ["AIRFLOW_CTX_DAG_ID"] == "test_dag_env_vars"
    assert os.environ["AIRFLOW_CTX_TASK_ID"] == "test_env_task"


def test_execute_success_task_with_rendered_map_index(create_runtime_ti, mock_supervisor_comms):
    """Test that the map index is rendered in the task context."""

    def test_function():
        return "test function"

    task = PythonOperator(
        task_id="test_task",
        python_callable=test_function,
        map_index_template="Hello! {{ run_id }}",
    )

    ti = create_runtime_ti(task=task, dag_id="dag_with_map_index_template")

    run(ti, ti.get_template_context(), log=mock.MagicMock())

    assert ti.rendered_map_index == "Hello! test_run"


def test_execute_failed_task_with_rendered_map_index(create_runtime_ti, mock_supervisor_comms):
    """Test that the map index is rendered in the task context."""

    task = BaseOperator(task_id="test_task", map_index_template="Hello! {{ run_id }}")

    ti = create_runtime_ti(task=task, dag_id="dag_with_map_index_template")

    run(ti, ti.get_template_context(), log=mock.MagicMock())

    assert ti.rendered_map_index == "Hello! test_run"


def test_rendered_map_index_updates_sent_progressively(create_runtime_ti, mock_supervisor_comms):
    """Test that rendered_map_index is rendered and potentially updated after execution."""

    def test_function(ti):
        # Simulate setting a context variable during execution
        ti.xcom_push(key="execution_result", value="completed")
        return "test function"

    task = PythonOperator(
        task_id="test_task",
        python_callable=test_function,
        map_index_template="Label: {{ task.task_id }}",
    )

    ti = create_runtime_ti(task=task, dag_id="dag_with_progressive_map_index")

    run(ti, ti.get_template_context(), log=mock.MagicMock())

    # Verify that rendered_map_index is set (existing behavior)
    assert ti.rendered_map_index == "Label: test_task"


class TestSerializeOutletEvents:
    """Tests for the wire format produced by ``_serialize_outlet_events``."""

    def test_emits_single_event_when_no_partition_keys(self):
        accessors = OutletEventAccessors()
        accessors[Asset(name="a")].extra = {"x": 1}

        events = list(_serialize_outlet_events(accessors))

        assert events == [{"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {"x": 1}}]

    def test_emits_one_event_per_partition_key(self):
        accessors = OutletEventAccessors()
        accessors[Asset(name="a")].partition_keys = {"us", "eu"}

        events = list(_serialize_outlet_events(accessors))

        assert sorted(events, key=lambda e: e["partition_key"]) == [
            {"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {}, "partition_key": "eu"},
            {"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {}, "partition_key": "us"},
        ]

    def test_dedupes_partition_keys_at_serialization(self):
        accessors = OutletEventAccessors()
        accessor = accessors[Asset(name="a")]
        accessor.add_partitions("us")
        accessor.add_partitions("us")
        accessor.add_partitions(["us", "eu"])

        events = list(_serialize_outlet_events(accessors))

        assert sorted(events, key=lambda e: e["partition_key"]) == [
            {"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {}, "partition_key": "eu"},
            {"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {}, "partition_key": "us"},
        ]


class TestRuntimeTaskInstance:
    def test_get_context_without_ti_context_from_server(self, mocked_parse, make_ti_context):
        """Test get_template_context without ti_context_from_server."""

        task = BaseOperator(task_id="hello")
        dag_id = "basic_task"

        # Assign task to Dag
        get_inline_dag(dag_id=dag_id, task=task)

        ti_id = uuid7()
        ti = TaskInstance(
            id=ti_id,
            task_id=task.task_id,
            dag_id=dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )
        start_date = timezone.datetime(2025, 1, 1)

        # Keep the context empty
        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True),
            task=task,
            _ti_context_from_server=None,
            start_date=start_date,
        )
        context = runtime_ti.get_template_context()

        # Verify the context keys and values
        assert context == {
            "params": {},
            "var": {
                "json": VariableAccessor(deserialize_json=True),
                "value": VariableAccessor(deserialize_json=False),
            },
            "conn": ConnectionAccessor(),
            "dag": runtime_ti.task.dag,
            "inlets": task.inlets,
            "inlet_events": InletEventsAccessors(inlets=[]),
            "macros": MacrosAccessor(),
            "map_index_template": task.map_index_template,
            "outlet_events": OutletEventAccessors(),
            "outlets": task.outlets,
            "run_id": "test_run",
            "task": task,
            "task_instance": runtime_ti,
            "task_state_store": TaskStateStoreAccessor(
                ti_id=ti_id, scope=TaskScope(dag_id=dag_id, run_id="test_run", task_id="hello")
            ),
            "ti": runtime_ti,
        }

    def test_get_context_with_ti_context_from_server(self, create_runtime_ti, mock_supervisor_comms):
        """Test the context keys are added when sent from API server (mocked)"""

        task = BaseOperator(task_id="hello")

        # Assume the context is sent from the API server
        # `task-sdk/tests/api/test_client.py::test_task_instance_start` checks the context is received
        # from the API server
        runtime_ti = create_runtime_ti(task=task, dag_id="basic_task")

        dr = runtime_ti._ti_context_from_server.dag_run

        mock_supervisor_comms.send.return_value = PrevSuccessfulDagRunResult(
            data_interval_end=dr.logical_date - timedelta(hours=1),
            data_interval_start=dr.logical_date - timedelta(hours=2),
            start_date=dr.start_date - timedelta(hours=1),
            end_date=dr.start_date,
        )

        context = runtime_ti.get_template_context()

        assert context == {
            "params": {},
            "var": {
                "json": VariableAccessor(deserialize_json=True),
                "value": VariableAccessor(deserialize_json=False),
            },
            "conn": ConnectionAccessor(),
            "dag": runtime_ti.task.dag,
            "inlets": task.inlets,
            "inlet_events": InletEventsAccessors(inlets=[]),
            "macros": MacrosAccessor(),
            "map_index_template": task.map_index_template,
            "outlet_events": OutletEventAccessors(),
            "outlets": task.outlets,
            "prev_data_interval_end_success": timezone.datetime(2024, 12, 1, 0, 0, 0),
            "prev_data_interval_start_success": timezone.datetime(2024, 11, 30, 23, 0, 0),
            "prev_end_date_success": timezone.datetime(2024, 12, 1, 1, 0, 0),
            "prev_start_date_success": timezone.datetime(2024, 12, 1, 0, 0, 0),
            "run_id": "test_run",
            "task": task,
            "task_instance": runtime_ti,
            "task_state_store": TaskStateStoreAccessor(
                ti_id=runtime_ti.id,
                scope=TaskScope(dag_id=runtime_ti.dag_id, run_id="test_run", task_id="hello"),
            ),
            "ti": runtime_ti,
            "dag_run": dr,
            "data_interval_end": timezone.datetime(2024, 12, 1, 1, 0, 0),
            "data_interval_start": timezone.datetime(2024, 12, 1, 1, 0, 0),
            "logical_date": timezone.datetime(2024, 12, 1, 1, 0, 0),
            "task_reschedule_count": 0,
            "triggering_asset_events": TriggeringAssetEventsAccessor.build(dr.consumed_asset_events),
            "ds": "2024-12-01",
            "ds_nodash": "20241201",
            "task_instance_key_str": "basic_task__hello__20241201",
            "ts": "2024-12-01T01:00:00+00:00",
            "ts_nodash": "20241201T010000",
            "ts_nodash_with_tz": "20241201T010000+0000",
            "partition_key": dr.partition_key,
            "partition_date": dr.partition_date,
        }

    def test_partition_key_in_context(self, create_runtime_ti, mock_supervisor_comms):
        """Test that partition_key from dag_run is exposed in the template context."""
        task = BaseOperator(task_id="hello")
        runtime_ti = create_runtime_ti(task=task, dag_id="basic_task")

        dr = runtime_ti._ti_context_from_server.dag_run

        mock_supervisor_comms.send.return_value = PrevSuccessfulDagRunResult(
            data_interval_end=dr.logical_date - timedelta(hours=1),
            data_interval_start=dr.logical_date - timedelta(hours=2),
            start_date=dr.start_date - timedelta(hours=1),
            end_date=dr.start_date,
        )

        context = runtime_ti.get_template_context()

        # Default: partition_key is None
        assert context["partition_key"] is None

        # Set partition_key on dag_run and verify it surfaces in context
        dr.partition_key = "some-partition"
        context = runtime_ti.get_template_context()
        assert context["partition_key"] == "some-partition"

    def test_partition_date_in_context(self, create_runtime_ti, mock_supervisor_comms):
        """Test that partition_date from dag_run is exposed in the template context."""
        task = BaseOperator(task_id="hello")
        runtime_ti = create_runtime_ti(task=task, dag_id="basic_task")

        dr = runtime_ti._ti_context_from_server.dag_run

        mock_supervisor_comms.send.return_value = PrevSuccessfulDagRunResult(
            data_interval_end=dr.logical_date - timedelta(hours=1),
            data_interval_start=dr.logical_date - timedelta(hours=2),
            start_date=dr.start_date - timedelta(hours=1),
            end_date=dr.start_date,
        )

        context = runtime_ti.get_template_context()

        # Default: partition_date is None
        assert context["partition_date"] is None

        # Set partition_date on dag_run and verify it surfaces in context
        partition_date = timezone.datetime(2026, 5, 20, 1, 0, 0)
        dr.partition_date = partition_date
        context = runtime_ti.get_template_context()
        assert context["partition_date"] == partition_date

        # Naive datetime is coerced to tz-aware so Jinja `| ds` / `| ts` filters
        # operate on a real awareness boundary.
        dr.partition_date = datetime(2026, 5, 20, 1, 0, 0)
        context = runtime_ti.get_template_context()
        assert context["partition_date"].tzinfo is not None

    def test_lazy_loading_not_triggered_until_accessed(self, create_runtime_ti, mock_supervisor_comms):
        """Ensure lazy-loaded attributes are not resolved until accessed."""
        task = BaseOperator(task_id="hello")
        runtime_ti = create_runtime_ti(task=task, dag_id="basic_task")

        mock_supervisor_comms.send.return_value = PrevSuccessfulDagRunResult(
            data_interval_end=timezone.datetime(2025, 1, 1, 2, 0, 0),
            data_interval_start=timezone.datetime(2025, 1, 1, 1, 0, 0),
            start_date=timezone.datetime(2025, 1, 1, 1, 0, 0),
            end_date=timezone.datetime(2025, 1, 1, 2, 0, 0),
        )

        context = runtime_ti.get_template_context()

        # Assert lazy attributes are not resolved initially
        mock_supervisor_comms.send.assert_not_called()

        # Access a lazy-loaded attribute to trigger computation
        assert context["prev_data_interval_start_success"] == timezone.datetime(2025, 1, 1, 1, 0, 0)

        # Now the lazy attribute should trigger the call
        mock_supervisor_comms.send.assert_called_once()

    def test_get_connection_from_context(self, create_runtime_ti, mock_supervisor_comms):
        """Test that the connection is fetched from the API server via the Supervisor lazily when accessed"""

        task = BaseOperator(task_id="hello")

        conn = ConnectionResult(
            conn_id="test_conn",
            conn_type="mysql",
            host="mysql",
            schema="airflow",
            login="root",
            password="password",
            port=1234,
            extra='{"extra_key": "extra_value"}',
        )

        runtime_ti = create_runtime_ti(task=task, dag_id="test_get_connection_from_context")
        mock_supervisor_comms.send.return_value = conn

        context = runtime_ti.get_template_context()

        # Assert that the connection is not fetched from the API server yet!
        # The connection should be only fetched connection is accessed
        mock_supervisor_comms.send.assert_not_called()

        # Access the connection from the context
        conn_from_context = context["conn"].test_conn

        mock_supervisor_comms.send.assert_any_call(GetConnection(conn_id="test_conn"))

        assert conn_from_context == Connection(
            conn_id="test_conn",
            conn_type="mysql",
            description=None,
            host="mysql",
            schema="airflow",
            login="root",
            password="password",
            port=1234,
            extra='{"extra_key": "extra_value"}',
        )

        dejson_from_conn = conn_from_context.extra_dejson
        assert dejson_from_conn == {"extra_key": "extra_value"}

    def test_template_render(self, create_runtime_ti):
        task = BaseOperator(task_id="test_template_render_task")

        runtime_ti = create_runtime_ti(task=task, dag_id="test_template_render")
        template_context = runtime_ti.get_template_context()
        result = runtime_ti.task.render_template(
            "Task: {{ dag.dag_id }} -> {{ task.task_id }}", template_context
        )
        assert result == "Task: test_template_render -> test_template_render_task"

    @pytest.mark.parametrize(
        ("content", "expected_output"),
        [
            ('{{ conn.get("a_connection").host }}', "hostvalue"),
            ('{{ conn.get("a_connection", "unused_fallback").host }}', "hostvalue"),
            ("{{ conn.a_connection.host }}", "hostvalue"),
            ("{{ conn.a_connection.login }}", "loginvalue"),
            ("{{ conn.a_connection.schema }}", "schemavalues"),
            ("{{ conn.a_connection.password }}", "passwordvalue"),
            ('{{ conn.a_connection.extra_dejson["extra__asana__workspace"] }}', "extra1"),
            ("{{ conn.a_connection.extra_dejson.extra__asana__workspace }}", "extra1"),
        ],
    )
    def test_template_with_connection(
        self, content, expected_output, create_runtime_ti, mock_supervisor_comms
    ):
        """
        Test the availability of connections in templates
        """
        task = BaseOperator(task_id="hello")
        runtime_ti = create_runtime_ti(task=task, dag_id="test_template_with_connection")

        conn = ConnectionResult(
            conn_id="a_connection",
            conn_type="a_type",
            host="hostvalue",
            login="loginvalue",
            password="passwordvalue",
            schema="schemavalues",
            extra='{"extra__asana__workspace": "extra1"}',
            port=None,
        )

        mock_supervisor_comms.send.return_value = conn

        context = runtime_ti.get_template_context()
        result = runtime_ti.task.render_template(content, context)
        assert result == expected_output

    @pytest.mark.parametrize(
        ("accessor_type", "var_value", "expected_value"),
        [
            pytest.param("value", "test_value", "test_value"),
            pytest.param(
                "json",
                '{\r\n  "key1": "value1",\r\n  "key2": "value2",\r\n  "enabled": true,\r\n  "threshold": 42\r\n}',
                {"key1": "value1", "key2": "value2", "enabled": True, "threshold": 42},
            ),
        ],
    )
    def test_get_variable_from_context(
        self, create_runtime_ti, mock_supervisor_comms, accessor_type, var_value: str, expected_value
    ):
        """Test that the variable is fetched from the API server via the Supervisor lazily when accessed"""

        task = BaseOperator(task_id="hello")
        runtime_ti = create_runtime_ti(task=task)

        var = VariableResult(key="test_key", value=var_value)

        mock_supervisor_comms.send.return_value = var

        context = runtime_ti.get_template_context()

        # Assert that the variable is not fetched from the API server yet!
        # The variable should be only fetched connection is accessed
        mock_supervisor_comms.send.assert_not_called()

        # Access the variable from the context
        var_from_context = context["var"][accessor_type].test_key

        mock_supervisor_comms.send.assert_any_call(GetVariable(key="test_key"))

        assert var_from_context == expected_value

    @pytest.mark.parametrize(
        "map_indexes",
        [
            pytest.param(-1, id="not_mapped_index"),
            pytest.param(1, id="single_map_index"),
            pytest.param([0, 1], id="multiple_map_indexes"),
            pytest.param((0, 1), id="any_iterable_multi_indexes"),
            pytest.param(None, id="index_none"),
            pytest.param(NOTSET, id="index_not_set"),
        ],
    )
    @pytest.mark.parametrize(
        "task_ids",
        [
            pytest.param("push_task", id="single_task"),
            pytest.param(["push_task1", "push_task2"], id="tid_multiple_tasks"),
            pytest.param({"push_task1", "push_task2"}, id="tid_any_iterable"),
            pytest.param(None, id="tid_none"),
            pytest.param(NOTSET, id="tid_not_set"),
        ],
    )
    @pytest.mark.parametrize(
        "xcom_values",
        [
            pytest.param("hello", id="string_value"),
            pytest.param("'hello'", id="quoted_string_value"),
            pytest.param({"key": "value"}, id="json_value"),
            pytest.param([], id="empty_list_no_xcoms_found"),
            pytest.param((1, 2, 3), id="tuple_int_value"),
            pytest.param([1, 2, 3], id="list_int_value"),
            pytest.param(42, id="int_value"),
            pytest.param(True, id="boolean_value"),
            pytest.param(pd.DataFrame({"col1": [1, 2], "col2": [3, 4]}), id="dataframe_value"),
        ],
    )
    def test_xcom_pull(
        self,
        create_runtime_ti,
        mock_supervisor_comms,
        spy_agency,
        xcom_values,
        task_ids,
        map_indexes,
    ):
        """
        Test that a task makes an expected call to the Supervisor to pull XCom values
        based on various task_ids, map_indexes, and xcom_values configurations.
        """
        map_indexes_kwarg = {} if map_indexes is NOTSET else {"map_indexes": map_indexes}
        task_ids_kwarg = {} if task_ids is NOTSET else {"task_ids": task_ids}
        spy_agency.spy_on(deserialize)

        class CustomOperator(BaseOperator):
            def execute(self, context):
                value = context["ti"].xcom_pull(key="key", **task_ids_kwarg, **map_indexes_kwarg)
                print(f"Pulled XCom Value: {value}")

        test_task_id = "pull_task"
        task = CustomOperator(task_id=test_task_id)

        # In case of the specific map_index we should check it is passed to TI.
        # ``None`` is not a valid TaskInstance.map_index value, but xcom_pull's
        # behaviour with ``map_indexes=None`` is independent of the TI's own map_index.
        extra_for_ti = {"map_index": map_indexes} if isinstance(map_indexes, int) else {}
        runtime_ti = create_runtime_ti(task=task, **extra_for_ti)

        ser_value = BaseXCom.serialize_value(xcom_values)

        def mock_send_side_effect(*args, **kwargs):
            msg = kwargs.get("msg") or args[0]
            print(f"{args=}, {kwargs=}, {msg=}")
            if isinstance(msg, GetXComSequenceSlice):
                return XComSequenceSliceResult(root=[ser_value])
            return XComResult(key="key", value=ser_value)

        mock_supervisor_comms.send.side_effect = mock_send_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
        spy_agency.assert_spy_called_with(deserialize, ser_value)

        if not isinstance(task_ids, Iterable) or isinstance(task_ids, str):
            task_ids = [task_ids]

        if not isinstance(map_indexes, Iterable):
            map_indexes = [map_indexes]

        for task_id_raw in task_ids:
            # Without task_ids (or None) expected behavior is to pull with calling task_id
            task_id = task_id_raw if is_arg_set(task_id_raw) and task_id_raw is not None else test_task_id

            for map_index in map_indexes:
                if map_index == NOTSET:
                    mock_supervisor_comms.send.assert_any_call(
                        msg=GetXComSequenceSlice(
                            key="key",
                            dag_id="test_dag",
                            run_id="test_run",
                            task_id=task_id,
                            start=None,
                            stop=None,
                            step=None,
                        ),
                    )
                else:
                    expected_map_index = map_index if map_index is not None else None
                    mock_supervisor_comms.send.assert_any_call(
                        msg=GetXCom(
                            key="key",
                            dag_id="test_dag",
                            run_id="test_run",
                            task_id=task_id,
                            map_index=expected_map_index,
                        ),
                    )

    @pytest.mark.asyncio
    async def test_axcom_pull(
        self,
        create_runtime_ti,
        mock_supervisor_comms,
        spy_agency,
    ):
        """
        Test that a task makes an expected call to the Supervisor to pull an XCom value asynchronously.
        """
        spy_agency.spy_on(deserialize)

        task = BaseOperator(task_id="pull_task")
        runtime_ti = create_runtime_ti(task=task)

        xcom_value = "hello"
        ser_value = BaseXCom.serialize_value(xcom_value)
        mock_supervisor_comms.asend.return_value = XComResult(key="key", value=ser_value)

        value = await runtime_ti.axcom_pull(key="key", task_ids="push_task", map_indexes=-1)

        assert value == xcom_value
        spy_agency.assert_spy_called_with(deserialize, ser_value)
        mock_supervisor_comms.asend.assert_called_once_with(
            GetXCom(
                key="key",
                dag_id="test_dag",
                run_id="test_run",
                task_id="push_task",
                map_index=-1,
            ),
        )

    @pytest.mark.parametrize(
        ("task_ids", "map_indexes", "expected_value"),
        [
            pytest.param("task_a", 0, {"a": 1, "b": 2}, id="task_id is str, map_index is int"),
            pytest.param("task_a", [0], [{"a": 1, "b": 2}], id="task_id is str, map_index is list"),
            pytest.param("task_a", None, {"a": 1, "b": 2}, id="task_id is str, map_index is None"),
            pytest.param(["task_a"], 0, [{"a": 1, "b": 2}], id="task_id is list, map_index is int"),
            pytest.param(["task_a"], [0], [{"a": 1, "b": 2}], id="task_id is list, map_index is list"),
            pytest.param(["task_a"], None, [{"a": 1, "b": 2}], id="task_id is list, map_index is None"),
            pytest.param(
                ["task_a"], NOTSET, [{"a": 1, "b": 2}], id="task_id is list, map_index is ArgNotSet"
            ),
            pytest.param(None, 0, {"a": 1, "b": 2}, id="task_id is None, map_index is int"),
            pytest.param(None, [0], [{"a": 1, "b": 2}], id="task_id is None, map_index is list"),
            pytest.param(None, None, {"a": 1, "b": 2}, id="task_id is None, map_index is None"),
            pytest.param(
                ["task_a", "task_b"],
                NOTSET,
                [{"a": 1, "b": 2}, {"c": 3, "d": 4}],
                id="multiple task_ids, map_index is ArgNotSet",
            ),
            pytest.param("task_a", NOTSET, {"a": 1, "b": 2}, id="task_id is str, map_index is ArgNotSet"),
            pytest.param(None, NOTSET, {"a": 1, "b": 2}, id="task_id is None, map_index is ArgNotSet"),
        ],
    )
    def test_xcom_pull_return_values(
        self,
        create_runtime_ti,
        mock_supervisor_comms,
        task_ids,
        map_indexes,
        expected_value,
    ):
        """
        Tests return value of xcom_pull under various combinations of task_ids and map_indexes.
        Also verifies the correct XCom method (get_one vs get_all) is called.
        """

        class CustomOperator(BaseOperator):
            def execute(self, context):
                print("This is a custom operator")

        test_task_id = "pull_task"
        task = CustomOperator(task_id=test_task_id)
        runtime_ti = create_runtime_ti(task=task)

        with patch.object(XCom, "get_one") as mock_get_one, patch.object(XCom, "get_all") as mock_get_all:
            if map_indexes == NOTSET:
                # Use side_effect to return different values for different tasks
                def mock_get_all_side_effect(task_id, **kwargs):
                    if task_id == "task_b":
                        return [{"c": 3, "d": 4}]
                    return [{"a": 1, "b": 2}]

                mock_get_all.side_effect = mock_get_all_side_effect
                mock_get_one.return_value = None
            else:
                mock_get_one.return_value = {"a": 1, "b": 2}
                mock_get_all.return_value = None

            xcom = runtime_ti.xcom_pull(key="key", task_ids=task_ids, map_indexes=map_indexes)
            assert xcom == expected_value
            if map_indexes == NOTSET:
                assert mock_get_all.called
                assert not mock_get_one.called
            else:
                assert mock_get_one.called
                assert not mock_get_all.called

    @pytest.mark.parametrize(
        "api_return_value",
        [
            pytest.param(("data", "test_value"), id="api returns tuple"),
            pytest.param({"data": "test_value"}, id="api returns dict"),
            pytest.param(None, id="api returns None, no xcom found"),
        ],
    )
    def test_xcom_pull_with_no_map_index(
        self,
        api_return_value,
        create_runtime_ti,
        mock_supervisor_comms,
    ):
        """
        Test xcom_pull when map_indexes is not specified, so that XCom.get_all is called.
        The test also tests if the response is deserialized and returned.
        """
        test_task_id = "pull_task"
        task = BaseOperator(task_id=test_task_id)
        runtime_ti = create_runtime_ti(task=task)

        ser_value = BaseXCom.serialize_value(api_return_value)

        def mock_send_side_effect(*args, **kwargs):
            msg = kwargs.get("msg") or args[0]
            if isinstance(msg, GetXComSequenceSlice):
                return XComSequenceSliceResult(root=[ser_value])
            return XComResult(key="test_key", value=None)

        mock_supervisor_comms.send.side_effect = mock_send_side_effect
        result = runtime_ti.xcom_pull(key="test_key", task_ids="task_a")

        # if the API returns a tuple or dict, the below assertion assures that the value is deserialized correctly by XCom.get_all
        assert result == api_return_value

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetXComSequenceSlice(
                key="test_key",
                dag_id=runtime_ti.dag_id,
                run_id=runtime_ti.run_id,
                task_id="task_a",
                start=None,
                stop=None,
                step=None,
            ),
        )

    def test_get_param_from_context(
        self, mocked_parse, make_ti_context, mock_supervisor_comms, create_runtime_ti
    ):
        """Test that a params can be retrieved from context."""

        class CustomOperator(BaseOperator):
            def execute(self, context):
                value = context["params"]
                print("The dag params are", value)

        task = CustomOperator(task_id="print-params")
        runtime_ti = create_runtime_ti(
            dag_id="basic_param_dag",
            task=task,
            conf={
                "x": 3,
                "text": "Hello World!",
                "flag": False,
                "a_simple_list": ["one", "two", "three", "actually one value is made per line"],
            },
        )
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        assert runtime_ti.task.dag.params == {
            "x": 3,
            "text": "Hello World!",
            "flag": False,
            "a_simple_list": ["one", "two", "three", "actually one value is made per line"],
        }

    @pytest.mark.parametrize(
        ("logical_date", "check"),
        (
            pytest.param(None, pytest.raises(KeyError), id="no-logical-date"),
            pytest.param(timezone.datetime(2024, 12, 3), contextlib.nullcontext(), id="with-logical-date"),
        ),
    )
    def test_no_logical_date_key_error(
        self, mocked_parse, make_ti_context, mock_supervisor_comms, create_runtime_ti, logical_date, check
    ):
        """Test that a params can be retrieved from context."""

        class CustomOperator(BaseOperator):
            def execute(self, context):
                for key in ("ds", "ds_nodash", "ts", "ts_nodash", "ts_nodash_with_tz"):
                    with check:
                        context[key]
                # We should always be able to get this
                assert context["task_instance_key_str"]

        task = CustomOperator(task_id="print-params")
        runtime_ti = create_runtime_ti(
            dag_id="basic_param_dag",
            logical_date=logical_date,
            task=task,
            conf={
                "x": 3,
                "text": "Hello World!",
                "flag": False,
                "a_simple_list": ["one", "two", "three", "actually one value is made per line"],
            },
        )
        _, msg, _ = run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
        assert isinstance(msg, SucceedTask)

    def test_task_run_with_operator_extra_links(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        """Test that a task can run with operator extra links defined and can set an xcom."""
        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        class DummyTestOperator(BaseOperator):
            operator_extra_links = (AirflowLink(),)

            def execute(self, context):
                print("Hello from custom operator", self.operator_extra_links)

        task = DummyTestOperator(task_id="task_with_operator_extra_links")

        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        run(runtime_ti, context=context, log=mock.MagicMock())

        mock_supervisor_comms.send.assert_called_once_with(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )

        with mock.patch.object(XCom, "_set_xcom_in_db") as mock_xcom_set:
            finalize(
                runtime_ti,
                log=mock.MagicMock(),
                state=TaskInstanceState.SUCCESS,
                context=runtime_ti.get_template_context(),
            )
            mock_xcom_set.assert_called_once_with(
                key="_link_AirflowLink",
                value="https://airflow.apache.org",
                dag_id=runtime_ti.dag_id,
                task_id=runtime_ti.task_id,
                run_id=runtime_ti.run_id,
                map_index=runtime_ti.map_index,
            )

    def test_task_failed_with_operator_extra_links(
        self, create_runtime_ti, mock_supervisor_comms, time_machine
    ):
        """Test that operator extra links are pushed to xcoms even when task fails."""
        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        class DummyTestOperator(BaseOperator):
            operator_extra_links = (AirflowLink(),)

            def execute(self, context):
                raise ValueError("Task failed intentionally")

        task = DummyTestOperator(task_id="task_with_operator_extra_links")
        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        runtime_ti.start_date = instant
        runtime_ti.end_date = instant

        state, _, error = run(runtime_ti, context=context, log=mock.MagicMock())
        assert state == TaskInstanceState.FAILED
        assert error is not None

        with mock.patch.object(XCom, "_set_xcom_in_db") as mock_xcom_set:
            finalize(
                runtime_ti,
                log=mock.MagicMock(),
                state=TaskInstanceState.FAILED,
                context=context,
                error=error,
            )
            assert mock_xcom_set.mock_calls == [
                call(
                    key="_link_AirflowLink",
                    value="https://airflow.apache.org",
                    dag_id=runtime_ti.dag_id,
                    task_id=runtime_ti.task_id,
                    run_id=runtime_ti.run_id,
                    map_index=runtime_ti.map_index,
                )
            ]

    def test_operator_extra_links_exception_handling(
        self, create_runtime_ti, mock_supervisor_comms, time_machine
    ):
        """Test that exceptions in get_link() don't prevent other links from being pushed."""
        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        class FailingLink(BaseOperatorLink):
            """A link that raises an exception when get_link is called."""

            name = "failing_link"

            def get_link(self, operator, *, ti_key):
                raise ValueError("Link generation failed")

        class DummyTestOperator(BaseOperator):
            operator_extra_links = (FailingLink(), AirflowLink())

            def execute(self, context):
                pass

        task = DummyTestOperator(task_id="task_with_multiple_links")
        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        runtime_ti.start_date = instant
        runtime_ti.end_date = instant

        with mock.patch.object(XCom, "_set_xcom_in_db") as mock_xcom_set:
            finalize(
                runtime_ti,
                log=mock.MagicMock(),
                state=TaskInstanceState.SUCCESS,
                context=context,
            )
            assert mock_xcom_set.mock_calls == [
                call(
                    key="_link_AirflowLink",
                    value="https://airflow.apache.org",
                    dag_id=runtime_ti.dag_id,
                    task_id=runtime_ti.task_id,
                    run_id=runtime_ti.run_id,
                    map_index=runtime_ti.map_index,
                )
            ]

    @pytest.mark.parametrize(
        ("cmd", "rendered_cmd"),
        [
            pytest.param("echo 'hi'", "echo 'hi'", id="no_template_fields"),
            pytest.param(SET_DURING_EXECUTION, SET_DURING_EXECUTION.serialize(), id="with_default"),
        ],
    )
    def test_overwrite_rtif_after_execution_sets_rtif(
        self, create_runtime_ti, mock_supervisor_comms, cmd, rendered_cmd
    ):
        """Test that the RTIF is overwritten after execution for certain operators."""

        class CustomOperator(BaseOperator):
            overwrite_rtif_after_execution = True
            template_fields = ["bash_command"]

            def __init__(self, bash_command, *args, **kwargs):
                self.bash_command = bash_command
                super().__init__(*args, **kwargs)

        task = CustomOperator(task_id="hello", bash_command=cmd)
        runtime_ti = create_runtime_ti(task=task)

        finalize(
            runtime_ti,
            state=TaskInstanceState.SUCCESS,
            context=runtime_ti.get_template_context(),
            log=mock.MagicMock(),
        )

        mock_supervisor_comms.send.assert_called_with(
            msg=SetRenderedFields(rendered_fields={"bash_command": rendered_cmd})
        )

    def test_overwrite_rtif_after_execution_handles_errors_gracefully(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """
        Test that errors during SetRenderedFields in finalize() don't mask the original task error.
        """

        class TaskWithRTIF(BaseOperator):
            overwrite_rtif_after_execution = True
            template_fields = ["command"]

            def __init__(self, command, *args, **kwargs):
                self.command = command
                super().__init__(*args, **kwargs)

        task = TaskWithRTIF(task_id="test_task", command="test command")
        runtime_ti = create_runtime_ti(task=task)
        mock_log = mock.MagicMock()

        # mock the SetRenderedFields call to fail with API_SERVER_ERROR
        mock_supervisor_comms.send.side_effect = AirflowRuntimeError(
            error=ErrorResponse(
                error=ErrorType.API_SERVER_ERROR,
                detail={
                    "status_code": 404,
                    "message": "Not Found",
                    "detail": {"detail": "Not Found"},
                },
            )
        )

        finalize(
            runtime_ti,
            state=TaskInstanceState.FAILED,
            context=runtime_ti.get_template_context(),
            log=mock_log,
            error=Exception("Task execution failed"),
        )

        mock_log.exception.assert_called_once_with(
            "Failed to set rendered fields during finalization",
            ti=runtime_ti,
            task=task,
        )

    @pytest.mark.parametrize(
        ("task_reschedule_count", "expected_date"),
        [
            (
                0,
                None,
            ),
            (
                1,
                timezone.datetime(2025, 1, 1),
            ),
        ],
    )
    def test_get_first_reschedule_date(
        self, create_runtime_ti, mock_supervisor_comms, task_reschedule_count, expected_date
    ):
        """Test that the first reschedule date is fetched from the Supervisor."""
        task = BaseOperator(task_id="hello")
        runtime_ti = create_runtime_ti(task=task, task_reschedule_count=task_reschedule_count)

        mock_supervisor_comms.send.return_value = TaskRescheduleStartDate(
            start_date=timezone.datetime(2025, 1, 1)
        )

        context = runtime_ti.get_template_context()
        assert runtime_ti.get_first_reschedule_date(context=context) == expected_date

    def test_get_ti_count(self, mock_supervisor_comms):
        """Test that get_ti_count sends the correct request and returns the count."""
        mock_supervisor_comms.send.return_value = TICount(count=2)

        count = RuntimeTaskInstance.get_ti_count(
            dag_id="test_dag",
            task_ids=["task1", "task2"],
            task_group_id="group1",
            logical_dates=[timezone.datetime(2024, 1, 1)],
            run_ids=["run1"],
            states=["success", "failed"],
        )

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetTICount(
                dag_id="test_dag",
                task_ids=["task1", "task2"],
                task_group_id="group1",
                logical_dates=[timezone.datetime(2024, 1, 1)],
                run_ids=["run1"],
                states=["success", "failed"],
            ),
        )
        assert count == 2

    def test_get_dr_count(self, mock_supervisor_comms):
        """Test that get_dr_count sends the correct request and returns the count."""
        mock_supervisor_comms.send.return_value = DRCount(count=2)

        count = RuntimeTaskInstance.get_dr_count(
            dag_id="test_dag",
            logical_dates=[timezone.datetime(2024, 1, 1)],
            run_ids=["run1"],
            states=["success", "failed"],
        )

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetDRCount(
                dag_id="test_dag",
                logical_dates=[timezone.datetime(2024, 1, 1)],
                run_ids=["run1"],
                states=["success", "failed"],
            ),
        )
        assert count == 2

    def test_get_dagrun_state(self, mock_supervisor_comms):
        """Test that get_dagrun_state sends the correct request and returns the state."""
        mock_supervisor_comms.send.return_value = DagRunStateResult(state="running")

        state = RuntimeTaskInstance.get_dagrun_state(
            dag_id="test_dag",
            run_id="run1",
        )

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetDagRunState(dag_id="test_dag", run_id="run1"),
        )
        assert state == "running"

    def test_get_task_states(self, mock_supervisor_comms):
        """Test that get_task_states sends the correct request and returns the states."""
        mock_supervisor_comms.send.return_value = TaskStatesResult(task_states={"run1": {"task1": "running"}})

        states = RuntimeTaskInstance.get_task_states(
            dag_id="test_dag",
            task_ids=["task1"],
            run_ids=["run1"],
        )

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetTaskStates(
                dag_id="test_dag",
                task_ids=["task1"],
                run_ids=["run1"],
            ),
        )
        assert states == {"run1": {"task1": "running"}}

    def test_get_previous_dagrun_basic(self, create_runtime_ti, mock_supervisor_comms):
        """Test that get_previous_dagrun sends the correct request without state filter."""

        task = BaseOperator(task_id="hello")
        dag_id = "test_dag"
        runtime_ti = create_runtime_ti(task=task, dag_id=dag_id, logical_date=timezone.datetime(2025, 1, 2))

        dag_run_data = DagRun(
            dag_id=dag_id,
            run_id="prev_run",
            logical_date=timezone.datetime(2025, 1, 1),
            start_date=timezone.datetime(2025, 1, 1),
            run_after=timezone.datetime(2025, 1, 1),
            run_type="scheduled",
            state="success",
            consumed_asset_events=[],
            data_interval_start=None,
            data_interval_end=None,
            end_date=None,
            partition_key=None,
        )

        mock_supervisor_comms.send.return_value = PreviousDagRunResult(dag_run=dag_run_data)

        dr = runtime_ti.get_previous_dagrun()

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetPreviousDagRun(dag_id="test_dag", logical_date=timezone.datetime(2025, 1, 2), state=None),
        )
        assert dr.dag_id == "test_dag"
        assert dr.run_id == "prev_run"
        assert dr.state == "success"

    def test_get_previous_dagrun_with_state(self, create_runtime_ti, mock_supervisor_comms):
        """Test that get_previous_dagrun sends the correct request with state filter."""

        task = BaseOperator(task_id="hello")
        dag_id = "test_dag"
        runtime_ti = create_runtime_ti(task=task, dag_id=dag_id, logical_date=timezone.datetime(2025, 1, 2))

        dag_run_data = DagRun(
            dag_id=dag_id,
            run_id="prev_success_run",
            logical_date=timezone.datetime(2025, 1, 1),
            start_date=timezone.datetime(2025, 1, 1),
            run_after=timezone.datetime(2025, 1, 1),
            run_type="scheduled",
            state="success",
            consumed_asset_events=[],
            data_interval_start=None,
            data_interval_end=None,
            end_date=None,
            partition_key=None,
        )

        mock_supervisor_comms.send.return_value = PreviousDagRunResult(dag_run=dag_run_data)

        dr = runtime_ti.get_previous_dagrun(state="success")

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetPreviousDagRun(
                dag_id="test_dag", logical_date=timezone.datetime(2025, 1, 2), state="success"
            ),
        )
        assert dr.dag_id == "test_dag"
        assert dr.run_id == "prev_success_run"
        assert dr.state == "success"

    def test_get_previous_ti_basic(self, create_runtime_ti, mock_supervisor_comms):
        """Test that get_previous_ti sends the correct request without filters."""

        task = BaseOperator(task_id="test_task")
        dag_id = "test_dag"
        runtime_ti = create_runtime_ti(task=task, dag_id=dag_id, logical_date=timezone.datetime(2025, 1, 2))

        ti_data = PreviousTIResponse(
            task_id="test_task",
            dag_id=dag_id,
            run_id="prev_run",
            logical_date=timezone.datetime(2025, 1, 1),
            start_date=timezone.datetime(2025, 1, 1, 12, 0, 0),
            end_date=timezone.datetime(2025, 1, 1, 12, 5, 0),
            state="success",
            try_number=1,
            map_index=-1,
            duration=300.0,
        )

        mock_supervisor_comms.send.return_value = PreviousTIResult(task_instance=ti_data)

        prev_ti = runtime_ti.get_previous_ti()

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetPreviousTI(
                dag_id="test_dag",
                task_id="test_task",
                logical_date=timezone.datetime(2025, 1, 2),
                map_index=-1,
                state=None,
            ),
        )
        assert prev_ti.task_id == "test_task"
        assert prev_ti.dag_id == "test_dag"
        assert prev_ti.run_id == "prev_run"
        assert prev_ti.state == "success"

    def test_get_previous_ti_with_state(self, create_runtime_ti, mock_supervisor_comms):
        """Test get_previous_ti with state filter."""

        task = BaseOperator(task_id="test_task")
        dag_id = "test_dag"
        runtime_ti = create_runtime_ti(task=task, dag_id=dag_id, logical_date=timezone.datetime(2025, 1, 2))

        ti_data = PreviousTIResponse(
            task_id="test_task",
            dag_id=dag_id,
            run_id="prev_success_run",
            logical_date=timezone.datetime(2025, 1, 1),
            start_date=timezone.datetime(2025, 1, 1, 12, 0, 0),
            end_date=timezone.datetime(2025, 1, 1, 12, 5, 0),
            state="success",
            try_number=1,
            map_index=-1,
            duration=300.0,
        )

        mock_supervisor_comms.send.return_value = PreviousTIResult(task_instance=ti_data)

        prev_ti = runtime_ti.get_previous_ti(state=TaskInstanceState.SUCCESS)

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetPreviousTI(
                dag_id="test_dag",
                task_id="test_task",
                logical_date=timezone.datetime(2025, 1, 2),
                map_index=-1,
                state=TaskInstanceState.SUCCESS,
            ),
        )
        assert prev_ti.state == "success"
        assert prev_ti.run_id == "prev_success_run"

    def test_get_previous_ti_with_map_index(self, create_runtime_ti, mock_supervisor_comms):
        """Test get_previous_ti with explicit map_index filter."""

        task = BaseOperator(task_id="test_task")
        dag_id = "test_dag"
        runtime_ti = create_runtime_ti(
            task=task, dag_id=dag_id, logical_date=timezone.datetime(2025, 1, 2), map_index=0
        )

        ti_data = PreviousTIResponse(
            task_id="test_task",
            dag_id=dag_id,
            run_id="prev_run",
            logical_date=timezone.datetime(2025, 1, 1),
            start_date=timezone.datetime(2025, 1, 1, 12, 0, 0),
            end_date=timezone.datetime(2025, 1, 1, 12, 5, 0),
            state="success",
            try_number=1,
            map_index=1,
            duration=300.0,
        )

        mock_supervisor_comms.send.return_value = PreviousTIResult(task_instance=ti_data)

        # Query for a different map_index than current TI
        prev_ti = runtime_ti.get_previous_ti(map_index=1)

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetPreviousTI(
                dag_id="test_dag",
                task_id="test_task",
                logical_date=timezone.datetime(2025, 1, 2),
                map_index=1,
                state=None,
            ),
        )
        assert prev_ti.map_index == 1

    def test_get_previous_ti_not_found(self, create_runtime_ti, mock_supervisor_comms):
        """Test get_previous_ti when no previous TI exists."""

        task = BaseOperator(task_id="test_task")
        dag_id = "test_dag"
        runtime_ti = create_runtime_ti(task=task, dag_id=dag_id, logical_date=timezone.datetime(2025, 1, 2))

        mock_supervisor_comms.send.return_value = PreviousTIResult(task_instance=None)

        prev_ti = runtime_ti.get_previous_ti()

        assert prev_ti is None

    @pytest.mark.parametrize(
        "map_index",
        [
            pytest.param(-1, id="map_index_negative_one"),
            pytest.param(0, id="map_index_zero"),
            pytest.param(5, id="map_index_positive"),
        ],
    )
    def test_xcom_clearing_includes_map_index(self, create_runtime_ti, mock_supervisor_comms, map_index):
        """Test that XCom clearing during task execution properly passes map_index."""

        task = BaseOperator(task_id="test_task")
        runtime_ti = create_runtime_ti(task=task, map_index=map_index)

        # Mock the ti_context_from_server to include xcom_keys_to_clear
        runtime_ti._ti_context_from_server = TIRunContext(
            dag_run=runtime_ti._ti_context_from_server.dag_run,
            task_reschedule_count=0,
            max_tries=1,
            should_retry=False,
            xcom_keys_to_clear=["key1", "key2", "key3"],
        )

        with mock.patch.object(XCom, "delete") as mock_delete:
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

            # Verify that XCom.delete was called for each key with the correct map_index
            expected_calls = [
                mock.call(
                    key="key1",
                    dag_id=runtime_ti.dag_id,
                    task_id=runtime_ti.task_id,
                    run_id=runtime_ti.run_id,
                    map_index=map_index,
                ),
                mock.call(
                    key="key2",
                    dag_id=runtime_ti.dag_id,
                    task_id=runtime_ti.task_id,
                    run_id=runtime_ti.run_id,
                    map_index=map_index,
                ),
                mock.call(
                    key="key3",
                    dag_id=runtime_ti.dag_id,
                    task_id=runtime_ti.task_id,
                    run_id=runtime_ti.run_id,
                    map_index=map_index,
                ),
            ]
            mock_delete.assert_has_calls(expected_calls)

    def test_xcom_clearing_without_keys_to_clear(self, create_runtime_ti, mock_supervisor_comms):
        """Test that no XCom clearing occurs when xcom_keys_to_clear is empty."""
        task = BaseOperator(task_id="test_task")
        runtime_ti = create_runtime_ti(task=task)

        assert runtime_ti._ti_context_from_server.xcom_keys_to_clear is None

        with mock.patch.object(XCom, "delete") as mock_delete:
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

            mock_delete.assert_not_called()

    def test_xcom_push_pull_with_slash_in_key(self, create_runtime_ti, mock_supervisor_comms):
        """
        Ensure that XCom keys containing slashes are correctly quoted/unquoted
        and do not break API routes (no 400/404).
        """

        class PushOperator(BaseOperator):
            def execute(self, context):
                context["ti"].xcom_push(key="some/key/with/slash", value="slash_value")

        task = PushOperator(task_id="push_task")
        runtime_ti = create_runtime_ti(task=task, dag_id="test_dag")

        # Run the task (which should trigger xcom_push)
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        # Verify supervisor received a SetXCom with quoted key
        called_args = [
            call.kwargs.get("msg") or call.args[0] for call in mock_supervisor_comms.send.call_args_list
        ]
        assert any(getattr(arg, "key", None) == "some/key/with/slash" for arg in called_args)

        ser_value = BaseXCom.serialize_value("slash_value")
        mock_supervisor_comms.send.reset_mock()
        mock_supervisor_comms.send.return_value = XComSequenceSliceResult(
            key="some/key/with/slash",
            root=[ser_value],
        )

        pulled_value = runtime_ti.xcom_pull(key="some/key/with/slash", task_ids="push_task")
        assert pulled_value == "slash_value"

        # Key should NOT be quoted here - client API will handle encoding
        mock_supervisor_comms.send.assert_any_call(
            GetXComSequenceSlice(
                key="some/key/with/slash",
                dag_id="test_dag",
                run_id="test_run",
                task_id="push_task",
                map_index=0,
                include_prior_dates=False,
                start=None,
                stop=None,
                step=None,
                type="GetXComSequenceSlice",
            )
        )

    def test_taskflow_dict_return_with_slash_key(self, create_runtime_ti, mock_supervisor_comms):
        """
        High-level: Ensure TaskFlow returning dict with slash in key doesn't 404 during XCom push.
        """

        @dag_decorator(schedule=None, start_date=timezone.datetime(2024, 12, 3))
        def dag_with_slash_key():
            @task_decorator
            def dict_task():
                return {"key with slash /": "Some Value"}

            return dict_task()  # returns XComArg

        dag_obj = dag_with_slash_key()
        task_op = dag_obj.get_task("dict_task")
        runtime_ti = create_runtime_ti(task=task_op, dag_id=dag_obj.dag_id)

        # Run task instance → should trigger TaskFlow dict expansion + XCom push
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        # Mock supervisor response to simulate retrieval
        ser_value = BaseXCom.serialize_value("Some Value")
        mock_supervisor_comms.send.reset_mock()
        mock_supervisor_comms.send.return_value = XComSequenceSliceResult(
            key="key/slash",
            root=[ser_value],
        )

        pulled = runtime_ti.xcom_pull(key="key/slash", task_ids="dict_task")
        assert pulled == "Some Value"

    @pytest.mark.enable_redact
    def test_rendered_templates_mask_secrets_with_truncation(self, create_runtime_ti, mock_supervisor_comms):
        """Test that secrets are masked before truncation when rendered fields exceed max_templated_field_length."""
        from airflow.sdk._shared.secrets_masker import _secrets_masker

        secret_url = "postgresql+psycopg2://username:testpass123@test.domain.com/testdb"
        _secrets_masker().add_mask(secret_url, None)

        class CustomOperator(BaseOperator):
            template_fields = ("env_vars", "region")

            def __init__(self, env_vars, region, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.env_vars = env_vars
                self.region = region

            def execute(self, context):
                pass

        # generate 50 env_vars to exceed default char limit of 4096 (50 * 87 chars ≈ 4350 chars)
        env_vars = {f"TEST_URL_{i}": secret_url for i in range(50)}

        task = CustomOperator(
            task_id="test_truncation_masking",
            env_vars=env_vars,
            region="us-west-2",
        )

        runtime_ti = create_runtime_ti(task=task, dag_id="test_truncation_masking_dag")
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        msg = next(
            c.kwargs["msg"]
            for c in mock_supervisor_comms.send.mock_calls
            if c.kwargs.get("msg") and getattr(c.kwargs["msg"], "type", None) == "SetRenderedFields"
        )
        rendered_fields = msg.rendered_fields

        # region is short enough to not be truncated
        assert rendered_fields["region"] == "us-west-2"

        # env_vars exceeds max_templated_field_length and must be truncated with secrets redacted
        env_vars_value = rendered_fields["env_vars"]
        assert isinstance(env_vars_value, str)
        assert env_vars_value.startswith(
            "Truncated. You can change this behaviour in [core]max_templated_field_length. "
        )
        assert env_vars_value.endswith("...")
        assert "***" in env_vars_value  # secrets are redacted before truncation

    @pytest.mark.enable_redact
    def test_rendered_templates_mask_nested_keys_with_truncation(
        self, create_runtime_ti, mock_supervisor_comms, monkeypatch
    ):
        """Nested sensitive-key masking applies consistently across the truncation path.

        A value under a documented sensitive key (``password``, ``token``, ``secret``,
        ``api_key``) is masked recursively by ``redact()`` when the structured value
        is walked. The oversized branch must redact while still structured so that
        nested-key context is preserved before stringification — otherwise the post-
        stringify ``redact()`` call only sees the outer field name and the recursive
        walker cannot reach the inner key.
        """
        from airflow.sdk._shared.secrets_masker import _secrets_masker

        # Earlier tests in this file (e.g. test_get_connection_from_context) call
        # mask_secret(conn.password) where the fixture's password value is the literal
        # "password"; that registers "password" as a regex pattern in the singleton
        # masker. Without isolation, str(redacted) gets that regex applied and the
        # dict KEY name "password" itself becomes "***", obscuring whether the
        # structured nested-key walk fired. Reset the regex patterns for this test
        # (monkeypatch restores them on teardown) so the assertion can distinguish
        # value-masking (what we are testing) from key-token replacement.
        masker = _secrets_masker()
        monkeypatch.setattr(masker, "patterns", set())
        monkeypatch.setattr(masker, "replacer", None)
        # The SDK masker starts with an empty sensitive-fields list in the test runtime
        # (settings.py has not run); register `password` explicitly so the structured
        # walker has something to match. Production workers get this from settings.py.
        monkeypatch.setattr(
            masker,
            "sensitive_variables_fields",
            list(masker.sensitive_variables_fields) + ["password"],
        )

        nested_value = "REGRESSION-FIXTURE-NESTED-PASSWORD-VALUE"

        class CustomOperator(BaseOperator):
            template_fields = ("env_vars",)

            def __init__(self, env_vars, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.env_vars = env_vars

            def execute(self, context):
                pass

        # Nested 'password' key under enough padding to exceed default 4096-char limit.
        env_vars = {
            "DB": {"password": nested_value, "host": "db.internal", "zz_pad": "A" * 5000},
        }

        task = CustomOperator(task_id="test_nested_truncation_masking", env_vars=env_vars)

        runtime_ti = create_runtime_ti(task=task, dag_id="test_nested_truncation_masking_dag")
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        msg = next(
            c.kwargs["msg"]
            for c in mock_supervisor_comms.send.mock_calls
            if c.kwargs.get("msg") and getattr(c.kwargs["msg"], "type", None) == "SetRenderedFields"
        )
        env_vars_value = msg.rendered_fields["env_vars"]

        assert isinstance(env_vars_value, str)
        assert env_vars_value.startswith(
            "Truncated. You can change this behaviour in [core]max_templated_field_length. "
        )
        assert nested_value not in env_vars_value
        assert "'password': '***'" in env_vars_value

    @pytest.mark.enable_redact
    def test_rendered_templates_masks_secrets_in_complex_objects(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """Test that secrets in complex objects like V1EnvVar are masked well."""
        from kubernetes import client as k8s

        from airflow.sdk._shared.secrets_masker import _secrets_masker

        secret1 = "This is a longer test phrase. We are checking if this handles regular sentences."
        secret2 = "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."
        _secrets_masker().add_mask(secret1, None)
        _secrets_masker().add_mask(secret2, None)

        class CustomOperator(BaseOperator):
            template_fields = ("env_vars",)

            def __init__(self, env_vars, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.env_vars = env_vars

            def execute(self, context):
                pass

        env_vars = [
            k8s.V1EnvVar(name="var1", value="This is a test phrase."),
            k8s.V1EnvVar(name="var2", value=secret1),
            k8s.V1EnvVar(name="var3", value=secret2),
        ]

        task = CustomOperator(
            task_id="test_complex_object_masking",
            env_vars=env_vars,
        )

        runtime_ti = create_runtime_ti(task=task, dag_id="test_complex_object_dag")
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        rendered_fields = mock_supervisor_comms.send.mock_calls[0].kwargs["msg"].rendered_fields
        assert rendered_fields is not None
        # K8s V1EnvVar objects expose .to_dict(); the recursive walk normalizes the list of objects
        # into a list of plain dicts so the result is directly JSON-encodable and redact can mask secrets in nested values.
        assert rendered_fields["env_vars"] == [
            {"name": "var1", "value": "This is a test phrase.", "value_from": None},
            {"name": "var2", "value": "***", "value_from": None},
            {"name": "var3", "value": "***", "value_from": None},
        ]

    def test_nested_template_field_renderer_respects_redaction(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """
        Ensure nested template_fields_renderers paths still serialize
        and redact correctly.
        """

        class CustomOperator(BaseOperator):
            template_fields = ("config",)
            template_fields_renderers = {"config.nested.secret": "json"}

            def __init__(self, config, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.config = config

            def execute(self, context):
                pass

        config = {"nested": {"secret": "top_secret"}}

        task = CustomOperator(
            task_id="nested_redact_test",
            config=config,
        )

        runtime_ti = create_runtime_ti(task=task)

        with mock.patch(
            "airflow.sdk._shared.secrets_masker.redact",
            side_effect=lambda val, _: f"MASKED:{val}",
        ):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        assert any(
            "config.nested.secret" in call.kwargs.get("msg").rendered_fields
            for call in mock_supervisor_comms.send.mock_calls
            if hasattr(call.kwargs.get("msg"), "rendered_fields")
        )

    def test_rendered_fields_validates_json_value_types(self, create_runtime_ti, mock_supervisor_comms):
        """
        Ensure validated JSON-compatible types are preserved after redaction.
        """

        class CustomOperator(BaseOperator):
            template_fields = ("data",)

            def __init__(self, data, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.data = data

            def execute(self, context):
                pass

        complex_value = {"key": [1, 2, {"nested": "value"}]}

        task = CustomOperator(task_id="json_validation_test", data=complex_value)
        runtime_ti = create_runtime_ti(task=task)

        with mock.patch(
            "airflow.sdk._shared.secrets_masker.redact",
            side_effect=lambda val, _: val,
        ):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        assert any(
            call.kwargs.get("msg").rendered_fields["data"] == complex_value
            for call in mock_supervisor_comms.send.mock_calls
            if hasattr(call.kwargs.get("msg"), "rendered_fields")
        )

    def test_get_dag(self, mock_supervisor_comms):
        """Test that get_dag sends the correct request and returns the dag."""
        mock_supervisor_comms.send.return_value = DagResult(
            dag_id="test_dag",
            is_paused=False,
            bundle_name="dags-folder",
            bundle_version="bundle-version",
            relative_fileloc="dags/example.py",
            owners="owner_1",
            tags=["a_tag", "z_tag"],
            next_dagrun=datetime(2026, 4, 13, tzinfo=dt_timezone.utc),
        )

        response = RuntimeTaskInstance.get_dag(
            dag_id="test_dag",
        )

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetDag(dag_id="test_dag"),
        )
        assert response.dag_id == "test_dag"
        assert response.is_paused is False


class TestXComAfterTaskExecution:
    @pytest.mark.parametrize(
        ("do_xcom_push", "should_push_xcom", "expected_xcom_value"),
        [
            pytest.param(False, False, None, id="do_xcom_push_false"),
            pytest.param(True, True, "Hello World!", id="do_xcom_push_true"),
        ],
    )
    def test_xcom_push_flag(
        self,
        create_runtime_ti,
        mock_supervisor_comms,
        spy_agency,
        do_xcom_push: bool,
        should_push_xcom: bool,
        expected_xcom_value,
    ):
        """Test that the do_xcom_push flag controls whether the task pushes to XCom."""

        class CustomOperator(BaseOperator):
            def execute(self, context):
                return "Hello World!"

        task = CustomOperator(task_id="hello", do_xcom_push=do_xcom_push)

        runtime_ti = create_runtime_ti(task=task)

        spy_agency.spy_on(_push_xcom_if_needed, call_original=True)
        spy_agency.spy_on(_xcom_push, call_original=False)

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        spy_agency.assert_spy_called(_push_xcom_if_needed)

        if should_push_xcom:
            spy_agency.assert_spy_called_with(
                _xcom_push, runtime_ti, BaseXCom.XCOM_RETURN_KEY, expected_xcom_value
            )
        else:
            spy_agency.assert_spy_not_called(_xcom_push)

    def test_xcom_with_multiple_outputs(self, create_runtime_ti, spy_agency):
        """Test that the task pushes to XCom when multiple outputs are returned."""
        result = {"key1": "value1", "key2": "value2"}

        class CustomOperator(BaseOperator):
            def execute(self, context):
                return result

        task = CustomOperator(
            task_id="test_xcom_push_with_multiple_outputs", do_xcom_push=True, multiple_outputs=True
        )

        runtime_ti = create_runtime_ti(task=task)

        spy_agency.spy_on(_xcom_push, call_original=False)
        _push_xcom_if_needed(result=result, ti=runtime_ti, log=mock.MagicMock())

        expected_calls = [
            ("key1", "value1"),
            ("key2", "value2"),
            (BaseXCom.XCOM_RETURN_KEY, result),
        ]
        spy_agency.assert_spy_call_count(_xcom_push, len(expected_calls))
        for key, value in expected_calls:
            spy_agency.assert_spy_called_with(_xcom_push, runtime_ti, key, value, mapped_length=None)

    def test_xcom_with_mapped_length(self, create_runtime_ti):
        """Test that the task pushes to XCom with mapped length."""
        result = {"key1": "value1", "key2": "value2"}

        class CustomOperator(BaseOperator):
            def execute(self, context):
                return result

        task = CustomOperator(
            task_id="test_xcom_push_with_mapped_length",
            do_xcom_push=True,
        )

        runtime_ti = create_runtime_ti(task=task)

        with mock.patch.object(XCom, "set") as mock_xcom_set:
            _xcom_push(runtime_ti, BaseXCom.XCOM_RETURN_KEY, result, mapped_length=7)
            mock_xcom_set.assert_called_once_with(
                key=BaseXCom.XCOM_RETURN_KEY,
                value=result,
                dag_id=runtime_ti.dag_id,
                task_id=runtime_ti.task_id,
                run_id=runtime_ti.run_id,
                map_index=runtime_ti.map_index,
                dag_result=False,
                _mapped_length=7,
            )

    def test_xcom_with_multiple_outputs_and_no_mapping_result(self, create_runtime_ti, spy_agency):
        """Test that error is raised when multiple outputs are returned without mapping."""
        result = "value1"

        class CustomOperator(BaseOperator):
            def execute(self, context):
                return result

        task = CustomOperator(
            task_id="test_xcom_push_with_multiple_outputs", do_xcom_push=True, multiple_outputs=True
        )

        runtime_ti = create_runtime_ti(task=task)

        spy_agency.spy_on(runtime_ti.xcom_push, call_original=False)
        with pytest.raises(
            TypeError,
            match=f"Returned output was type {type(result)} expected dictionary for multiple_outputs",
        ):
            _push_xcom_if_needed(result=result, ti=runtime_ti, log=mock.MagicMock())

    def test_xcom_with_multiple_outputs_and_key_is_not_string(self, create_runtime_ti, spy_agency):
        """Test that error is raised when multiple outputs are returned and key isn't string."""
        result = {2: "value1", "key2": "value2"}

        class CustomOperator(BaseOperator):
            def execute(self, context):
                return result

        task = CustomOperator(
            task_id="test_xcom_push_with_multiple_outputs", do_xcom_push=True, multiple_outputs=True
        )

        runtime_ti = create_runtime_ti(task=task)

        spy_agency.spy_on(runtime_ti.xcom_push, call_original=False)

        with pytest.raises(TypeError) as exc_info:
            _push_xcom_if_needed(result=result, ti=runtime_ti, log=mock.MagicMock())

        assert str(exc_info.value) == (
            f"Returned dictionary keys must be strings when using multiple_outputs, found 2 ({int}) instead"
        )

    def test_xcom_push_to_custom_xcom_backend(
        self, create_runtime_ti, mock_supervisor_comms, mock_xcom_backend
    ):
        """Test that a task pushes a xcom to the custom xcom backend."""

        class CustomOperator(BaseOperator):
            def execute(self, context):
                return "pushing to xcom backend!"

        task = CustomOperator(task_id="pull_task")
        runtime_ti = create_runtime_ti(task=task)

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_xcom_backend.set.assert_called_once_with(
            key=BaseXCom.XCOM_RETURN_KEY,
            value="pushing to xcom backend!",
            dag_id="test_dag",
            task_id="pull_task",
            run_id="test_run",
            map_index=-1,
            dag_result=False,
            _mapped_length=None,
        )

        # assert that we didn't call the API when XCom backend is configured
        assert not any(
            x
            == mock.call(
                msg=SetXCom(
                    key="key",
                    value="pushing to xcom backend!",
                    dag_id="test_dag",
                    run_id="test_run",
                    task_id="pull_task",
                    map_index=-1,
                ),
            )
            for x in mock_supervisor_comms.send.call_args_list
        )

    def test_xcom_pull_from_custom_xcom_backend(
        self, create_runtime_ti, mock_supervisor_comms, mock_xcom_backend
    ):
        """Test that a task pulls the expected XCom value if it exists, but from custom xcom backend."""

        class CustomOperator(BaseOperator):
            def execute(self, context):
                context["ti"].xcom_pull(task_ids="pull_task", key="key")

        task = CustomOperator(task_id="pull_task")
        runtime_ti = create_runtime_ti(task=task)
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_xcom_backend.get_all.assert_called_once_with(
            key="key",
            dag_id="test_dag",
            task_id="pull_task",
            run_id="test_run",
            include_prior_dates=False,
        )

        assert not any(
            x
            == mock.call(
                msg=GetXCom(
                    key="key",
                    dag_id="test_dag",
                    run_id="test_run",
                    task_id="pull_task",
                    map_index=-1,
                ),
            )
            for x in mock_supervisor_comms.send.call_args_list
        )

    def test_get_all_uses_custom_deserialize_value(self, mock_supervisor_comms):
        """
        Tests that XCom.get_all() calls the custom deserialize_value method.
        """

        class CustomXCom(BaseXCom):
            @classmethod
            def deserialize_value(cls, result):
                """Custom deserialization that adds a prefix to show it was called."""
                original_value = super().deserialize_value(result)
                return f"from custom xcom deserialize:{original_value}"

        serialized_values = ["value1", "value2", "value3"]
        mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=serialized_values)

        result = CustomXCom.get_all(key="test_key", dag_id="test_dag", task_id="test_task", run_id="test_run")

        expected = [
            "from custom xcom deserialize:value1",
            "from custom xcom deserialize:value2",
            "from custom xcom deserialize:value3",
        ]
        assert result == expected

    def test_xcom_pull_with_explicit_dag_id_and_run_id(self, create_runtime_ti, mock_supervisor_comms):
        task = BaseOperator(task_id="parent_task")
        runtime_ti = create_runtime_ti(task=task, dag_id="parent_dag", run_id="parent_run")
        value = {"child_result": "hello world"}
        ser_value = BaseXCom.serialize_value(value)

        mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=[ser_value])

        assert (
            runtime_ti.xcom_pull(
                task_ids="child_task",
                dag_id="child_dag",
                run_id="child_run",
            )
            == value
        )
        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetXComSequenceSlice(
                key="return_value",
                dag_id="child_dag",
                run_id="child_run",
                task_id="child_task",
                start=None,
                stop=None,
                step=None,
                include_prior_dates=False,
            ),
        )

    @pytest.mark.parametrize(
        ("include_prior_dates", "expected_value"),
        [
            pytest.param(True, True, id="include_prior_dates_true"),
            pytest.param(False, False, id="include_prior_dates_false"),
            pytest.param(None, False, id="include_prior_dates_default"),
        ],
    )
    def test_xcom_pull_with_include_prior_dates(
        self,
        create_runtime_ti,
        mock_supervisor_comms,
        include_prior_dates,
        expected_value,
    ):
        """Test that xcom_pull with include_prior_dates parameter correctly behaves as we expect."""
        task = BaseOperator(task_id="pull_task")
        runtime_ti = create_runtime_ti(task=task)

        value = {"previous_run_data": "test_value"}
        ser_value = BaseXCom.serialize_value(value)

        def mock_send_side_effect(*args, **kwargs):
            msg = kwargs.get("msg") or args[0]
            if isinstance(msg, GetXComSequenceSlice):
                assert msg.include_prior_dates is expected_value, (
                    f"include_prior_dates should be {expected_value} in GetXComSequenceSlice"
                )
                return XComSequenceSliceResult(root=[ser_value])
            return XComResult(key="test_key", value=None)

        mock_supervisor_comms.send.side_effect = mock_send_side_effect
        kwargs = {"key": "test_key", "task_ids": "previous_task"}
        if include_prior_dates is not None:
            kwargs["include_prior_dates"] = include_prior_dates
        result = runtime_ti.xcom_pull(**kwargs)
        assert result == value

        mock_supervisor_comms.send.assert_called_once_with(
            msg=GetXComSequenceSlice(
                key="test_key",
                dag_id=runtime_ti.dag_id,
                run_id=runtime_ti.run_id,
                task_id="previous_task",
                start=None,
                stop=None,
                step=None,
                include_prior_dates=expected_value,
            ),
        )


def _recording_email_backend(
    to: list[str] | Iterable[str],
    subject: str,
    html_content: str,
    files: list[str] | None = None,
    dryrun: bool = False,
    cc: str | Iterable[str] | None = None,
    bcc: str | Iterable[str] | None = None,
    mime_subtype: str = "mixed",
    mime_charset: str = "utf-8",
    conn_id: str | None = None,
    custom_headers: dict[str, Any] | None = None,
    **kwargs,
) -> None:
    """
    Legacy ``[email] email_backend`` stub, patched with an autospecced mock in tests.

    Mirrors the ``airflow.utils.email.send_email`` signature so the autospec enforces the
    calling convention ``_LegacyEmailBackendNotifier`` has to honour. Duplicated here rather
    than imported because the Task SDK must not depend on ``airflow-core``.
    """
    raise AssertionError("should be patched in the test")


class TestEmailNotifications:
    FROM = "from@airflow"

    @pytest.mark.parametrize(
        ("emails", "sent"),
        [
            pytest.param(
                "test@example.com",
                True,
                id="one-email",
            ),
            pytest.param(
                ["test@example.com"],
                True,
                id="one-email-as-list",
            ),
            pytest.param(
                ["test@example.com", "test2@example.com"],
                True,
                id="multiple-email-as-list",
            ),
            pytest.param(None, False, id="no-email"),
            pytest.param([], False, id="no-email-as-list"),
        ],
    )
    def test_email_on_retry(self, emails, sent, create_runtime_ti, mock_supervisor_comms):
        """Test email notification on task retry."""
        from airflow.sdk.execution_time.task_runner import finalize, run

        class ZeroDivsionOperator(BaseOperator):
            def execute(self, context):
                1 // 0

        task = ZeroDivsionOperator(
            task_id="divide_by_zero_task",
            email=emails,
            email_on_retry=True,
            retries=2,
        )

        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        log = mock.MagicMock()

        with conf_vars({("email", "from_email"): self.FROM}):
            with mock.patch("airflow.providers.smtp.notifications.smtp.SmtpNotifier") as mock_smtp_notifier:
                state, _, error = run(runtime_ti, context, log)
                finalize(runtime_ti, state, context, log, error)

                if not sent:
                    mock_smtp_notifier.assert_not_called()
                else:
                    mock_smtp_notifier.assert_called_once()
                    kwargs = mock_smtp_notifier.call_args.kwargs
                    assert kwargs["from_email"] == self.FROM
                    assert kwargs["to"] == emails
                    assert (
                        kwargs["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>'
                    )

    @pytest.mark.parametrize(
        ("emails", "sent"),
        [
            pytest.param(
                "test@example.com",
                True,
                id="one-email",
            ),
            pytest.param(
                ["test@example.com"],
                True,
                id="one-email-as-list",
            ),
            pytest.param(
                ["test@example.com", "test2@example.com"],
                True,
                id="multiple-email-as-list",
            ),
            pytest.param(None, False, id="no-email"),
            pytest.param([], False, id="no-email-as-list"),
        ],
    )
    def test_email_on_failure(self, emails, sent, create_runtime_ti, mock_supervisor_comms):
        """Test email notification on task failure."""
        from airflow.sdk.exceptions import AirflowFailException
        from airflow.sdk.execution_time.task_runner import finalize, run

        class FailingOperator(BaseOperator):
            def execute(self, context):
                raise AirflowFailException("Task failed on purpose")

        task = FailingOperator(
            task_id="failing_task",
            email=emails,
            email_on_failure=True,
        )

        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        log = mock.MagicMock()

        with conf_vars({("email", "from_email"): self.FROM}):
            with mock.patch("airflow.providers.smtp.notifications.smtp.SmtpNotifier") as mock_smtp_notifier:
                state, _, error = run(runtime_ti, context, log)
                finalize(runtime_ti, state, context, log, error)

                if not sent:
                    mock_smtp_notifier.assert_not_called()
                else:
                    mock_smtp_notifier.assert_called_once()
                    kwargs = mock_smtp_notifier.call_args.kwargs
                    assert kwargs["from_email"] == self.FROM
                    assert kwargs["to"] == emails
                    assert (
                        kwargs["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>'
                    )

    def test_email_with_custom_templates(self, create_runtime_ti, mock_supervisor_comms, tmp_path):
        """Test email notification respects custom subject and html_content templates."""
        from airflow.sdk.exceptions import AirflowFailException

        subject_template = tmp_path / "custom_subject.jinja2"
        html_template = tmp_path / "custom_html.html"

        subject_template.write_text("Custom Subject: Task {{ti.task_id}} Failed\n")
        html_template.write_text(
            "<h1>Custom Template</h1><p>Task: {{ti.task_id}}</p><p>Error: {{exception_html}}</p>"
        )

        class FailingOperator(BaseOperator):
            def execute(self, context):
                raise AirflowFailException("Task failed for template test")

        task = FailingOperator(
            task_id="template_test_task",
            email=["test@example.com"],
            email_on_failure=True,
        )

        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        log = mock.MagicMock()

        with conf_vars(
            {
                ("email", "subject_template"): str(subject_template),
                ("email", "html_content_template"): str(html_template),
                ("email", "from_email"): self.FROM,
            }
        ):
            with mock.patch("airflow.providers.smtp.notifications.smtp.SmtpNotifier") as mock_smtp_notifier:
                state, _, error = run(runtime_ti, context, log)
                finalize(runtime_ti, state, context, log, error)

                mock_smtp_notifier.assert_called_once()
                kwargs = mock_smtp_notifier.call_args.kwargs

                assert kwargs["subject"] == "Custom Subject: Task {{ti.task_id}} Failed\n"
                assert (
                    kwargs["html_content"]
                    == "<h1>Custom Template</h1><p>Task: {{ti.task_id}}</p><p>Error: {{exception_html}}</p>"
                )
                assert kwargs["from_email"] == self.FROM

    def test_custom_email_backend_is_used(self, create_runtime_ti, mock_supervisor_comms):
        """A custom ``[email] email_backend`` is wrapped and invoked with rendered fields."""
        backend = mock.create_autospec(_recording_email_backend)

        class FailingOperator(BaseOperator):
            def execute(self, context):
                raise AirflowFailException("Task failed on purpose")

        task = FailingOperator(
            task_id="legacy_backend_task",
            email=["test@example.com"],
            email_on_failure=True,
        )

        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        log = mock.MagicMock()

        with conf_vars(
            {
                ("email", "email_backend"): f"{__name__}._recording_email_backend",
                ("email", "email_conn_id"): "my_smtp",
                ("email", "from_email"): self.FROM,
            }
        ):
            with mock.patch(f"{__name__}._recording_email_backend", backend):
                with mock.patch(
                    "airflow.providers.smtp.notifications.smtp.SmtpNotifier", autospec=True
                ) as mock_smtp_notifier:
                    state, _, error = run(runtime_ti, context, log)
                    finalize(runtime_ti, state, context, log, error)

        # The default SMTP notifier must not be used when a custom backend is configured.
        mock_smtp_notifier.assert_not_called()
        backend.assert_called_once()
        args, kwargs = backend.call_args
        # send_email(to, subject, html_content, conn_id=..., from_email=...)
        assert args[0] == ["test@example.com"]
        assert kwargs["conn_id"] == "my_smtp"
        assert kwargs["from_email"] == self.FROM

    def test_unresolvable_email_backend_is_logged(self, create_runtime_ti, mock_supervisor_comms):
        """An unimportable custom backend is logged, and does not silently fall back to SMTP."""

        class FailingOperator(BaseOperator):
            def execute(self, context):
                raise AirflowFailException("Task failed on purpose")

        task = FailingOperator(
            task_id="bad_backend_task",
            email=["test@example.com"],
            email_on_failure=True,
        )

        runtime_ti = create_runtime_ti(task=task)
        context = runtime_ti.get_template_context()
        log = mock.MagicMock()

        with conf_vars({("email", "email_backend"): "airflow.does.not.Exist"}):
            # SmtpNotifier is patched to succeed, so a logged exception can only come from
            # resolving the custom backend -- without that, this would pass pre-fix too.
            with mock.patch(
                "airflow.providers.smtp.notifications.smtp.SmtpNotifier", autospec=True
            ) as mock_smtp_notifier:
                state, _, error = run(runtime_ti, context, log)
                # Must not raise even though the backend cannot be loaded.
                finalize(runtime_ti, state, context, log, error)

        mock_smtp_notifier.assert_not_called()
        log.exception.assert_called()

    @pytest.mark.enable_redact
    def test_rendered_templates_mask_secrets(self, create_runtime_ti, mock_supervisor_comms):
        """Test that secrets registered with mask_secret() are redacted in rendered template fields."""
        from unittest.mock import call

        from airflow.sdk._shared.secrets_masker import _secrets_masker
        from airflow.sdk.log import mask_secret

        _secrets_masker().add_mask("admin_user_12345", None)

        class CustomOperator(BaseOperator):
            template_fields = ("username", "region")

            def __init__(self, username, region, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.username = username
                self.region = region

            def execute(self, context):
                # Only mask username
                mask_secret(self.username)

        task = CustomOperator(
            task_id="test_masking",
            username="admin_user_12345",
            region="us-west-2",
        )

        runtime_ti = create_runtime_ti(task=task, dag_id="test_secrets_in_rtif")
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        assert (
            call(MaskSecret(value="admin_user_12345", name=None, type="MaskSecret"))
            in mock_supervisor_comms.send.mock_calls
        )
        # Region should not be masked
        assert (
            call(MaskSecret(value="us-west-2", name=None, type="MaskSecret"))
            not in mock_supervisor_comms.send.mock_calls
        )

        assert (
            call(
                msg=SetRenderedFields(
                    rendered_fields={"username": "***", "region": "us-west-2"},
                    type="SetRenderedFields",
                )
            )
            in mock_supervisor_comms.send.mock_calls
        )


class TestDagParamRuntime:
    DEFAULT_ARGS = {
        "owner": "test",
        "depends_on_past": True,
        "start_date": datetime.now(tz=timezone.utc),
        "retries": 1,
        "retry_delay": timedelta(minutes=1),
    }
    VALUE = 42

    def test_dag_param_resolves_from_task(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        """Test dagparam resolves on operator execution"""
        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        dag = DAG(dag_id="dag_with_dag_params", start_date=timezone.datetime(2024, 12, 3))
        dag.param("value", default="NOTSET")

        class CustomOperator(BaseOperator):
            def execute(self, context):
                assert dag.params["value"] == "NOTSET"

        task = CustomOperator(task_id="task_with_dag_params")
        runtime_ti = create_runtime_ti(task=task, dag_id="dag_with_dag_params")

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_called_once_with(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )

    def test_dag_param_dag_overwrite(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        """Test dag param is overwritten from dagrun config"""
        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        dag = DAG(dag_id="dag_with_dag_params_overwrite", start_date=timezone.datetime(2024, 12, 3))
        dag.param("value", default="NOTSET")

        class CustomOperator(BaseOperator):
            def execute(self, context):
                # important to use self.dag here
                assert self.dag.params["value"] == "new_value"

        # asserting on the default value when not set in dag run
        assert dag.params["value"] == "NOTSET"
        task = CustomOperator(task_id="task_with_dag_params_overwrite")

        # we reparse the dag here, and if conf passed, added as params
        runtime_ti = create_runtime_ti(
            task=task, dag_id="dag_with_dag_params_overwrite", conf={"value": "new_value"}
        )
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
        mock_supervisor_comms.send.assert_called_once_with(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )

    def test_dag_param_dag_default(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        """Test that dag param is correctly resolved by operator"""
        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        dag = DAG(
            dag_id="dag_with_dag_params_default",
            start_date=timezone.datetime(2024, 12, 3),
            params={"value": "test"},
        )

        class CustomOperator(BaseOperator):
            def execute(self, context):
                assert dag.params["value"] == "test"

        assert dag.params["value"] == "test"
        task = CustomOperator(task_id="task_with_dag_params_default")
        runtime_ti = create_runtime_ti(task=task, dag_id="dag_with_dag_params_default")

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
        mock_supervisor_comms.send.assert_called_once_with(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )

    def test_dag_param_resolves(
        self, create_runtime_ti, mock_supervisor_comms, time_machine, make_ti_context
    ):
        """Test that dag param is correctly resolved by operator"""

        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        @dag_decorator(schedule=None, start_date=timezone.datetime(2024, 12, 3))
        def dag_with_dag_params(value="NOTSET"):
            @task_decorator
            def dummy_task(val):
                return val

            class CustomOperator(BaseOperator):
                def execute(self, context):
                    assert self.dag.params["value"] == "NOTSET"

            _ = dummy_task(value)
            custom_task = CustomOperator(task_id="task_with_dag_params")
            self.operator = custom_task

        dag_with_dag_params()

        runtime_ti = create_runtime_ti(task=self.operator, dag_id="dag_with_dag_params")

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_called_once_with(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )

    def test_dag_param_dagrun_parameterized(
        self, create_runtime_ti, mock_supervisor_comms, time_machine, make_ti_context
    ):
        """Test that dag param is correctly overwritten when set in dag run"""

        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        @dag_decorator(schedule=None, start_date=timezone.datetime(2024, 12, 3))
        def dag_with_dag_params(value=self.VALUE):
            @task_decorator
            def dummy_task(val):
                return val

            assert isinstance(value, DagParam)

            class CustomOperator(BaseOperator):
                def execute(self, context):
                    assert self.dag.params["value"] == "new_value"

            _ = dummy_task(value)
            custom_task = CustomOperator(task_id="task_with_dag_params")
            self.operator = custom_task

        dag_with_dag_params()

        runtime_ti = create_runtime_ti(
            task=self.operator, dag_id="dag_with_dag_params", conf={"value": "new_value"}
        )

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_called_once_with(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )

    @pytest.mark.parametrize("value", [VALUE, 0])
    def test_set_params_for_dag(
        self, create_runtime_ti, mock_supervisor_comms, time_machine, make_ti_context, value
    ):
        """Test that dag param is correctly set when using dag decorator"""

        instant = timezone.datetime(2024, 12, 3, 10, 0)
        time_machine.move_to(instant, tick=False)

        @dag_decorator(schedule=None, start_date=timezone.datetime(2024, 12, 3))
        def dag_with_param(value=value):
            @task_decorator
            def return_num(num):
                return num

            xcom_arg = return_num(value)
            self.operator = xcom_arg.operator

        dag_with_param()

        runtime_ti = create_runtime_ti(task=self.operator, dag_id="dag_with_param", conf={"value": value})

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SucceedTask(state=TaskInstanceState.SUCCESS, end_date=instant, task_outlets=[], outlet_events=[]),
        )


class TestTaskRunnerCallsListeners:
    class CustomListener:
        def __init__(self):
            self.state = []
            self.component = None
            self.error = None

        @hookimpl
        def on_starting(self, component):
            self.component = component

        @hookimpl
        def on_task_instance_running(self, previous_state, task_instance):
            self.state.append(TaskInstanceState.RUNNING)

        @hookimpl
        def on_task_instance_success(self, previous_state, task_instance):
            self.state.append(TaskInstanceState.SUCCESS)

        @hookimpl
        def on_task_instance_failed(self, previous_state, task_instance, error):
            self.state.append(TaskInstanceState.FAILED)
            self.error = error

        @hookimpl
        def on_task_instance_skipped(self, previous_state, task_instance):
            self.state.append(TaskInstanceState.SKIPPED)

        @hookimpl
        def before_stopping(self, component):
            self.component = component

    class CustomOutletEventsListener:
        def __init__(self):
            self.outlet_events = []
            self.error = None

        def _add_outlet_events(self, context):
            outlets = context["outlets"]
            for outlet in outlets:
                self.outlet_events.append(context["outlet_events"][outlet])

        @hookimpl
        def on_task_instance_running(self, previous_state, task_instance):
            context = task_instance.get_template_context()
            self._add_outlet_events(context)

        @hookimpl
        def on_task_instance_success(self, previous_state, task_instance):
            context = task_instance.get_template_context()
            self._add_outlet_events(context)

        @hookimpl
        def on_task_instance_failed(self, previous_state, task_instance, error):
            context = task_instance.get_template_context()
            self._add_outlet_events(context)
            self.error = error

    def test_task_runner_calls_on_startup_before_stopping(
        self, make_ti_context, mocked_parse, mock_supervisor_comms, listener_manager
    ):
        listener = self.CustomListener()
        listener_manager(listener)

        class CustomOperator(BaseOperator):
            def execute(self, context):
                self.value = "something"

        task = CustomOperator(
            task_id="test_task_runner_calls_listeners", do_xcom_push=True, multiple_outputs=True
        )
        what = StartupDetails(
            ti=TaskInstance(
                id=uuid7(),
                task_id="templated_task",
                dag_id="basic_dag",
                run_id="c",
                try_number=1,
                dag_version_id=uuid7(),
                queue="default",
            ),
            dag_rel_path="",
            bundle_info=FAKE_BUNDLE,
            ti_context=make_ti_context(),
            start_date=timezone.utcnow(),
            sentry_integration="",
        )

        mock_supervisor_comms._get_response.return_value = what
        mocked_parse(what, "basic_dag", task)

        runtime_ti, context, log = startup(get_startup_details())
        assert runtime_ti is not None
        assert isinstance(listener.component, TaskRunnerMarker)
        del listener.component

        state, _, _ = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log)
        assert isinstance(listener.component, TaskRunnerMarker)

    def test_task_runner_calls_listeners_success(self, mocked_parse, mock_supervisor_comms, listener_manager):
        listener = self.CustomListener()
        listener_manager(listener)

        class CustomOperator(BaseOperator):
            def execute(self, context):
                self.value = "something"

        task = CustomOperator(
            task_id="test_task_runner_calls_listeners", do_xcom_push=True, multiple_outputs=True
        )
        dag = get_inline_dag(dag_id="test_dag", task=task)
        ti = TaskInstance(
            id=uuid7(),
            task_id=task.task_id,
            dag_id=dag.dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )

        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow()
        )
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, _ = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log)

        assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.SUCCESS]

    @pytest.mark.parametrize(
        "exception",
        [
            ValueError("oops"),
            SystemExit("oops"),
            AirflowException("oops"),
        ],
    )
    def test_task_runner_calls_listeners_failed(
        self, mocked_parse, mock_supervisor_comms, exception, listener_manager
    ):
        listener = self.CustomListener()
        listener_manager(listener)

        class CustomOperator(BaseOperator):
            def execute(self, context):
                raise exception

        task = CustomOperator(
            task_id="test_task_runner_calls_listeners_failed", do_xcom_push=True, multiple_outputs=True
        )
        dag = get_inline_dag(dag_id="test_dag", task=task)
        ti = TaskInstance(
            id=uuid7(),
            task_id=task.task_id,
            dag_id=dag.dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )

        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow()
        )
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, error = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log, error)

        assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.FAILED]
        assert listener.error == error

    def test_task_runner_calls_listeners_failed_when_terminal_send_fails(
        self, mocked_parse, mock_supervisor_comms, listener_manager
    ):
        """Callbacks/listeners must still fire when the FAILED terminal-state
        IPC send to the supervisor fails. The fail-closed exit is deferred to
        main() (signalled via `_terminal_state_send_failed` on the ti) so
        finalize() runs first.
        """
        listener = self.CustomListener()
        listener_manager(listener)

        class CustomOperator(BaseOperator):
            def execute(self, context):
                raise RuntimeError("task body failed")

        task = CustomOperator(task_id="failing_with_broken_ipc")
        dag = get_inline_dag(dag_id="test_dag", task=task)
        ti = TaskInstance(
            id=uuid7(),
            task_id=task.task_id,
            dag_id=dag.dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )

        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow()
        )

        def send_side_effect(msg=None, **kwargs):
            if isinstance(msg, TaskState):
                raise BrokenPipeError("supervisor IPC broken")
            return mock.DEFAULT

        mock_supervisor_comms.send.side_effect = send_side_effect

        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, error = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log, error)

        assert state == TaskInstanceState.FAILED
        assert runtime_ti._terminal_state_send_failed is True
        assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.FAILED]
        assert listener.error == error

    def test_task_runner_calls_listeners_skipped(self, mocked_parse, mock_supervisor_comms, listener_manager):
        listener = self.CustomListener()
        listener_manager(listener)

        class CustomOperator(BaseOperator):
            def execute(self, context):
                raise AirflowSkipException("Task intentionally skipped")

        task = CustomOperator(
            task_id="test_task_runner_calls_listeners_skipped", do_xcom_push=True, multiple_outputs=True
        )
        dag = get_inline_dag(dag_id="test_dag", task=task)
        ti = TaskInstance(
            id=uuid7(),
            task_id=task.task_id,
            dag_id=dag.dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )

        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow()
        )
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, _ = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log)

        assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.SKIPPED]

    def test_listener_access_outlet_event_on_running_and_success(
        self, mocked_parse, mock_supervisor_comms, listener_manager
    ):
        """Test listener can access outlet events through invoking get_template_context() while task running and success"""
        listener = self.CustomOutletEventsListener()
        listener_manager(listener)

        test_asset = Asset("test-asset")
        test_key = AssetUniqueKey(name="test-asset", uri="test-asset")
        test_extra = {"name1": "value1", "nested_obj": {"name2": "value2"}}

        class Producer(BaseOperator):
            def execute(self, context):
                outlet_events = context["outlet_events"]
                outlet_events[test_asset].extra = test_extra

        task = Producer(
            task_id="test_listener_access_outlet_event_on_running_and_success", outlets=[test_asset]
        )
        dag = get_inline_dag(dag_id="test_dag", task=task)
        ti = TaskInstance(
            id=uuid7(),
            task_id=task.task_id,
            dag_id=dag.dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )

        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow()
        )

        log = mock.MagicMock()
        context = runtime_ti.get_template_context()

        with mock.patch(
            "airflow.sdk.execution_time.task_runner._validate_task_inlets_and_outlets"
        ) as validate_mock:
            state, _, _ = run(runtime_ti, context, log)

        validate_mock.assert_called_once()

        outlet_event_accessor = listener.outlet_events.pop()
        assert outlet_event_accessor.key == test_key
        assert outlet_event_accessor.extra == test_extra

        finalize(runtime_ti, state, context, log)

        outlet_event_accessor = listener.outlet_events.pop()
        assert outlet_event_accessor.key == test_key
        assert outlet_event_accessor.extra == test_extra

    @pytest.mark.parametrize(
        "exception",
        [
            ValueError("oops"),
            SystemExit("oops"),
            AirflowException("oops"),
        ],
        ids=["ValueError", "SystemExit", "AirflowException"],
    )
    def test_listener_access_outlet_event_on_failed(
        self, mocked_parse, mock_supervisor_comms, exception, listener_manager
    ):
        """Test listener can access outlet events through invoking get_template_context() while task failed"""
        listener = self.CustomOutletEventsListener()
        listener_manager(listener)

        test_asset = Asset("test-asset")
        test_key = AssetUniqueKey(name="test-asset", uri="test-asset")
        test_extra = {"name1": "value1", "nested_obj": {"name2": "value2"}}

        class Producer(BaseOperator):
            def execute(self, context):
                outlet_events = context["outlet_events"]
                outlet_events[test_asset].extra = test_extra
                raise exception

        task = Producer(task_id="test_listener_access_outlet_event_on_failed", outlets=[test_asset])
        dag = get_inline_dag(dag_id="test_dag", task=task)
        ti = TaskInstance(
            id=uuid7(),
            task_id=task.task_id,
            dag_id=dag.dag_id,
            run_id="test_run",
            try_number=1,
            dag_version_id=uuid7(),
        )

        runtime_ti = RuntimeTaskInstance.model_construct(
            **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow()
        )

        log = mock.MagicMock()
        context = runtime_ti.get_template_context()

        with mock.patch(
            "airflow.sdk.execution_time.task_runner._validate_task_inlets_and_outlets"
        ) as validate_mock:
            state, _, error = run(runtime_ti, context, log)

        validate_mock.assert_called_once()

        outlet_event_accessor = listener.outlet_events.pop()
        assert outlet_event_accessor.key == test_key
        assert outlet_event_accessor.extra == test_extra

        finalize(runtime_ti, state, context, log, error)

        outlet_event_accessor = listener.outlet_events.pop()
        assert outlet_event_accessor.key == test_key
        assert outlet_event_accessor.extra == test_extra

        assert listener.error == error


@pytest.mark.usefixtures("mock_supervisor_comms")
class TestTaskRunnerCallsCallbacks:
    class _Failure(Exception):
        """Exception raised in a failed execution and received by the failure callback."""

    def _execute_success(self, context):
        self.results.append("execute success")

    def _execute_skipped(self, context):
        from airflow.sdk.exceptions import AirflowSkipException

        self.results.append("execute skipped")
        raise AirflowSkipException

    def _execute_failure(self, context):
        self.results.append("execute failure")
        raise self._Failure("sorry!")

    @pytest.mark.parametrize(
        ("execute_impl", "should_retry", "expected_state", "expected_results"),
        [
            pytest.param(
                _execute_success,
                False,
                TaskInstanceState.SUCCESS,
                ["on-execute callback", "execute success", "on-success callback"],
                id="success",
            ),
            pytest.param(
                _execute_skipped,
                False,
                TaskInstanceState.SKIPPED,
                ["on-execute callback", "execute skipped", "on-skipped callback"],
                id="skipped",
            ),
            pytest.param(
                _execute_failure,
                False,
                TaskInstanceState.FAILED,
                ["on-execute callback", "execute failure", "on-failure callback"],
                id="failure",
            ),
            pytest.param(
                _execute_failure,
                True,
                TaskInstanceState.UP_FOR_RETRY,
                ["on-execute callback", "execute failure", "on-retry callback"],
                id="retry",
            ),
        ],
    )
    def test_task_runner_calls_callback(
        self,
        create_runtime_ti,
        execute_impl,
        should_retry,
        expected_state,
        expected_results,
    ):
        collected_results = []

        def custom_callback(context, *, kind):
            collected_results.append(f"on-{kind} callback")

        def failure_callback(context):
            custom_callback(context, kind="failure")
            assert isinstance(context["exception"], self._Failure)

        class CustomOperator(BaseOperator):
            results = collected_results
            execute = execute_impl

        task = CustomOperator(
            task_id="task",
            on_execute_callback=functools.partial(custom_callback, kind="execute"),
            on_skipped_callback=functools.partial(custom_callback, kind="skipped"),
            on_success_callback=functools.partial(custom_callback, kind="success"),
            on_failure_callback=failure_callback,
            on_retry_callback=functools.partial(custom_callback, kind="retry"),
        )
        runtime_ti = create_runtime_ti(dag_id="dag", task=task, should_retry=should_retry)
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, error = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log, error)

        assert state == expected_state
        assert collected_results == expected_results

    @pytest.mark.parametrize(
        ("base_url", "expected_url"),
        [
            ("http://localhost:8080/", "http://localhost:8080/dags/test_dag/runs/test_run/tasks/test_task"),
            ("http://localhost:8080", "http://localhost:8080/dags/test_dag/runs/test_run/tasks/test_task"),
            (
                "https://airflow.example.com/",
                "https://airflow.example.com/dags/test_dag/runs/test_run/tasks/test_task",
            ),
            (
                "https://airflow.example.com",
                "https://airflow.example.com/dags/test_dag/runs/test_run/tasks/test_task",
            ),
        ],
        ids=["localhost_with_slash", "localhost_no_slash", "domain_with_slash", "domain_no_slash"],
    )
    def test_runtime_task_instance_log_url_property(self, create_runtime_ti, base_url, expected_url):
        """Test that RuntimeTaskInstance.log_url property correctly handles various base_url formats."""
        task = BaseOperator(task_id="test_task")
        runtime_ti = create_runtime_ti(task=task, dag_id="test_dag", run_id="test_run", try_number=0)

        with conf_vars({("api", "base_url"): base_url}):
            log_url = runtime_ti.log_url
            assert log_url == expected_url

    def test_task_runner_on_failure_callback_context(self, create_runtime_ti):
        """Test that on_failure_callback context has end_date and duration."""

        def failure_callback(context):
            ti = context["task_instance"]
            assert isinstance(ti.end_date, datetime)
            duration = (ti.end_date - ti.start_date).total_seconds()
            assert duration is not None
            assert duration >= 0

        class FailingOperator(BaseOperator):
            def execute(self, context):
                raise AirflowException("Failing task")

        task = FailingOperator(task_id="failing_task", on_failure_callback=failure_callback)
        runtime_ti = create_runtime_ti(dag_id="dag", task=task)
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, error = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log, error)

        assert state == TaskInstanceState.FAILED

    def test_task_runner_on_success_callback_context(self, create_runtime_ti):
        """Test that on_success_callback context has end_date and duration."""
        callback_data = {}  # Store callback data for inspection

        def success_callback(context):
            ti = context["task_instance"]
            callback_data["end_date"] = ti.end_date
            callback_data["duration"] = (ti.end_date - ti.start_date).total_seconds() if ti.end_date else None
            callback_data["start_date"] = ti.start_date

        class SuccessOperator(BaseOperator):
            def execute(self, context):
                return "success"

        task = SuccessOperator(task_id="success_task", on_success_callback=success_callback)
        runtime_ti = create_runtime_ti(dag_id="dag", task=task)
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()

        state, _, error = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log, error)

        assert state == TaskInstanceState.SUCCESS

        # Verify callback was called and data was captured
        assert "end_date" in callback_data, "Success callback should have been called"
        assert isinstance(callback_data["end_date"], datetime), (
            f"end_date should be datetime, got {type(callback_data['end_date'])}"
        )
        assert callback_data["duration"] is not None, (
            f"duration should not be None, got {callback_data['duration']}"
        )
        assert callback_data["duration"] >= 0, f"duration should be >= 0, got {callback_data['duration']}"

    def test_task_runner_both_callbacks_have_timing_info(self, create_runtime_ti):
        """Test that both success and failure callbacks receive accurate timing information."""
        import time

        success_data = {}
        failure_data = {}

        def success_callback(context):
            ti = context["task_instance"]
            success_data["end_date"] = ti.end_date
            success_data["start_date"] = ti.start_date
            success_data["duration"] = (ti.end_date - ti.start_date).total_seconds() if ti.end_date else None

        def failure_callback(context):
            ti = context["task_instance"]
            failure_data["end_date"] = ti.end_date
            failure_data["start_date"] = ti.start_date
            failure_data["duration"] = (ti.end_date - ti.start_date).total_seconds() if ti.end_date else None

        # Test success callback
        class SuccessOperator(BaseOperator):
            def execute(self, context):
                time.sleep(0.01)  # Add small delay to ensure measurable duration
                return "success"

        success_task = SuccessOperator(task_id="success_task", on_success_callback=success_callback)
        success_runtime_ti = create_runtime_ti(dag_id="dag", task=success_task)
        success_log = mock.MagicMock()
        success_context = success_runtime_ti.get_template_context()

        success_state, _, success_error = run(success_runtime_ti, success_context, success_log)
        finalize(success_runtime_ti, success_state, success_context, success_log, success_error)

        # Test failure callback
        class FailureOperator(BaseOperator):
            def execute(self, context):
                time.sleep(0.01)  # Add small delay to ensure measurable duration
                raise AirflowException("Test failure")

        failure_task = FailureOperator(task_id="failure_task", on_failure_callback=failure_callback)
        failure_runtime_ti = create_runtime_ti(dag_id="dag", task=failure_task)
        failure_log = mock.MagicMock()
        failure_context = failure_runtime_ti.get_template_context()

        failure_state, _, failure_error = run(failure_runtime_ti, failure_context, failure_log)
        finalize(failure_runtime_ti, failure_state, failure_context, failure_log, failure_error)

        # Assertions for success callback
        assert success_state == TaskInstanceState.SUCCESS
        assert "end_date" in success_data, "Success callback should have been called"
        assert isinstance(success_data["end_date"], datetime)
        assert isinstance(success_data["start_date"], datetime)
        assert success_data["duration"] is not None
        assert success_data["duration"] >= 0.01, (
            f"Success duration should be >= 0.01, got {success_data['duration']}"
        )

        # Assertions for failure callback
        assert failure_state == TaskInstanceState.FAILED
        assert "end_date" in failure_data, "Failure callback should have been called"
        assert isinstance(failure_data["end_date"], datetime)
        assert isinstance(failure_data["start_date"], datetime)
        assert failure_data["duration"] is not None
        assert failure_data["duration"] >= 0.01, (
            f"Failure duration should be >= 0.01, got {failure_data['duration']}"
        )

    @pytest.mark.parametrize(
        (
            "callback_to_test",
            "execute_impl",
            "should_retry",
            "expected_state",
            "expected_results",
            "extra_exceptions",
        ),
        [
            pytest.param(
                "on_success_callback",
                _execute_success,
                False,
                TaskInstanceState.SUCCESS,
                ["on-execute 1", "on-execute 3", "execute success", "on-success 1", "on-success 3"],
                [],
                id="success",
            ),
            pytest.param(
                "on_skipped_callback",
                _execute_skipped,
                False,
                TaskInstanceState.SKIPPED,
                ["on-execute 1", "on-execute 3", "execute skipped", "on-skipped 1", "on-skipped 3"],
                [],
                id="skipped",
            ),
            pytest.param(
                "on_failure_callback",
                _execute_failure,
                False,
                TaskInstanceState.FAILED,
                ["on-execute 1", "on-execute 3", "execute failure", "on-failure 1", "on-failure 3"],
                [(1, mock.call("Task failed with exception"))],
                id="failure",
            ),
            pytest.param(
                "on_retry_callback",
                _execute_failure,
                True,
                TaskInstanceState.UP_FOR_RETRY,
                ["on-execute 1", "on-execute 3", "execute failure", "on-retry 1", "on-retry 3"],
                [(1, mock.call("Task failed with exception"))],
                id="retry",
            ),
        ],
    )
    def test_task_runner_not_fail_on_failed_callback(
        self,
        create_runtime_ti,
        callback_to_test,
        execute_impl,
        should_retry,
        expected_state,
        expected_results,
        extra_exceptions,
    ):
        collected_results = []

        def custom_callback_1(context, *, kind):
            collected_results.append(f"on-{kind} 1")

        def custom_callback_2(context, *, kind):
            raise Exception("sorry!")

        def custom_callback_3(context, *, kind):
            collected_results.append(f"on-{kind} 3")

        class CustomOperator(BaseOperator):
            results = collected_results
            execute = execute_impl

        task = CustomOperator(
            task_id="task",
            on_execute_callback=[
                functools.partial(custom_callback_1, kind="execute"),
                functools.partial(custom_callback_2, kind="execute"),
                functools.partial(custom_callback_3, kind="execute"),
            ],
            on_skipped_callback=[
                functools.partial(custom_callback_1, kind="skipped"),
                functools.partial(custom_callback_2, kind="skipped"),
                functools.partial(custom_callback_3, kind="skipped"),
            ],
            on_success_callback=[
                functools.partial(custom_callback_1, kind="success"),
                functools.partial(custom_callback_2, kind="success"),
                functools.partial(custom_callback_3, kind="success"),
            ],
            on_failure_callback=[
                functools.partial(custom_callback_1, kind="failure"),
                functools.partial(custom_callback_2, kind="failure"),
                functools.partial(custom_callback_3, kind="failure"),
            ],
            on_retry_callback=[
                functools.partial(custom_callback_1, kind="retry"),
                functools.partial(custom_callback_2, kind="retry"),
                functools.partial(custom_callback_3, kind="retry"),
            ],
        )
        runtime_ti = create_runtime_ti(dag_id="dag", task=task, should_retry=should_retry)
        log = mock.MagicMock()
        context = runtime_ti.get_template_context()
        state, _, error = run(runtime_ti, context, log)
        finalize(runtime_ti, state, context, log, error)

        assert state == expected_state, error
        assert collected_results == expected_results

        expected_exception_logs = [
            mock.call("Failed to run task callback", kind="on_execute_callback", index=1, callback=mock.ANY),
            mock.call("Failed to run task callback", kind=callback_to_test, index=1, callback=mock.ANY),
        ]
        for index, calls in extra_exceptions:
            expected_exception_logs.insert(index, calls)
        assert log.exception.mock_calls == expected_exception_logs


class TestTriggerDagRunOperator:
    """Tests to verify various aspects of TriggerDagRunOperator"""

    @time_machine.travel("2025-01-01 00:00:00", tick=False)
    def test_handle_trigger_dag_run(self, create_runtime_ti, mock_supervisor_comms):
        """Test that TriggerDagRunOperator (with default args) sends the correct message to the Supervisor"""
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
        )
        ti = create_runtime_ti(dag_id="test_handle_trigger_dag_run", run_id="test_run", task=task)

        log = mock.MagicMock()

        state, msg, _ = run(ti, ti.get_template_context(), log)

        assert state == TaskInstanceState.SUCCESS
        assert msg.state == TaskInstanceState.SUCCESS

        expected_calls = [
            mock.call.send(
                msg=TriggerDagRun(
                    dag_id="test_dag",
                    run_id="test_run_id",
                    reset_dag_run=False,
                    logical_date=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
                ),
            ),
            mock.call.send(
                msg=SetXCom(
                    key="trigger_run_id",
                    value="test_run_id",
                    dag_id="test_handle_trigger_dag_run",
                    task_id="test_task",
                    run_id="test_run",
                    map_index=-1,
                ),
            ),
        ]
        mock_supervisor_comms.assert_has_calls(expected_calls)

    @pytest.mark.parametrize(
        ("skip_when_already_exists", "expected_state"),
        [
            (True, TaskInstanceState.SKIPPED),
            (False, TaskInstanceState.FAILED),
        ],
    )
    @time_machine.travel("2025-01-01 00:00:00", tick=False)
    def test_handle_trigger_dag_run_conflict(
        self, skip_when_already_exists, expected_state, create_runtime_ti, mock_supervisor_comms
    ):
        """Test that TriggerDagRunOperator (when dagrun already exists) sends the correct message to the Supervisor"""
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
            skip_when_already_exists=skip_when_already_exists,
        )
        ti = create_runtime_ti(dag_id="test_handle_trigger_dag_run_conflict", run_id="test_run", task=task)

        log = mock.MagicMock()
        mock_supervisor_comms.send.return_value = ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS)
        state, msg, _ = run(ti, ti.get_template_context(), log)

        assert state == expected_state
        assert msg.state == expected_state

        expected_calls = [
            mock.call.send(
                msg=TriggerDagRun(
                    dag_id="test_dag",
                    logical_date=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
                    run_id="test_run_id",
                    reset_dag_run=False,
                ),
            ),
        ]
        mock_supervisor_comms.assert_has_calls(expected_calls)

    @time_machine.travel("2025-01-01 00:00:00", tick=False)
    def test_handle_trigger_dag_run_surfaces_original_error(self, create_runtime_ti, mock_supervisor_comms):
        """
        When an ``except`` handler in ``run()`` raises, the original exception must be
        surfaced as the task failure, not an ``UnboundLocalError`` on ``state``.
        """
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        class _TriggerSendError(Exception):
            pass

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="missing_dag",
            trigger_run_id="test_run_id",
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_surfaces_original_error",
            run_id="test_run",
            task=task,
        )

        def _send(msg=None, **kwargs):
            # Fail only the TriggerDagRun comms (the call that would 404 for a
            # missing DAG); let every earlier send behave as the default mock.
            if isinstance(msg, TriggerDagRun):
                raise _TriggerSendError("simulated 404 for missing target DAG")
            return mock.DEFAULT

        mock_supervisor_comms.send.side_effect = _send

        log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)

        state, _, error = run(ti, ti.get_template_context(), log)

        assert state == TaskInstanceState.FAILED
        assert isinstance(error, _TriggerSendError)

    @pytest.mark.parametrize(
        ("should_retry", "expected_state"),
        [
            (True, TaskInstanceState.UP_FOR_RETRY),
            (False, TaskInstanceState.FAILED),
        ],
    )
    @time_machine.travel("2025-01-01 00:00:00", tick=False)
    def test_handle_trigger_dag_run_missing_dag_honours_retries_and_callbacks(
        self, should_retry, expected_state, create_runtime_ti, mock_supervisor_comms
    ):
        """
        A 404 from the API server for a missing target Dag must be handled like any other
        task failure: the retry decision is made and ``finalize()`` fires the callbacks.

        Regression test for https://github.com/apache/airflow/issues/70683 -- the
        ``AirflowRuntimeError`` was raised from inside ``run()``'s
        ``except DagRunTriggerException`` handler, so it escaped ``run()`` without
        evaluating retries or running ``on_failure_callback`` / ``on_retry_callback``.
        """
        callbacks_run = []

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="this_dag_does_not_exist",
            trigger_run_id="test_run_id",
            on_failure_callback=lambda context: callbacks_run.append("failure"),
            on_retry_callback=lambda context: callbacks_run.append("retry"),
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_missing_dag",
            run_id="test_run",
            task=task,
            should_retry=should_retry,
        )

        not_found = AirflowRuntimeError(
            error=ErrorResponse(
                error=ErrorType.API_SERVER_ERROR,
                detail={
                    "status_code": 404,
                    "message": "Client error message: Dag with dag_id: 'this_dag_does_not_exist' not found",
                },
            )
        )

        def _send(msg=None, **kwargs):
            if isinstance(msg, TriggerDagRun):
                raise not_found
            return mock.DEFAULT

        mock_supervisor_comms.send.side_effect = _send

        log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger)
        context = ti.get_template_context()

        state, msg, error = run(ti, context, log)

        assert state == expected_state
        assert error is not_found
        # The terminal message is what the server acts on, so assert the shape too:
        # RetryTask drives UP_FOR_RETRY, TaskState(FAILED) ends the try.
        if should_retry:
            assert isinstance(msg, RetryTask)
        else:
            assert isinstance(msg, TaskState)
            assert msg.state == TaskInstanceState.FAILED

        context["exception"] = error
        finalize(ti, state, context, log, error)

        assert callbacks_run == ["retry" if should_retry else "failure"]

    @pytest.mark.parametrize(
        ("allowed_states", "failed_states", "target_dr_state", "expected_task_state"),
        [
            (None, None, DagRunState.FAILED, TaskInstanceState.FAILED),
            (None, None, DagRunState.SUCCESS, TaskInstanceState.SUCCESS),
            ([DagRunState.FAILED], [], DagRunState.FAILED, DagRunState.SUCCESS),
            ([DagRunState.FAILED], None, DagRunState.FAILED, DagRunState.FAILED),
            ([DagRunState.SUCCESS], None, DagRunState.FAILED, DagRunState.FAILED),
        ],
    )
    @time_machine.travel("2025-01-01 00:00:00", tick=False)
    def test_handle_trigger_dag_run_wait_for_completion(
        self,
        allowed_states,
        failed_states,
        target_dr_state,
        expected_task_state,
        create_runtime_ti,
        mock_supervisor_comms,
    ):
        """
        Test that TriggerDagRunOperator (with wait_for_completion) sends the correct message to the Supervisor

        It also polls the Supervisor for the DagRun state until it completes execution.
        """
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
            poke_interval=5,
            wait_for_completion=True,
            allowed_states=allowed_states,
            failed_states=failed_states,
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_wait_for_completion", run_id="test_run", task=task
        )

        log = mock.MagicMock()
        mock_supervisor_comms.send.side_effect = [
            # Set RTIF
            None,
            # Account for the extra link XCom message sent by TriggerDagRunLink
            None,
            # Successful Dag Run trigger
            OKResponse(ok=True),
            # Set XCOM,
            None,
            # Dag Run is still running
            DagRunStateResult(state=DagRunState.RUNNING),
            # Dag Run completes execution on the next poll
            DagRunStateResult(state=target_dr_state),
            # Succeed/Fail task
            None,
        ]
        with mock.patch("time.sleep", return_value=None):
            state, msg, _ = run(ti, ti.get_template_context(), log)

        assert state == expected_task_state
        assert msg.state == expected_task_state

        expected_calls = [
            mock.call.send(
                msg=SetXCom(
                    key="_link_TriggerDagRunLink",
                    value="/dags/test_dag/runs/test_run_id",
                    dag_id="test_handle_trigger_dag_run_wait_for_completion",
                    task_id="test_task",
                    run_id="test_run",
                    map_index=-1,
                ),
            ),
            mock.call.send(
                msg=TriggerDagRun(
                    dag_id="test_dag",
                    run_id="test_run_id",
                    logical_date=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
                ),
            ),
            mock.call.send(
                msg=SetXCom(
                    key="trigger_run_id",
                    value="test_run_id",
                    dag_id="test_handle_trigger_dag_run_wait_for_completion",
                    task_id="test_task",
                    run_id="test_run",
                    map_index=-1,
                ),
            ),
            mock.call.send(
                msg=GetDagRunState(
                    dag_id="test_dag",
                    run_id="test_run_id",
                ),
            ),
            mock.call.send(
                msg=GetDagRunState(
                    dag_id="test_dag",
                    run_id="test_run_id",
                ),
            ),
        ]
        mock_supervisor_comms.assert_has_calls(expected_calls)

    def test_handle_trigger_dag_run_wait_for_completion_failed_state_retries(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """When triggered DAG reaches a failed state, task should honor retry eligibility."""
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
            poke_interval=5,
            wait_for_completion=True,
            allowed_states=[DagRunState.SUCCESS],
            failed_states=[DagRunState.FAILED],
            deferrable=False,
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_wait_for_completion_failed_state_retries",
            run_id="test_run",
            task=task,
            should_retry=True,
        )

        def _send_side_effect(*args, **kwargs):
            msg = kwargs.get("msg")
            if msg is None and args:
                msg = args[0]
            if isinstance(msg, TriggerDagRun):
                return OKResponse(ok=True)
            if isinstance(msg, GetDagRunState):
                return DagRunStateResult(state=DagRunState.FAILED)
            return None

        mock_supervisor_comms.send.side_effect = _send_side_effect

        log = mock.MagicMock()
        with mock.patch("time.sleep", return_value=None):
            state, _, _ = run(ti, ti.get_template_context(), log)

        assert state == TaskInstanceState.UP_FOR_RETRY

    def test_handle_trigger_dag_run_wait_for_completion_failed_state_retry_policy_fail(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """A retry_policy returning FAIL fails the task even when retries are still available.

        This mirrors the deferrable path, where DagStateTrigger -> execute_complete raises
        AirflowException and the configured retry_policy is consulted via run()'s handler.
        """
        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
            poke_interval=5,
            wait_for_completion=True,
            allowed_states=[DagRunState.SUCCESS],
            failed_states=[DagRunState.FAILED],
            deferrable=False,
            retry_policy=ExceptionRetryPolicy(
                rules=[
                    RetryRule(exception=AirflowException, action=RetryAction.FAIL, reason="not retryable")
                ],
            ),
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_wait_for_completion_failed_state_retry_policy_fail",
            run_id="test_run",
            task=task,
            should_retry=True,
        )

        def _send_side_effect(*args, **kwargs):
            msg = kwargs.get("msg")
            if msg is None and args:
                msg = args[0]
            if isinstance(msg, TriggerDagRun):
                return OKResponse(ok=True)
            if isinstance(msg, GetDagRunState):
                return DagRunStateResult(state=DagRunState.FAILED)
            return None

        mock_supervisor_comms.send.side_effect = _send_side_effect

        log = mock.MagicMock()
        with mock.patch("time.sleep", return_value=None):
            state, _, _ = run(ti, ti.get_template_context(), log)

        assert state == TaskInstanceState.FAILED

    def test_handle_trigger_dag_run_wait_for_completion_failed_state_retry_policy_delay(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """A retry_policy returning RETRY forwards its custom delay and reason on the RetryTask."""
        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
            poke_interval=5,
            wait_for_completion=True,
            allowed_states=[DagRunState.SUCCESS],
            failed_states=[DagRunState.FAILED],
            deferrable=False,
            retry_policy=ExceptionRetryPolicy(
                rules=[
                    RetryRule(
                        exception=AirflowException,
                        action=RetryAction.RETRY,
                        retry_delay=timedelta(minutes=7),
                        reason="backing off",
                    ),
                ],
            ),
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_wait_for_completion_failed_state_retry_policy_delay",
            run_id="test_run",
            task=task,
            should_retry=True,
        )

        def _send_side_effect(*args, **kwargs):
            msg = kwargs.get("msg")
            if msg is None and args:
                msg = args[0]
            if isinstance(msg, TriggerDagRun):
                return OKResponse(ok=True)
            if isinstance(msg, GetDagRunState):
                return DagRunStateResult(state=DagRunState.FAILED)
            return None

        mock_supervisor_comms.send.side_effect = _send_side_effect

        log = mock.MagicMock()
        with mock.patch("time.sleep", return_value=None):
            state, msg, _ = run(ti, ti.get_template_context(), log)

        assert state == TaskInstanceState.UP_FOR_RETRY
        assert msg.retry_delay_seconds == timedelta(minutes=7).total_seconds()
        assert msg.retry_reason == "backing off"

    @pytest.mark.parametrize(
        ("allowed_states", "failed_states", "intermediate_state"),
        [
            ([DagRunState.SUCCESS], None, TaskInstanceState.SUCCESS),
        ],
    )
    def test_handle_trigger_dag_run_deferred(
        self,
        allowed_states,
        failed_states,
        intermediate_state,
        create_runtime_ti,
        mock_supervisor_comms,
    ):
        """
        Test that TriggerDagRunOperator does not defer when wait_for_completion=False
        """
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="test_run_id",
            poke_interval=5,
            wait_for_completion=False,
            allowed_states=allowed_states,
            failed_states=failed_states,
            deferrable=True,
        )
        ti = create_runtime_ti(dag_id="test_handle_trigger_dag_run_deferred", run_id="test_run", task=task)

        log = mock.MagicMock()
        with mock.patch("time.sleep", return_value=None):
            state, msg, _ = run(ti, ti.get_template_context(), log)

        assert state == intermediate_state

    @time_machine.travel("2025-01-01 00:00:00", tick=False)
    def test_handle_trigger_dag_run_deferred_with_reset_uses_run_id_only(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """
        Test that TriggerDagRunOperator with deferrable=True and reset_dag_run=True
        creates a DagStateTrigger with execution_dates=None.

        This prevents the bug where reset_dag_run preserves the original logical_date
        in the database but the trigger queries with a newly calculated logical_date,
        causing a mismatch that makes the trigger never find the dag run.
        """
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

        task = TriggerDagRunOperator(
            task_id="test_task",
            trigger_dag_id="test_dag",
            trigger_run_id="fixed_run_id",
            wait_for_completion=True,
            deferrable=True,
            reset_dag_run=True,
            poke_interval=5,
        )
        ti = create_runtime_ti(
            dag_id="test_handle_trigger_dag_run_deferred_reset", run_id="test_run", task=task
        )

        log = mock.MagicMock()
        state, msg, _ = run(ti, ti.get_template_context(), log)

        # Task should be deferred
        assert state == TaskInstanceState.DEFERRED
        assert isinstance(msg, DeferTask)

        # Verify the DeferTask message structure
        assert msg.classpath == "airflow.providers.standard.triggers.external_task.DagStateTrigger"
        assert msg.next_method == "execute_complete"

        # Critical assertion: execution_dates should be None to avoid logical_date mismatch
        # when reset_dag_run=True. The run_id alone is sufficient and unique.
        trigger_kwargs = msg.trigger_kwargs
        assert trigger_kwargs["execution_dates"] is None, (
            "execution_dates should be None when using run_ids. "
            "When reset_dag_run=True, the logical_date in the database may differ from "
            "the newly calculated logical_date, causing the trigger to never find the run."
        )
        assert trigger_kwargs["run_ids"] == ["fixed_run_id"]
        assert trigger_kwargs["dag_id"] == "test_dag"
        assert trigger_kwargs["poll_interval"] == 5

        # Also verify it was sent to supervisor
        mock_supervisor_comms.send.assert_any_call(msg)


class TestTaskInstanceMetrics:
    def test_ti_start_metric_emitted(self, create_runtime_ti, mock_supervisor_comms):
        """Test that ti.start metric is emitted at the beginning of task."""
        task = PythonOperator(task_id="test", python_callable=lambda: "success")
        ti = create_runtime_ti(task=task)

        with mock.patch("airflow.sdk._shared.observability.metrics.stats._get_backend") as mock_get_backend:
            backend = mock.MagicMock(spec=StatsLogger)
            mock_get_backend.return_value = backend
            run(ti, context=ti.get_template_context(), log=mock.MagicMock())

            # verify ti.start was called in legacy format
            backend.incr.assert_any_call(f"ti.start.{ti.dag_id}.{ti.task_id}")
            # verify ti.start was called in tagged format
            backend.incr.assert_any_call(
                "ti.start",
                tags={"dag_id": ti.dag_id, "task_id": ti.task_id},
            )

    @pytest.mark.parametrize(
        ("task_callable", "expected_state"),
        [
            pytest.param(lambda: "success", "success", id="success"),
            pytest.param(lambda: (_ for _ in ()).throw(AirflowSkipException()), "skipped", id="skipped"),
            pytest.param(lambda: (_ for _ in ()).throw(AirflowFailException("fail")), "failed", id="failed"),
            pytest.param(lambda: 1 / 0, "failed", id="zero_division"),
        ],
    )
    def test_ti_finish_metric_emitted_for_terminal_states(
        self, task_callable, expected_state, create_runtime_ti, mock_supervisor_comms
    ):
        """Test that ti.finish metric is emitted for all terminal states."""
        task = PythonOperator(task_id="test", python_callable=task_callable)
        ti = create_runtime_ti(task=task)

        with mock.patch("airflow.sdk._shared.observability.metrics.stats._get_backend") as mock_get_backend:
            backend = mock.MagicMock(spec=StatsLogger)
            mock_get_backend.return_value = backend
            run(ti, context=ti.get_template_context(), log=mock.MagicMock())

            # verify ti.finish was called in legacy format
            backend.incr.assert_any_call(f"ti.finish.{ti.dag_id}.{ti.task_id}.{expected_state}")
            # verify ti.finish was called in tagged format
            backend.incr.assert_any_call(
                "ti.finish",
                tags={"dag_id": ti.dag_id, "task_id": ti.task_id, "state": expected_state},
            )

    def test_operator_successes_metrics_emitted(self, create_runtime_ti, mock_supervisor_comms):
        """Test that operator_successes and ti_successes metrics are emitted on task success."""
        task = PythonOperator(task_id="test", python_callable=lambda: "success")
        ti = create_runtime_ti(task=task)

        with mock.patch("airflow.sdk._shared.observability.metrics.stats._get_backend") as mock_get_backend:
            backend = mock.MagicMock(spec=StatsLogger)
            mock_get_backend.return_value = backend
            run(ti, context=ti.get_template_context(), log=mock.MagicMock())

            stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}

            # verify operator_successes in legacy format
            backend.incr.assert_any_call("operator_successes_PythonOperator", tags=stats_tags)
            # verify operator_successes in tagged format
            backend.incr.assert_any_call(
                "operator_successes",
                tags={**stats_tags, "operator_name": "PythonOperator"},
            )
            backend.incr.assert_any_call("ti_successes", tags=stats_tags)

    def test_operator_failures_metrics_emitted(self, create_runtime_ti, mock_supervisor_comms):
        """Test that operator_failures and ti_failures metrics are emitted on task failure."""
        task = PythonOperator(task_id="test", python_callable=lambda: 1 / 0)
        ti = create_runtime_ti(task=task)

        with mock.patch("airflow.sdk._shared.observability.metrics.stats._get_backend") as mock_get_backend:
            backend = mock.MagicMock(spec=StatsLogger)
            mock_get_backend.return_value = backend
            run(ti, context=ti.get_template_context(), log=mock.MagicMock())

            stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}

            # verify operator_failures in legacy format
            backend.incr.assert_any_call("operator_failures_PythonOperator", tags=stats_tags)
            # verify operator_failures in tagged format
            backend.incr.assert_any_call(
                "operator_failures",
                tags={**stats_tags, "operator_name": "PythonOperator"},
            )
            backend.incr.assert_any_call("ti_failures", tags=stats_tags)

    @pytest.mark.parametrize(
        ("team_name", "expected_tags_extra"),
        [
            pytest.param("my_team", {"team_name": "my_team"}, id="with_team"),
            pytest.param(None, {}, id="without_team"),
        ],
    )
    def test_ti_start_metric_respects_team_name(
        self, team_name, expected_tags_extra, create_runtime_ti, mock_supervisor_comms
    ):
        task = PythonOperator(task_id="test", python_callable=lambda: "success")
        ti = create_runtime_ti(task=task)
        if team_name:
            ti._ti_context_from_server.dag_run.team_name = team_name

        with mock.patch("airflow.sdk._shared.observability.metrics.stats._get_backend") as mock_get_backend:
            backend = mock.MagicMock(spec=StatsLogger)
            mock_get_backend.return_value = backend
            run(ti, context=ti.get_template_context(), log=mock.MagicMock())

            expected = {"dag_id": ti.dag_id, "task_id": ti.task_id, **expected_tags_extra}
            backend.incr.assert_any_call("ti.start", tags=expected)

    @pytest.mark.parametrize(
        ("task_callable", "operator_metric", "ti_metric"),
        [
            pytest.param(lambda: "success", "operator_successes", "ti_successes", id="success"),
            pytest.param(lambda: 1 / 0, "operator_failures", "ti_failures", id="failure"),
        ],
    )
    def test_operator_metrics_respect_team_name(
        self, task_callable, operator_metric, ti_metric, create_runtime_ti, mock_supervisor_comms
    ):
        task = PythonOperator(task_id="test", python_callable=task_callable)
        ti = create_runtime_ti(task=task)
        ti._ti_context_from_server.dag_run.team_name = "team_a"

        with mock.patch("airflow.sdk._shared.observability.metrics.stats._get_backend") as mock_get_backend:
            backend = mock.MagicMock(spec=StatsLogger)
            mock_get_backend.return_value = backend
            run(ti, context=ti.get_template_context(), log=mock.MagicMock())

            stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id, "team_name": "team_a"}
            backend.incr.assert_any_call(
                operator_metric,
                tags={**stats_tags, "operator_name": "PythonOperator"},
            )
            backend.incr.assert_any_call(ti_metric, tags=stats_tags)


class TestDetailSpan:
    """Tests for the detail_span decorator / context manager."""

    @pytest.fixture(autouse=True)
    def _sampled_carrier_provider(self):
        """Make new_dagrun_trace_carrier produce a SAMPLED carrier.

        new_dagrun_trace_carrier consults the global tracer provider's sampler to
        decide the carrier's SAMPLED flag. In the test process the global provider
        is a no-op ProxyTracerProvider (no sampler) -> unsampled carrier, which
        would make the parent span (and its detail children) non-recording. Patch
        the lookup to a real SDK provider whose default sampler
        (parentbased_always_on) samples the root, mirroring "otel on" in production.
        """
        provider = TracerProvider()
        with mock.patch(
            "airflow._shared.observability.traces.trace.get_tracer_provider",
            return_value=provider,
        ):
            yield

    def test_level_1_no_child_span_as_context_manager(self):
        """At detail level 1, entering detail_span should not create a real recorded span."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=1)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                with detail_span("child") as span:
                    assert span is trace.INVALID_SPAN

        # Only the "parent" span should be recorded; no "child".
        names = [s.name for s in exporter.get_finished_spans()]
        assert "child" not in names

    def test_level_2_creates_child_span_as_context_manager(self):
        """At detail level 2, detail_span should create a real recorded child span."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=2)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                with detail_span("child"):
                    pass

        names = [s.name for s in exporter.get_finished_spans()]
        assert "child" in names

    def test_decorator_at_level_1_does_not_create_span(self):
        """@detail_span at level 1 should not produce a recorded span."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=1)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        @detail_span("decorated")
        def my_func():
            return 42

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                result = my_func()

        assert result == 42
        names = [s.name for s in exporter.get_finished_spans()]
        assert "decorated" not in names

    def test_decorator_at_level_2_creates_span_and_preserves_return_value(self):
        """@detail_span at level 2 creates a span and the wrapped function's return value is preserved."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=2)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        @detail_span("decorated")
        def my_func(x):
            return x * 2

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                result = my_func(7)

        assert result == 14
        names = [s.name for s in exporter.get_finished_spans()]
        assert "decorated" in names

    def test_exception_in_context_manager_propagates(self):
        """Exceptions inside `with detail_span(...)` propagate normally."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=2)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                with pytest.raises(ValueError, match="boom"):
                    with detail_span("child"):
                        raise ValueError("boom")


class TestRunExecuteCallable:
    """Tests for ``_run_execute_callable``.

    It runs the task's execute callable inside an isolated contextvars copy (with
    the ExecutorSafeguard tracker set), applies the execution timeout when one is
    configured, and wraps the call in a ``task.execute`` detail span.
    """

    @pytest.fixture(autouse=True)
    def _sampled_carrier_provider(self):
        """Make new_dagrun_trace_carrier produce a SAMPLED carrier (see TestDetailSpan)."""
        provider = TracerProvider()
        with mock.patch(
            "airflow._shared.observability.traces.trace.get_tracer_provider",
            return_value=provider,
        ):
            yield

    @staticmethod
    def _make_task(execution_timeout=None):
        task = mock.MagicMock(spec=BaseOperator)
        task.execution_timeout = execution_timeout
        return task

    def test_runs_in_isolated_context_with_safeguard_tracker_set(self):
        """The callable runs in an internal context copy that has the safeguard tracker set and does not leak."""
        var = contextvars.ContextVar("marker")
        var.set("outer")
        task = self._make_task()
        seen = {}

        def execute(context):
            var.set("inner")
            seen["tracker"] = ExecutorSafeguard.tracker.get(None)
            return context["value"] * 2

        result = _run_execute_callable(context={"value": 21}, execute=execute, task=task)

        assert result == 42
        # The safeguard tracker is set to the task inside the copy used to run execute.
        assert seen["tracker"] is task
        # The mutation happened inside the copy, so it does not leak to the caller's context.
        assert var.get() == "outer"
        # The .set was confined to the copy, so the tracker never leaked to the caller's context.
        assert ExecutorSafeguard.tracker.get(None) is not task
        task.on_kill.assert_not_called()

    def test_applies_execution_timeout(self):
        """When a timeout is set and the callable overruns, AirflowTaskTimeout is raised and on_kill is called."""
        task = self._make_task(execution_timeout=timedelta(milliseconds=10))

        def execute(context):
            time.sleep(2)

        with pytest.raises(AirflowTaskTimeout):
            _run_execute_callable(context={}, execute=execute, task=task)

        task.on_kill.assert_called_once()

    def test_fast_fails_when_timeout_already_elapsed(self):
        """A non-positive timeout fast-fails before running the callable and still calls on_kill."""
        task = self._make_task(execution_timeout=timedelta(seconds=-1))
        execute = mock.MagicMock()

        with pytest.raises(AirflowTaskTimeout):
            _run_execute_callable(context={}, execute=execute, task=task)

        execute.assert_not_called()
        task.on_kill.assert_called_once()

    def test_emits_task_execute_span_at_detail_level_2(self):
        """At detail level 2, running the callable produces a recorded ``task.execute`` span."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=2)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        task = self._make_task()

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task)

        assert result == "ok"
        names = [s.name for s in exporter.get_finished_spans()]
        assert "task.execute" in names

    def test_operator_child_spans_nest_under_task_execute(self):
        """Spans the operator emits during execute nest under ``task.execute``, not its caller.

        The contextvars snapshot is taken inside ``_run_execute_callable`` after the
        ``task.execute`` span is current, so a span started during execute parents to
        ``task.execute`` rather than to the surrounding span.
        """
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=2)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        task = self._make_task()

        def execute(context):
            with t.start_as_current_span("operator_child"):
                return "ok"

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                result = _run_execute_callable(context={}, execute=execute, task=task)

        assert result == "ok"
        spans = {s.name: s for s in exporter.get_finished_spans()}
        assert spans["operator_child"].parent.span_id == spans["task.execute"].context.span_id

    def test_no_task_execute_span_at_detail_level_1(self):
        """At detail level 1, no ``task.execute`` span is recorded but the callable still runs."""
        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        t = provider.get_tracer("test")
        carrier = new_dagrun_trace_carrier(task_span_detail_level=1)
        parent_ctx = TraceContextTextMapPropagator().extract(carrier)

        task = self._make_task()

        with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t):
            with t.start_as_current_span("parent", context=parent_ctx):
                result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task)

        assert result == "ok"
        names = [s.name for s in exporter.get_finished_spans()]
        assert "task.execute" not in names


def test_dag_add_result(create_runtime_ti, mock_supervisor_comms):
    with DAG(dag_id="test_dag_add_result") as dag:
        task = PythonOperator(task_id="t", python_callable=lambda: 123)
        dag.add_result(task.output)

    ti = create_runtime_ti(task=task)
    run(ti, context=ti.get_template_context(), log=mock.MagicMock())

    mock_supervisor_comms.send.assert_any_call(
        SetXCom(
            key="return_value",
            value=123,
            dag_id="test_dag_add_result",
            run_id="test_run",
            task_id="t",
            map_index=-1,
            dag_result=True,
        )
    )


class TestTaskInstanceStateOperations:
    """Tests to verify that tasks can perform state operations (task / asset) via the supervisor."""

    def test_task_can_set_and_get_state(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        class MyOperator(BaseOperator):
            def execute(self, context):
                ts = context["task_state_store"]
                ts.set("job_id", "spark_app_001")
                return ts.get("job_id")

        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)
        frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(frozen_dt, tick=False)

        with conf_vars({("state_store", "default_retention_days"): "30"}):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetTaskStateStore(
                ti_id=runtime_ti.id,
                key="job_id",
                value="spark_app_001",
                expires_at=frozen_dt + timedelta(days=30),
            )
        )
        mock_supervisor_comms.send.assert_any_call(GetTaskStateStore(ti_id=runtime_ti.id, key="job_id"))

    def test_task_state_get_returns_default_when_key_missing(self, create_runtime_ti, mock_supervisor_comms):
        captured = {}

        class MyOperator(BaseOperator):
            def execute(self, context):
                captured["result"] = context["task_state_store"].get(
                    "watermark", default="2026-01-01T00:00:00+00:00"
                )

        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.TASK_STORE_NOT_FOUND, detail={"key": "watermark"}
        )
        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        assert captured["result"] == "2026-01-01T00:00:00+00:00"

    def test_task_state_set_sends_typed_values(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        """set() accepts any JsonValue — dict, int, list — not just strings."""

        class MyOperator(BaseOperator):
            def execute(self, context):
                ts = context["task_state_store"]
                ts.set("retry_count", 3)
                ts.set("poll_result", {"status": "succeeded", "rows": 1234})
                ts.set("checkpoints", [1, 2, 3])

        frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(frozen_dt, tick=False)
        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)

        with conf_vars({("state_store", "default_retention_days"): "30"}):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        expires_at = frozen_dt + timedelta(days=30)
        mock_supervisor_comms.send.assert_any_call(
            SetTaskStateStore(ti_id=runtime_ti.id, key="retry_count", value=3, expires_at=expires_at)
        )
        mock_supervisor_comms.send.assert_any_call(
            SetTaskStateStore(
                ti_id=runtime_ti.id,
                key="poll_result",
                value={"status": "succeeded", "rows": 1234},
                expires_at=expires_at,
            )
        )
        mock_supervisor_comms.send.assert_any_call(
            SetTaskStateStore(ti_id=runtime_ti.id, key="checkpoints", value=[1, 2, 3], expires_at=expires_at)
        )

    def test_task_can_set_state_with_retention(self, create_runtime_ti, mock_supervisor_comms, time_machine):
        class MyOperator(BaseOperator):
            def execute(self, context):
                context["task_state_store"].set("job_id", "spark_app_001", retention=timedelta(days=7))

        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)
        frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(frozen_dt, tick=False)

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetTaskStateStore(
                ti_id=runtime_ti.id,
                key="job_id",
                value="spark_app_001",
                expires_at=frozen_dt + timedelta(days=7),
            )
        )

    def test_task_can_delete_state(self, create_runtime_ti, mock_supervisor_comms):
        class MyOperator(BaseOperator):
            def execute(self, context):
                context["task_state_store"].delete("job_id")

        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(DeleteTaskStateStore(ti_id=runtime_ti.id, key="job_id"))

    def test_task_can_clear_state(self, create_runtime_ti, mock_supervisor_comms):
        class MyOperator(BaseOperator):
            def execute(self, context):
                context["task_state_store"].clear()

        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())
        mock_supervisor_comms.send.assert_any_call(ClearTaskStateStore(ti_id=runtime_ti.id))

    @staticmethod
    def _watcher_side_effect(msg=None, *args, **kwargs):
        actual = msg or (args[0] if args else None)
        if isinstance(actual, ValidateInletsAndOutlets):
            return InactiveAssetsResult(inactive_assets=[])
        if isinstance(actual, GetAssetByUri):
            # GetAssetByUri has no .name field. Mirroring AssetModel behaviour:
            # when only uri is provided, name defaults to uri.
            return AssetResult(name=actual.uri, uri=actual.uri, group="asset")
        return OKResponse(ok=True)

    def test_asset_state_store_get_and_set(self, create_runtime_ti, mock_supervisor_comms):
        watched = Asset(name="my_asset", uri="s3://bucket/data")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"].set("watermark", "2026-04-30")
                context["asset_state_store"].get("watermark")

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByName(name="my_asset", key="watermark", value="2026-04-30")
        )
        mock_supervisor_comms.send.assert_any_call(GetAssetStateStoreByName(name="my_asset", key="watermark"))

    def test_asset_state_store_get_returns_default_when_key_missing(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        watched = Asset(name="my_asset", uri="s3://bucket/data")
        captured = {}

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                captured["result"] = context["asset_state_store"].get(
                    "watermark", default="2026-01-01T00:00:00+00:00"
                )

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = lambda msg: (
            ErrorResponse(error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "watermark"})
            if isinstance(msg, GetAssetStateStoreByName)
            else TestTaskInstanceStateOperations._watcher_side_effect(msg)
        )
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        assert captured["result"] == "2026-01-01T00:00:00+00:00"

    def test_asset_state_store_delete(self, create_runtime_ti, mock_supervisor_comms):
        watched = Asset(name="my_asset", uri="s3://bucket/data")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"].delete("watermark")

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            DeleteAssetStateStoreByName(name="my_asset", key="watermark")
        )

    def test_asset_state_store_clear(self, create_runtime_ti, mock_supervisor_comms):
        watched = Asset(name="my_asset", uri="s3://bucket/data")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"].clear()

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(ClearAssetStateStoreByName(name="my_asset"))

    def test_asset_state_store_uri_ref_inlet(self, create_runtime_ti, mock_supervisor_comms):
        watched = AssetUriRef(uri="s3://bucket/data")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"].set("watermark", "2026-04-30")
                context["asset_state_store"].get("watermark")

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByUri(uri="s3://bucket/data", key="watermark", value="2026-04-30")
        )
        mock_supervisor_comms.send.assert_any_call(
            GetAssetStateStoreByUri(uri="s3://bucket/data", key="watermark")
        )

    def test_asset_state_store_alias_as_inlet(self, create_runtime_ti, mock_supervisor_comms):
        alias = AssetAlias(name="my_alias")
        resolved = Asset(name="resolved_asset", uri="s3://bucket/resolved")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"][resolved].set("watermark", "2026-05-01")

        def side_effect(msg):
            if isinstance(msg, GetAssetsByAlias):
                return AssetsByAliasResult(
                    assets=[AssetResult(name=resolved.name, uri=resolved.uri, group="asset")]
                )
            return TestTaskInstanceStateOperations._watcher_side_effect(msg)

        task = WatcherOperator(task_id="t", inlets=[alias])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByName(name="resolved_asset", key="watermark", value="2026-05-01")
        )

    def test_asset_state_store_alias_inlet_no_resolved_assets(self, create_runtime_ti, mock_supervisor_comms):
        alias = AssetAlias(name="empty_alias")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                # asset_state is in context but it is empty because alias resolved to nothing
                assert "asset_state_store" in context

        def side_effect(msg):
            if isinstance(msg, GetAssetsByAlias):
                return AssetsByAliasResult(assets=[])
            return TestTaskInstanceStateOperations._watcher_side_effect(msg)

        task = WatcherOperator(task_id="t", inlets=[alias])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

    def test_asset_state_store_keyed_access_single_inlet(self, create_runtime_ti, mock_supervisor_comms):
        watched = Asset(name="my_asset", uri="s3://bucket/data")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                # accessing via asset name key
                context["asset_state_store"][watched].set("watermark", "2026-05-01")

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByName(name="my_asset", key="watermark", value="2026-05-01")
        )

    def test_asset_state_store_multi_inlet(self, create_runtime_ti, mock_supervisor_comms):
        asset_a = Asset(name="asset_a", uri="s3://bucket/a")
        asset_b = Asset(name="asset_b", uri="s3://bucket/b")

        class MultiInletOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"][asset_a].set("watermark_a", "2026-05-01")
                context["asset_state_store"][asset_b].set("watermark_b", "2026-05-02")

        task = MultiInletOperator(task_id="t", inlets=[asset_a, asset_b])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByName(name="asset_a", key="watermark_a", value="2026-05-01")
        )
        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByName(name="asset_b", key="watermark_b", value="2026-05-02")
        )

    def test_asset_state_store_set_sends_reference_via_custom_backend(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """When a worker backend is configured, asset state set() sends a reference, not the actual value."""
        watched = Asset(name="my_asset", uri="s3://bucket/data")

        class WatcherOperator(BaseOperator):
            def execute(self, context):
                context["asset_state_store"].set("watermark", "2026-05-01")

        task = WatcherOperator(task_id="t", inlets=[watched])
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        mock_backend = mock.MagicMock()
        mock_backend.serialize_asset_state_store_to_ref.return_value = "mem://my_asset/watermark"

        with mock.patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=mock_backend
        ):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_backend.serialize_asset_state_store_to_ref.assert_called_once_with(
            value="2026-05-01", key="watermark", scope=AssetScope(name="my_asset", uri=None)
        )
        mock_supervisor_comms.send.assert_any_call(
            SetAssetStateStoreByName(
                name="my_asset",
                key="watermark",
                value=_wrap_external_ref("mem://my_asset/watermark"),
            )
        )

    def test_task_state_set_sends_reference_via_custom_backend(
        self, create_runtime_ti, mock_supervisor_comms, time_machine
    ):
        """When a worker backend is configured, task state set() sends a reference, not the actual value."""

        class MyOperator(BaseOperator):
            def execute(self, context):
                context["task_state_store"].set("job_id", "app_001")

        frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(frozen_dt, tick=False)
        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)
        mock_supervisor_comms.send.side_effect = TestTaskInstanceStateOperations._watcher_side_effect

        mock_backend = mock.MagicMock()
        scope = TaskScope(
            dag_id=runtime_ti.dag_id,
            run_id=runtime_ti.run_id,
            task_id=runtime_ti.task_id,
            map_index=runtime_ti.map_index,
        )
        ref = f"mem://{scope.dag_id}/{scope.run_id}/{scope.task_id}/{scope.map_index}/job_id"
        mock_backend.serialize_task_state_store_to_ref.return_value = ref

        with mock.patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=mock_backend
        ):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_backend.serialize_task_state_store_to_ref.assert_called_once_with(
            value="app_001", key="job_id", scope=scope
        )
        mock_supervisor_comms.send.assert_any_call(
            SetTaskStateStore(
                ti_id=runtime_ti.id,
                key="job_id",
                value=_wrap_external_ref(ref),
                expires_at=frozen_dt + timedelta(days=30),
            )
        )

    @conf_vars({("state_store", "clear_on_success"): "True"})
    def test_clear_on_success_skips_backend_clear_when_no_custom_backend(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """clear_on_success does not call backend.clear() when no custom backend is configured."""

        class MyOperator(BaseOperator):
            def execute(self, context):
                pass

        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)

        # no need to patch because default of _get_worker_state_store_backend() is None
        run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        sent_types = [
            type(call.kwargs.get("msg") or (call.args[0] if call.args else None))
            for call in mock_supervisor_comms.send.call_args_list
        ]
        assert ClearTaskStateStore not in sent_types

    @conf_vars({("state_store", "clear_on_success"): "True"})
    def test_clear_on_success_clears_backend_without_comms_roundtrip(
        self, create_runtime_ti, mock_supervisor_comms
    ):
        """clear_on_success calls backend.clear() directly without sending ClearTaskStateStore comms."""
        mock_backend = mock.MagicMock()

        class MyOperator(BaseOperator):
            def execute(self, context):
                pass

        task = MyOperator(task_id="t")
        runtime_ti = create_runtime_ti(task=task)

        with (
            mock.patch(
                "airflow.sdk.execution_time.context._get_worker_state_store_backend",
                return_value=mock_backend,
            ),
            mock.patch(
                "airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
                return_value=mock_backend,
            ),
        ):
            run(runtime_ti, context=runtime_ti.get_template_context(), log=mock.MagicMock())

        mock_backend.clear.assert_called_once()
        sent_types = [
            type(call.kwargs.get("msg") or (call.args[0] if call.args else None))
            for call in mock_supervisor_comms.send.call_args_list
        ]
        assert ClearTaskStateStore not in sent_types


class _WalkerModelA(BaseModel):
    a: int


class _WalkerModelB(BaseModel):
    b: int


class _WalkerOperator(BaseOperator):
    """Operator that declares a Pydantic ``output_type`` for deserialization."""

    deserialization_allowed_class_fields = ("output_type",)

    def __init__(self, *, output_type=str, value=None, **kwargs):
        super().__init__(**kwargs)
        self.output_type = output_type
        self.value = value

    def execute(self, context):
        return None


class TestRegisterDeserializationAllowedClasses:
    """The worker-side walk registers operator-declared XCom classes for the whole DAG.

    ``allow_class`` is patched to a spy so these assert the walk *extracts* the right
    classes from both real and mapped operators; ``allow_class``'s own import
    validation is covered by the serde tests.
    """

    def test_registers_real_and_mapped_operators(self):

        with DAG("walker_dag") as dag:
            # Non-mapped producer: output_type is a plain attribute.
            _WalkerOperator(task_id="real", output_type=_WalkerModelA)
            # Mapped producer: output_type lives in partial_kwargs and __init__ never
            # runs at parse -- the case the old __init__ registration could not reach.
            _WalkerOperator.partial(task_id="mapped", output_type=_WalkerModelB).expand(value=[1, 2])

        registered: list[type] = []
        with patch("airflow.sdk.execution_time.task_runner.allow_class", side_effect=registered.append):
            _register_deserialization_allowed_classes(dag, structlog.get_logger())

        assert _WalkerModelA in registered, "real operator output_type not registered"
        assert _WalkerModelB in registered, "mapped operator output_type not registered"

    def test_default_operator_registers_nothing(self):

        with DAG("walker_dag_plain") as dag:
            BaseOperator(task_id="plain")

        registered: list[type] = []
        with patch("airflow.sdk.execution_time.task_runner.allow_class", side_effect=registered.append):
            _register_deserialization_allowed_classes(dag, structlog.get_logger())
        assert registered == []

    def test_bad_declaration_is_skipped_not_fatal(self):
        """A class that fails ``allow_class`` validation is logged and skipped, not raised."""

        with DAG("walker_dag_bad") as dag:
            _WalkerOperator(task_id="real", output_type=_WalkerModelA)

        with patch("airflow.sdk.execution_time.task_runner.allow_class", side_effect=ValueError("nope")):
            # Must not raise -- the walk swallows per-class registration errors.
            _register_deserialization_allowed_classes(dag, structlog.get_logger())
