# 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

from datetime import datetime
from types import SimpleNamespace
from typing import TYPE_CHECKING
from unittest import mock
from uuid import UUID, uuid4

import pytest
import uuid6
from fastapi import Request
from opentelemetry import trace as otel_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 import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from sqlalchemy import select, update
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session

from airflow._shared.observability.traces import OverrideableRandomIdGenerator
from airflow._shared.state import TaskScope
from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
from airflow.api_fastapi.execution_api.app import lifespan
from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span
from airflow.api_fastapi.execution_api.security import require_auth
from airflow.exceptions import AirflowSkipException
from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger
from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel
from airflow.models.dag import DagModel
from airflow.models.log import Log
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance
from airflow.models.taskinstancehistory import TaskInstanceHistory
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group
from airflow.state.metastore import MetastoreBackend
from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState

from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
    clear_db_assets,
    clear_db_dags,
    clear_db_logs,
    clear_db_runs,
    clear_db_serialized_dags,
    clear_rendered_ti_fields,
)
from unit.listeners import asset_listener

if TYPE_CHECKING:
    from airflow.sdk.api.client import Client

    from tests_common.pytest_plugin import DagMaker

pytestmark = pytest.mark.db_test


DEFAULT_START_DATE = timezone.parse("2024-10-31T11:00:00Z")
DEFAULT_END_DATE = timezone.parse("2024-10-31T12:00:00Z")
DEFAULT_RENDERED_MAP_INDEX = "test rendered map index"


def _where_column_keys(statement) -> set[str]:
    whereclause = getattr(statement, "whereclause", None)
    if whereclause is None:
        return set()

    keys: set[str] = set()
    stack = [whereclause]
    while stack:
        clause = stack.pop()
        left = getattr(clause, "left", None)
        if left is not None and getattr(left, "key", None) is not None:
            keys.add(left.key)
        stack.extend(clause.get_children())
    return keys


def _is_task_instance_update(statement) -> bool:
    return getattr(statement, "is_update", False) and statement.table.name == TaskInstance.__table__.name


def _is_select_for_update(statement) -> bool:
    return getattr(statement, "is_select", False) and "FOR UPDATE" in str(statement.compile()).upper()


def _create_asset_aliases(session, num: int = 2) -> None:
    asset_aliases = [
        AssetAliasModel(
            id=i,
            name=f"simple{i}",
            group="alias",
        )
        for i in range(1, 1 + num)
    ]
    session.add_all(asset_aliases)
    session.commit()


@pytest.fixture
def _use_real_jwt_bearer(exec_app):
    """Remove the mock require_auth override so the real JWT validation runs end-to-end."""
    exec_app.dependency_overrides.pop(require_auth, None)


@pytest.mark.usefixtures("_use_real_jwt_bearer")
def test_id_matches_sub_claim(client, session, create_task_instance):
    """Test that scope validation (ti:self) is enforced at the router level."""
    ti = create_task_instance(
        task_id="test_ti_run_state_conflict_if_not_queued",
        state="queued",
    )
    session.commit()

    validator = mock.AsyncMock(spec=JWTValidator)
    validator.avalidated_claims.return_value = {
        "sub": str(ti.id),
        "scope": "execution",
        "exp": 9999999999,
        "iat": 1000000000,
        "nbf": 1000000000,
    }
    lifespan.registry.register_value(JWTValidator, validator)

    payload = {
        "state": "running",
        "hostname": "random-hostname",
        "unixname": "random-unixname",
        "pid": 100,
        "start_date": "2024-10-31T12:00:00Z",
    }

    resp = client.patch("/execution/task-instances/9c230b40-da03-451d-8bd7-be30471be383/run", json=payload)
    assert resp.status_code == 403
    validator.avalidated_claims.reset_mock()

    resp = client.patch(f"/execution/task-instances/{ti.id}/run", json=payload)
    assert resp.status_code == 200, resp.json()
    validator.avalidated_claims.assert_awaited()


class TestTIRunState:
    def setup_method(self):
        clear_db_logs()
        clear_db_runs()
        clear_db_serialized_dags()
        clear_db_dags()
        clear_db_assets()

    def teardown_method(self):
        clear_db_logs()
        clear_db_runs()
        clear_db_serialized_dags()
        clear_db_dags()
        clear_db_assets()

    def test_ti_run_context_exposes_consumed_event_partition_key(self, client, session, create_task_instance):
        """The partition key of each consumed asset event is returned in the run context."""
        ti = create_task_instance(
            task_id="test_consumed_event_partition_key",
            state=State.QUEUED,
            session=session,
        )
        asset = AssetModel(name="upstream", uri="s3://bucket/upstream", group="asset", extra={})
        session.add_all([asset, AssetActive.for_asset(asset)])
        session.flush()
        ti.dag_run.consumed_asset_events.append(
            AssetEvent(asset_id=asset.id, source_dag_id="src", source_run_id="r1", partition_key="2024-01-15")
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "h",
                "unixname": "u",
                "pid": 1,
                "start_date": "2024-09-30T12:00:00Z",
            },
        )

        assert response.status_code == 200
        events = response.json()["dag_run"]["consumed_asset_events"]
        assert [e["partition_key"] for e in events] == ["2024-01-15"]

    @pytest.mark.parametrize(
        ("max_tries", "should_retry"),
        [
            pytest.param(0, False, id="max_retries=0"),
            pytest.param(3, True, id="should_retry"),
        ],
    )
    def test_ti_run_state_to_running(
        self,
        client,
        session,
        create_task_instance,
        time_machine,
        max_tries,
        should_retry,
    ):
        """
        Test that the Task Instance state is updated to running when the Task Instance is in a state where it can be
        marked as running.
        """
        instant_str = "2024-09-30T12:00:00Z"
        instant = timezone.parse(instant_str)
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_ti_run_state_to_running",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            start_date=instant,
            dag_id=str(uuid4()),
        )
        ti.max_tries = max_tries
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": instant_str,
            },
        )

        assert response.status_code == 200
        result = response.json()
        assert result == {
            "dag_run": {
                "dag_id": ti.dag_id,
                "run_id": "test",
                "clear_number": 0,
                "logical_date": instant_str,
                "data_interval_start": instant.subtract(days=1).to_iso8601_string(),
                "data_interval_end": instant_str,
                "run_after": instant_str,
                "start_date": instant_str,
                "state": "running",
                "end_date": None,
                "run_type": "manual",
                "conf": {},
                "triggering_user_name": None,
                "consumed_asset_events": [],
                "partition_key": None,
                "partition_date": None,
                "note": None,
                "team_name": None,
            },
            "task_reschedule_count": 0,
            "max_tries": max_tries,
            "should_retry": should_retry,
            "variables": [],
            "connections": [],
            "xcom_keys_to_clear": [],
        }
        # upstream_map_indexes is now computed by Task SDK, not returned by the server in HEAD version
        assert "upstream_map_indexes" not in result
        # execution-scoped tokens do not trigger a token swap
        assert "Refreshed-API-Token" not in response.headers

        # Refresh the Task Instance from the database so that we can check the updated values
        session.refresh(ti)
        assert ti.state == State.RUNNING
        assert ti.hostname == "random-hostname"
        assert ti.unixname == "random-unixname"
        assert ti.pid == 100

        response1 = response.json()

        # Test that if we make a second request (simulating a network glitch so the client issues a retry)
        # that it is accepted and we get the same info back

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": instant_str,
            },
        )
        assert response.status_code == 200
        assert response.json() == response1

        # But that for a different pid on the same host (etc) it fails
        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 101,
                "start_date": instant_str,
            },
        )
        assert response.status_code == 409

    def test_ti_run_returns_execution_token(
        self, client, exec_app, session, create_task_instance, time_machine
    ):
        """PATCH /run with a workload token should swap to an execution-scoped token."""
        instant = timezone.parse("2024-10-31T12:00:00Z")
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_exec_token",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            start_date=instant,
            dag_id=str(uuid4()),
        )
        session.commit()

        mock_gen = mock.MagicMock(spec=JWTGenerator)
        mock_gen.generate.return_value = "mock-execution-token"
        lifespan.registry.register_value(JWTGenerator, mock_gen)

        async def workload_token(request: Request) -> TIToken:
            ti_id = UUID(request.path_params.get("task_instance_id", "00000000-0000-0000-0000-000000000000"))
            return TIToken(id=ti_id, claims=TIClaims(scope="workload"))

        exec_app.dependency_overrides[require_auth] = workload_token

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "test-host",
                "unixname": "test-user",
                "pid": 100,
                "start_date": "2024-10-31T12:00:00Z",
            },
        )

        exec_app.dependency_overrides.pop(require_auth, None)

        assert response.status_code == 200
        assert "Refreshed-API-Token" in response.headers
        assert response.headers["Refreshed-API-Token"] == "mock-execution-token"
        mock_gen.generate.assert_called_once()
        extras = mock_gen.generate.call_args.kwargs["extras"]
        assert extras["scope"] == "execution"
        assert extras["sub"] == str(ti.id)

    def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker):
        """Test that dynamic task mapping works correctly with parse-time values."""
        with dag_maker("test_dynamic_task_mapping_with_parse_time_value", serialized=True):

            @task_group
            def task_group_1(arg1):
                @task
                def group1_task_1(arg1):
                    return {"a": arg1}

                @task
                def group1_task_2(arg2):
                    return arg2

                group1_task_2(group1_task_1(arg1))

            @task
            def task2():
                return None

            task_group_1.expand(arg1=[0, 1]) >> task2()

        dr = dag_maker.create_dagrun()
        for ti in dr.get_task_instances():
            ti.set_state(State.QUEUED)
        dag_maker.session.flush()

        for ti in dr.get_task_instances():
            response = client.patch(
                f"/execution/task-instances/{ti.id}/run",
                json={
                    "state": "running",
                    "hostname": "random-hostname",
                    "unixname": "random-unixname",
                    "pid": 100,
                    "start_date": "2024-09-30T12:00:00Z",
                },
            )

            assert response.status_code == 200

    def test_nested_mapped_task_group(self, client, dag_maker):
        """Test that nested mapped task groups work correctly."""
        with dag_maker("test_nested_mapped_tg", serialized=True):

            @task
            def alter_input(inp: str) -> str:
                return f"{inp}_Altered"

            @task
            def print_task(orig_input: str, altered_input: str) -> str:
                return f"orig:{orig_input},altered:{altered_input}"

            @task_group
            def inner_task_group(orig_input: str) -> None:
                altered_input = alter_input(orig_input)
                print_task(orig_input, altered_input)

            @task_group
            def expandable_task_group(param: str) -> None:
                inner_task_group(param)

            expandable_task_group.expand(param=["One", "Two", "Three"])

        dr = dag_maker.create_dagrun()

        # Set all alter_input tasks to success so print_task can run
        for ti in dr.get_task_instances():
            if "alter_input" in ti.task_id and ti.map_index >= 0:
                ti.state = State.SUCCESS
            elif "print_task" in ti.task_id and ti.map_index >= 0:
                ti.set_state(State.QUEUED)
        dag_maker.session.flush()

        # Get only the expanded print_task instances (not the template)
        print_task_tis = [
            ti for ti in dr.get_task_instances() if "print_task" in ti.task_id and ti.map_index >= 0
        ]

        for ti in print_task_tis:
            response = client.patch(
                f"/execution/task-instances/{ti.id}/run",
                json={
                    "state": "running",
                    "hostname": "random-hostname",
                    "unixname": "random-unixname",
                    "pid": 100,
                    "start_date": "2024-09-30T12:00:00Z",
                },
            )

            assert response.status_code == 200

    def test_dynamic_task_mapping_with_xcom(self, client: Client, dag_maker: DagMaker, session: Session):
        """Test that dynamic task mapping works correctly with XCom values."""
        from airflow.models.taskmap import TaskMap

        with dag_maker(session=session, serialized=True):

            @task
            def task_1():
                return [0, 1]

            @task_group
            def tg(x, y):
                @task
                def task_2():
                    pass

                task_2()

            @task
            def task_3():
                pass

            tg.expand(x=task_1(), y=[1, 2, 3]) >> task_3()

        dr = dag_maker.create_dagrun()

        decision = dr.task_instance_scheduling_decisions(session=session)

        # Simulate task_1 execution to produce TaskMap.
        (ti_1,) = decision.schedulable_tis
        ti_1.state = TaskInstanceState.SUCCESS
        session.add(TaskMap.from_task_instance_xcom(ti_1, [0, 1]))
        session.flush()

        # Now task_2 in mapped tagk group is expanded.
        decision = dr.task_instance_scheduling_decisions(session=session)
        for ti in decision.schedulable_tis:
            ti.state = TaskInstanceState.SUCCESS
        session.flush()

        decision = dr.task_instance_scheduling_decisions(session=session)
        (task_3_ti,) = decision.schedulable_tis
        task_3_ti.set_state(State.QUEUED)

        response = client.patch(
            f"/execution/task-instances/{task_3_ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": "2024-09-30T12:00:00Z",
            },
        )
        assert response.status_code == 200

    def test_dynamic_task_mapping_with_all_success_trigger_rule(self, dag_maker: DagMaker, session: Session):
        """Test that with ALL_SUCCESS trigger rule and skipped upstream, downstream should not run."""

        with dag_maker(session=session, serialized=True):

            @task
            def task_1():
                raise AirflowSkipException()

            @task_group
            def tg(x):
                @task
                def task_2():
                    raise AirflowSkipException()

                task_2()

            @task(trigger_rule=TriggerRule.ALL_SUCCESS)
            def task_3():
                pass

            @task
            def task_4():
                pass

            tg.expand(x=task_1()) >> [task_3(), task_4()]

        dr = dag_maker.create_dagrun()

        decision = dr.task_instance_scheduling_decisions(session=session)

        # Simulate task_1 skipped
        (ti_1,) = decision.schedulable_tis
        ti_1.state = TaskInstanceState.SKIPPED
        session.flush()

        # Now task_2 in mapped task group is not expanded and also skipped.
        decision = dr.task_instance_scheduling_decisions(session=session)
        for ti in decision.schedulable_tis:
            ti.state = TaskInstanceState.SKIPPED
        session.flush()

        decision = dr.task_instance_scheduling_decisions(session=session)
        assert decision.schedulable_tis == []

    @pytest.mark.parametrize(
        "trigger_rule",
        [
            TriggerRule.ALL_DONE,
            TriggerRule.ALL_DONE_SETUP_SUCCESS,
            TriggerRule.NONE_FAILED,
            TriggerRule.ALL_SKIPPED,
        ],
    )
    def test_dynamic_task_mapping_with_non_all_success_trigger_rule(
        self, client: Client, dag_maker: DagMaker, session: Session, trigger_rule: TriggerRule
    ):
        """Test that with non-ALL_SUCCESS trigger rule, downstream task should still run."""

        with dag_maker(session=session, serialized=True):

            @task
            def task_1():
                raise AirflowSkipException()

            @task_group
            def tg(x):
                @task
                def task_2():
                    raise AirflowSkipException()

                task_2()

            @task(trigger_rule=trigger_rule)
            def task_3():
                pass

            @task
            def task_4():
                pass

            tg.expand(x=task_1()) >> [task_3(), task_4()]

        dr = dag_maker.create_dagrun()

        decision = dr.task_instance_scheduling_decisions(session=session)

        # Simulate task_1 skipped
        (ti_1,) = decision.schedulable_tis
        ti_1.state = TaskInstanceState.SKIPPED
        session.flush()

        # Now task_2 in mapped tagk group is not expanded and also skipped..
        decision = dr.task_instance_scheduling_decisions(session=session)
        for ti in decision.schedulable_tis:
            ti.state = TaskInstanceState.SKIPPED
        session.flush()

        decision = dr.task_instance_scheduling_decisions(session=session)
        # only task_3 is schedulable
        (task_3_ti,) = decision.schedulable_tis
        assert task_3_ti.task_id == "task_3"
        task_3_ti.set_state(State.QUEUED)

        response = client.patch(
            f"/execution/task-instances/{task_3_ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": "2024-09-30T12:00:00Z",
            },
        )
        assert response.status_code == 200

    def test_next_kwargs_still_encoded(self, client, session, create_task_instance, time_machine):
        instant_str = "2024-09-30T12:00:00Z"
        instant = timezone.parse(instant_str)
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_next_kwargs_still_encoded",
            state=State.QUEUED,
            session=session,
            start_date=instant,
            dag_id=str(uuid4()),
        )

        ti.next_method = "execute_complete"
        # explicitly use serde serialized value before assigning since we use JSON/JSONB now
        # that this value comes serde serialized from the worker
        expected_next_kwargs = {
            "moment": {
                "__classname__": "pendulum.datetime.DateTime",
                "__version__": 2,
                "__data__": {
                    "timestamp": 1727697600.0,
                    "tz": {
                        "__classname__": "builtins.tuple",
                        "__version__": 1,
                        "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                    },
                },
            }
        }
        ti.next_kwargs = expected_next_kwargs

        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": instant_str,
            },
        )

        assert response.status_code == 200
        assert response.json() == {
            "dag_run": mock.ANY,
            "task_reschedule_count": 0,
            "max_tries": 0,
            "should_retry": False,
            "variables": [],
            "connections": [],
            "xcom_keys_to_clear": [],
            "next_method": "execute_complete",
            "next_kwargs": expected_next_kwargs,
            "start_date": None,
        }

    @pytest.mark.parametrize("resume", [True, False])
    def test_next_kwargs_determines_start_date_update(self, client, session, create_task_instance, resume):
        dag_start_time_str = "2024-09-30T12:00:00Z"
        dag_start_time = timezone.parse(dag_start_time_str)
        orig_task_start_time = dag_start_time.add(seconds=5)

        ti = create_task_instance(
            task_id="test_next_kwargs_still_encoded",
            state=State.QUEUED,
            session=session,
            start_date=orig_task_start_time,
            dag_id=str(uuid4()),
        )

        ti.start_date = orig_task_start_time
        ti.next_method = "execute_complete"

        second_start_time = orig_task_start_time.add(seconds=30)
        second_start_time_str = second_start_time.isoformat()

        # explicitly serialize using serde before assigning since we use JSON/JSONB now
        # this value comes serde serialized from the worker
        if resume:
            # expected format is now in serde serialized format
            expected_next_kwargs = {
                "moment": {
                    "__classname__": "pendulum.datetime.DateTime",
                    "__version__": 2,
                    "__data__": {
                        "timestamp": 1727697635.0,
                        "tz": {
                            "__classname__": "builtins.tuple",
                            "__version__": 1,
                            "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                        },
                    },
                }
            }
            ti.next_kwargs = expected_next_kwargs
            expected_start_date = orig_task_start_time
        else:
            expected_start_date = second_start_time
            expected_next_kwargs = None

        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": second_start_time_str,
            },
        )
        session.commit()

        assert response.status_code == 200
        result = response.json()
        assert timezone.parse(result["start_date"]) == orig_task_start_time
        result.pop("start_date")
        assert result == {
            "dag_run": mock.ANY,
            "task_reschedule_count": 0,
            "max_tries": 0,
            "should_retry": False,
            "variables": [],
            "connections": [],
            "xcom_keys_to_clear": [],
            "next_method": "execute_complete",
            "next_kwargs": expected_next_kwargs,
        }
        session.expunge_all()
        ti = session.get(TaskInstance, ti.id)
        assert ti.start_date == expected_start_date

    def test_ti_run_resume_returns_original_start_date_in_context(
        self,
        client,
        session,
        create_task_instance,
    ):
        original_start_date = timezone.parse("2024-09-30T12:00:05Z")
        payload_start_date = timezone.parse("2024-09-30T12:00:35Z")

        ti = create_task_instance(
            task_id="test_ti_run_resume_returns_original_start_date_in_context",
            state=State.QUEUED,
            session=session,
            start_date=original_start_date,
            dag_id=str(uuid4()),
        )
        ti.start_date = original_start_date
        ti.next_method = "execute_complete"
        ti.next_kwargs = {
            "moment": {
                "__classname__": "pendulum.datetime.DateTime",
                "__version__": 2,
                "__data__": {
                    "timestamp": 1727697605.0,
                    "tz": {
                        "__classname__": "builtins.tuple",
                        "__version__": 1,
                        "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                    },
                },
            }
        }
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": payload_start_date.isoformat(),
            },
        )

        assert response.status_code == 200
        result = response.json()
        assert result["start_date"] is not None
        assert timezone.parse(result["start_date"]) == original_start_date

        session.expunge_all()
        ti = session.get(TaskInstance, ti.id)
        assert ti is not None
        assert ti.start_date == original_start_date

    @pytest.mark.parametrize(
        "initial_ti_state",
        [s for s in TaskInstanceState if s not in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING)],
    )
    def test_ti_run_state_conflict_if_not_queued(
        self, client, session, create_task_instance, initial_ti_state
    ):
        """
        Test that a 409 error is returned when the Task Instance is not in a state where it can be marked as
        running. In this case, the Task Instance is first in NONE state so it cannot be marked as running.
        """
        ti = create_task_instance(
            task_id="test_ti_run_state_conflict_if_not_queued",
            state=initial_ti_state,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": "2024-10-31T12:00:00Z",
            },
        )

        assert response.status_code == 409
        assert response.json() == {
            "detail": {
                "message": "TI was not in a state where it could be marked as running",
                "previous_state": initial_ti_state,
                "reason": "invalid_state",
            }
        }

        assert session.scalar(select(TaskInstance.state).where(TaskInstance.id == ti.id)) == initial_ti_state

    def test_xcom_not_cleared_for_deferral(self, client, session, create_task_instance, time_machine):
        """
        Test that the xcoms are not cleared when the Task Instance state is re-running after deferral.
        """
        instant_str = "2024-09-30T12:00:00Z"
        instant = timezone.parse(instant_str)
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_xcom_not_cleared_for_deferral",
            state=State.RUNNING,
            session=session,
            start_date=instant,
            dag_id=str(uuid4()),
        )
        session.commit()

        # Move this task to deferred
        payload = {
            "state": "deferred",
            "trigger_kwargs": {"key": "value", "moment": "2024-12-18T00:00:00Z"},
            "trigger_timeout": "P1D",  # 1 day
            "classpath": "my-classpath",
            "next_method": "execute_callback",
        }

        response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
        assert response.status_code == 204
        assert response.text == ""
        session.expire_all()

        # Deferred -> Queued so that we can run it again
        query = update(TaskInstance).where(TaskInstance.id == ti.id).values(state="queued")
        session.execute(query)
        session.commit()

        # Lets stage a xcom push
        ti.xcom_push(key="key", value="value")

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": instant_str,
            },
        )

        assert response.status_code == 200
        assert ti.xcom_pull(task_ids="test_xcom_not_cleared_for_deferral", key="key") == "value"

    def test_ti_run_with_triggering_user_name(
        self,
        client,
        session,
        dag_maker,
        time_machine,
    ):
        """
        Test that the triggering_user_name field is correctly returned when it has a non-None value.
        """
        instant_str = "2024-09-30T12:00:00Z"
        instant = timezone.parse(instant_str)
        time_machine.move_to(instant, tick=False)

        with dag_maker(dag_id=str(uuid4()), session=session):
            EmptyOperator(task_id="test_ti_run_with_triggering_user_name")

        # Create DagRun with triggering_user_name set to a specific value
        dr = dag_maker.create_dagrun(
            run_id="test",
            logical_date=instant,
            state=DagRunState.RUNNING,
            start_date=instant,
            triggering_user_name="test_user",
        )

        ti = dr.get_task_instance(task_id="test_ti_run_with_triggering_user_name")
        ti.set_state(State.QUEUED)
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "test-hostname",
                "unixname": "test-unixname",
                "pid": 12345,
                "start_date": instant_str,
            },
        )

        assert response.status_code == 200
        json_response = response.json()

        # Verify the dag_run is present
        assert "dag_run" in json_response
        dag_run = json_response["dag_run"]

        # The triggering_user_name field should be present with the correct value
        assert dag_run["triggering_user_name"] == "test_user"

        # Verify other expected fields are still present
        assert dag_run["dag_id"] == ti.dag_id
        assert dag_run["run_id"] == "test"
        assert dag_run["state"] == "running"

    @pytest.mark.parametrize(
        ("multi_team_enabled", "expect_team"),
        [
            pytest.param("False", False, id="multi-team-disabled"),
            pytest.param("True", True, id="multi-team-enabled"),
        ],
    )
    def test_ti_run_populates_team_name(
        self,
        client,
        session,
        dag_maker,
        time_machine,
        multi_team_enabled,
        expect_team,
    ):
        """``team_name`` is resolved from the DAG's bundle when multi-team is enabled."""
        from airflow.models.dagbundle import DagBundleModel
        from airflow.models.team import Team

        instant = timezone.parse("2024-09-30T12:00:00Z")
        time_machine.move_to(instant, tick=False)

        dag_id = str(uuid4())
        with dag_maker(dag_id=dag_id, session=session):
            EmptyOperator(task_id="task")
        dr = dag_maker.create_dagrun(
            run_id="test", logical_date=instant, state=DagRunState.RUNNING, start_date=instant
        )
        ti = dr.get_task_instance(task_id="task")
        ti.set_state(State.QUEUED)

        bundle_name = f"bundle-{dag_id}"
        team_name = f"team-{dag_id[:8]}"
        bundle = DagBundleModel(name=bundle_name)
        bundle.teams.append(Team(name=team_name))
        session.add(bundle)
        session.flush()
        session.execute(update(DagModel).where(DagModel.dag_id == dag_id).values(bundle_name=bundle_name))
        session.commit()

        with conf_vars({("core", "multi_team"): multi_team_enabled}):
            response = client.patch(
                f"/execution/task-instances/{ti.id}/run",
                json={
                    "state": "running",
                    "hostname": "h",
                    "unixname": "u",
                    "pid": 1,
                    "start_date": "2024-09-30T12:00:00Z",
                },
            )

        assert response.status_code == 200
        assert response.json()["dag_run"]["team_name"] == (team_name if expect_team else None)

    def test_ti_run_creates_audit_log(self, client, session, create_task_instance, time_machine):
        """Test that transitioning to RUNNING creates an audit log record."""
        instant_str = "2024-09-30T12:00:00Z"
        instant = timezone.parse(instant_str)
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_ti_run_creates_audit_log",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            start_date=instant,
            dag_id=str(uuid4()),
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "random-hostname",
                "unixname": "random-unixname",
                "pid": 100,
                "start_date": instant_str,
            },
        )

        assert response.status_code == 200

        logs = session.scalars(select(Log).where(Log.dag_id == ti.dag_id)).all()
        assert len(logs) == 1
        assert logs[0].event == TaskInstanceState.RUNNING.value
        assert logs[0].task_id == ti.task_id
        assert logs[0].dag_id == ti.dag_id
        assert logs[0].run_id == ti.run_id
        assert logs[0].map_index == ti.map_index
        assert logs[0].try_number == ti.try_number
        assert logs[0].logical_date == instant
        assert logs[0].owner == ti.task.owner
        assert logs[0].extra == '{"host_name": "random-hostname"}'


class TestTIUpdateState:
    def setup_method(self):
        clear_db_assets()
        clear_db_logs()
        clear_db_runs()

    def teardown_method(self):
        clear_db_assets()
        clear_db_logs()
        clear_db_runs()

    @pytest.mark.parametrize(
        ("state", "end_date", "expected_state"),
        [
            (State.SUCCESS, DEFAULT_END_DATE, State.SUCCESS),
            (State.FAILED, DEFAULT_END_DATE, State.FAILED),
            (State.SKIPPED, DEFAULT_END_DATE, State.SKIPPED),
        ],
    )
    def test_ti_update_state_to_terminal(
        self, client, session, create_task_instance, state, end_date, expected_state
    ):
        ti = create_task_instance(
            task_id="test_ti_update_state_to_terminal",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": state,
                "end_date": end_date.isoformat(),
            },
        )

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()

        ti = session.get(TaskInstance, ti.id)
        assert ti.state == expected_state
        assert ti.end_date == end_date

    @pytest.mark.parametrize(
        ("payload", "expected_event"),
        [
            pytest.param(
                {"state": State.SUCCESS, "end_date": DEFAULT_END_DATE.isoformat()},
                State.SUCCESS,
                id="success",
            ),
            pytest.param(
                {"state": State.FAILED, "end_date": DEFAULT_END_DATE.isoformat()},
                State.FAILED,
                id="failed",
            ),
            pytest.param(
                {"state": State.SKIPPED, "end_date": DEFAULT_END_DATE.isoformat()},
                State.SKIPPED,
                id="skipped",
            ),
            pytest.param(
                {"state": State.UP_FOR_RETRY, "end_date": DEFAULT_END_DATE.isoformat()},
                TaskInstanceState.UP_FOR_RETRY.value,
                id="up_for_retry",
            ),
            pytest.param(
                {
                    "state": "deferred",
                    "trigger_kwargs": {"key": "value", "moment": "2026-02-18T00:00:00Z"},
                    "trigger_timeout": "P1D",
                    "classpath": "my-classpath",
                    "next_method": "execute_callback",
                },
                TaskInstanceState.DEFERRED.value,
                id="deferred",
            ),
            pytest.param(
                {
                    "state": "up_for_reschedule",
                    "reschedule_date": "2026-02-18T11:03:00+00:00",
                    "end_date": DEFAULT_END_DATE.isoformat(),
                },
                TaskInstanceState.UP_FOR_RESCHEDULE.value,
                id="up_for_reschedule",
            ),
        ],
    )
    def test_ti_update_state_creates_audit_log(
        self, client, session, create_task_instance, payload, expected_event
    ):
        """Test that state transition creates an audit log record."""
        ti = create_task_instance(
            task_id="test_ti_update_state_creates_audit_log",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
            hostname="random-hostname",
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json=payload,
        )

        assert response.status_code == 204

        logs = session.scalars(select(Log).where(Log.dag_id == ti.dag_id)).all()
        assert len(logs) == 1
        assert logs[0].event == expected_event
        assert logs[0].task_id == ti.task_id
        assert logs[0].dag_id == ti.dag_id
        assert logs[0].run_id == ti.run_id
        assert logs[0].map_index == ti.map_index
        assert logs[0].try_number == ti.try_number
        assert logs[0].logical_date == ti.dag_run.logical_date
        assert logs[0].owner == ti.task.owner
        assert logs[0].extra == '{"host_name": "random-hostname"}'

    @pytest.mark.parametrize(
        ("state", "end_date", "expected_state", "rendered_map_index"),
        [
            (State.SUCCESS, DEFAULT_END_DATE, State.SUCCESS, DEFAULT_RENDERED_MAP_INDEX),
            (State.FAILED, DEFAULT_END_DATE, State.FAILED, DEFAULT_RENDERED_MAP_INDEX),
            (State.SKIPPED, DEFAULT_END_DATE, State.SKIPPED, DEFAULT_RENDERED_MAP_INDEX),
        ],
    )
    def test_ti_update_state_to_terminal_with_rendered_map_index(
        self, client, session, create_task_instance, state, end_date, expected_state, rendered_map_index
    ):
        ti = create_task_instance(
            task_id="test_ti_update_state_to_terminal_with_rendered_map_index",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={"state": state, "end_date": end_date.isoformat(), "rendered_map_index": rendered_map_index},
        )

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()

        ti = session.get(TaskInstance, ti.id)
        assert ti.state == expected_state
        assert ti.end_date == end_date
        assert ti.rendered_map_index == rendered_map_index

    @pytest.mark.parametrize(
        "task_outlets",
        [
            pytest.param([{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}], id="asset"),
            pytest.param([{"name": "my-task", "type": "AssetNameRef"}], id="name-ref"),
            pytest.param([{"uri": "s3://bucket/my-task", "type": "AssetUriRef"}], id="uri-ref"),
        ],
    )
    @pytest.mark.parametrize(
        ("outlet_events", "expected_extra"),
        [
            pytest.param([], {}, id="default"),
            pytest.param(
                [
                    {
                        "dest_asset_key": {"name": "my-task", "uri": "s3://bucket/my-task"},
                        "extra": {"foo": 1},
                    },
                    {
                        "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"},
                        "extra": {"foo": 2},
                    },
                ],
                {"foo": 1},
                id="extra",
            ),
        ],
    )
    def test_ti_update_state_to_success_with_asset_events(
        self, client, session, create_task_instance, task_outlets, outlet_events, expected_extra
    ):
        asset = AssetModel(
            id=1,
            name="my-task",
            uri="s3://bucket/my-task",
            group="asset",
            extra={},
        )
        asset_active = AssetActive.for_asset(asset)
        session.add_all([asset, asset_active])

        ti = create_task_instance(
            task_id="test_ti_update_state_to_success_with_asset_events",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "success",
                "end_date": DEFAULT_END_DATE.isoformat(),
                "task_outlets": task_outlets,
                "outlet_events": outlet_events,
            },
        )

        assert response.status_code == 204
        assert response.text == ""
        session.expire_all()

        event = session.scalars(select(AssetEvent)).all()
        assert len(event) == 1
        assert event[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={})
        assert event[0].extra == expected_extra

    def test_ti_update_state_to_success_runs_deferred_asset_listener_callbacks(
        self, client, session, create_task_instance, listener_manager
    ):
        """The success endpoint runs the deferred asset listener callbacks after committing."""
        asset_listener.clear()
        listener_manager(asset_listener)

        asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task", group="asset", extra={})
        session.add_all([asset, AssetActive.for_asset(asset)])

        ti = create_task_instance(
            task_id="test_ti_update_state_to_success_runs_deferred_asset_listener_callbacks",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "success",
                "end_date": DEFAULT_END_DATE.isoformat(),
                "task_outlets": [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}],
                "outlet_events": [],
            },
        )

        assert response.status_code == 204

        # Notifications are deferred during registration and run by the endpoint after the
        # TI state is committed (and the task_instance row lock released).
        assert len(asset_listener.changed) == 1
        assert asset_listener.changed[0].uri == "s3://bucket/my-task"
        assert len(asset_listener.emitted) == 1

    def test_ti_update_state_to_success_rolls_back_partial_asset_registration(
        self, client, session, create_task_instance
    ):
        """A failure partway through asset registration rolls back the partial writes.

        The endpoint's exception handler marks the TI failed and commits; the explicit rollback
        ensures an asset event flushed before the failure is not committed alongside it.
        """
        asset = AssetModel(id=1, name="my-task", uri="s3://bucket/my-task", group="asset", extra={})
        session.add_all([asset, AssetActive.for_asset(asset)])

        ti = create_task_instance(
            task_id="test_ti_update_state_to_success_rolls_back_partial_asset_registration",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        def _partial_then_fail(ti, task_outlets, outlet_events, *, session):
            # Simulate a half-way registration: an asset event is flushed, then registration
            # fails before the endpoint commits.
            session.add(AssetEvent(asset_id=asset.id))
            session.flush()
            raise RuntimeError("boom partway through outlets")

        with mock.patch.object(TaskInstance, "register_asset_changes_in_db", side_effect=_partial_then_fail):
            response = client.patch(
                f"/execution/task-instances/{ti.id}/state",
                json={
                    "state": "success",
                    "end_date": DEFAULT_END_DATE.isoformat(),
                    "task_outlets": [{"name": "my-task", "uri": "s3://bucket/my-task", "type": "Asset"}],
                    "outlet_events": [],
                },
            )

        assert response.status_code == 204
        session.expire_all()
        # The partially-written asset event was rolled back, and the TI is marked failed.
        assert session.scalars(select(AssetEvent)).all() == []
        assert session.get(TaskInstance, ti.id).state == State.FAILED

    @pytest.mark.parametrize(
        ("outlet_events", "expected_extra"),
        [
            pytest.param([], None, id="default"),
            pytest.param(
                [
                    {
                        "dest_asset_key": {"name": "my-task", "uri": "s3://bucket/my-task"},
                        "source_alias_name": "simple1",
                        "extra": {"foo": 1},
                    },
                    {
                        "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"},
                        "extra": {"foo": 2},
                    },
                    {
                        "dest_asset_key": {"name": "my-task-2", "uri": "s3://bucket/my-task-2"},
                        "source_alias_name": "simple2",
                        "extra": {"foo": 3},
                    },
                ],
                {"foo": 1},
                id="extra",
            ),
        ],
    )
    def test_ti_update_state_to_success_with_asset_alias_events(
        self, client, session, create_task_instance, outlet_events, expected_extra
    ):
        asset = AssetModel(
            id=1,
            name="my-task",
            uri="s3://bucket/my-task",
            group="asset",
            extra={},
        )
        asset_active = AssetActive.for_asset(asset)
        session.add_all([asset, asset_active])

        _create_asset_aliases(session, num=2)

        ti = create_task_instance(
            task_id="test_ti_update_state_to_success_with_asset_events",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "success",
                "end_date": DEFAULT_END_DATE.isoformat(),
                "task_outlets": [{"name": "simple1", "type": "AssetAlias"}],
                "outlet_events": outlet_events,
            },
        )

        assert response.status_code == 204
        assert response.text == ""
        session.expire_all()

        events = session.scalars(select(AssetEvent)).all()
        if expected_extra is None:
            assert events == []
        else:
            assert len(events) == 1
            assert events[0].asset == AssetModel(name="my-task", uri="s3://bucket/my-task", extra={})
            assert events[0].extra == expected_extra

    @pytest.mark.parametrize(
        ("partition_key", "expected_status"),
        [
            pytest.param("a" * 250, 204, id="at_limit_accepted"),
            pytest.param("a" * 251, 422, id="over_limit_rejected"),
            pytest.param("", 422, id="empty_rejected"),
            pytest.param("   ", 422, id="whitespace_rejected"),
        ],
    )
    def test_ti_update_state_to_success_outlet_event_partition_key_validation(
        self, client, session, create_task_instance, partition_key, expected_status
    ):
        """Outlet event partition_key content is validated server-side; invalid keys return 422."""
        ti = create_task_instance(
            task_id="test_outlet_event_partition_key_validation",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "success",
                "end_date": DEFAULT_END_DATE.isoformat(),
                "task_outlets": [],
                "outlet_events": [
                    {
                        "dest_asset_key": {"name": "my-asset", "uri": "s3://bucket/my-asset"},
                        "extra": {},
                        "partition_key": partition_key,
                    }
                ],
            },
        )
        assert response.status_code == expected_status
        if expected_status == 422:
            detail = response.json()["detail"]
            assert detail["reason"] == "invalid_partition_key"

    def test_ti_update_state_not_found(self, client, session):
        """
        Test that a 404 error is returned when the Task Instance does not exist.
        """
        task_instance_id = UUID("0182e924-0f1e-77e6-ab50-e977118bc139")

        # Pre-condition: the Task Instance does not exist
        assert session.get(TaskInstance, task_instance_id) is None

        payload = {"state": "success", "end_date": "2024-10-31T12:30:00Z"}

        response = client.patch(f"/execution/task-instances/{task_instance_id}/state", json=payload)
        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task Instance not found",
        }

    def test_ti_update_state_running_errors(self, client, session, create_task_instance, time_machine):
        """
        Test that a 422 error is returned when the Task Instance state is RUNNING in the payload.

        Task should be set to Running state via the /execution/task-instances/{task_instance_id}/run endpoint.
        """
        ti = create_task_instance(
            task_id="test_ti_update_state_running_errors",
            state=State.QUEUED,
            session=session,
            start_date=DEFAULT_START_DATE,
        )

        session.commit()

        response = client.patch(f"/execution/task-instances/{ti.id}/state", json={"state": "running"})

        assert response.status_code == 422

    def test_ti_update_state_database_error(self, client, session, create_task_instance):
        """
        Test that a database error is handled correctly when updating the Task Instance state.
        """
        ti = create_task_instance(
            task_id="test_ti_update_state_database_error",
            state=State.QUEUED,
        )
        session.commit()
        payload = {
            "state": "success",
            "end_date": "2024-10-31T12:00:00Z",
        }

        with (
            mock.patch(
                "airflow.api_fastapi.common.db.common.Session.execute",
                side_effect=[
                    mock.Mock(
                        one=lambda: (
                            "running",
                            1,
                            0,
                            "dag",
                            "task",
                            "run",
                            -1,
                            "localhost",
                            timezone.utcnow(),
                            "test_owner",
                        )
                    ),  # First call returns "queued"
                    mock.Mock(
                        one=lambda: (
                            "running",
                            1,
                            0,
                            "dag",
                            "task",
                            "run",
                            -1,
                            "localhost",
                            timezone.utcnow(),
                            "test_owner",
                        )
                    ),  # Second call returns "queued"
                    SQLAlchemyError("Database error"),  # Last call raises an error
                ],
            ),
            mock.patch(
                "airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db",
            ) as mock_register_asset_changes_in_db,
        ):
            mock_register_asset_changes_in_db.return_value = None
            response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
            assert response.status_code == 500
            detail = response.json()["detail"]
            assert isinstance(detail, dict)
            assert detail.get("reason") == "Database error"

    def test_ti_run_database_error(self, client, session, create_task_instance):
        """
        Test that a database error is handled correctly when starting the Task Instance.
        """
        ti = create_task_instance(
            task_id="test_ti_run_database_error",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            dag_id=str(uuid4()),
        )
        session.commit()

        payload = {
            "state": "running",
            "hostname": "hostname",
            "unixname": "unixname",
            "pid": 123,
            "start_date": "2024-10-31T12:00:00Z",
        }

        with mock.patch(
            "airflow.api_fastapi.common.db.common.Session.execute",
            side_effect=[
                mock.Mock(
                    one=mock.Mock(
                        return_value=SimpleNamespace(
                            state="queued",
                            dag_id="dag",
                            run_id="run",
                            task_id="task",
                            map_index=-1,
                            try_number=1,
                            max_tries=0,
                            start_date=None,
                            next_method=None,
                            hostname=None,
                            unixname=None,
                            pid=None,
                            next_kwargs=None,
                            logical_date=timezone.utcnow(),
                            owners="test_owner",
                        )
                    )
                ),
                SQLAlchemyError("Database error"),
            ],
        ):
            response = client.patch(f"/execution/task-instances/{ti.id}/run", json=payload)

        assert response.status_code == 500
        detail = response.json()["detail"]
        assert isinstance(detail, dict)
        assert detail.get("reason") == "Database error"

    @pytest.mark.parametrize("queues_enabled", [False, True])
    def test_ti_update_state_to_deferred(
        self, client, session, create_task_instance, time_machine, queues_enabled: bool
    ):
        """
        Test that tests if the transition to deferred state is handled correctly.
        """
        with conf_vars({("triggerer", "queues_enabled"): str(queues_enabled)}):
            ti = create_task_instance(
                task_id="test_ti_update_state_to_deferred",
                state=State.RUNNING,
                session=session,
            )
            session.commit()

            instant = timezone.datetime(2024, 11, 22)
            time_machine.move_to(instant, tick=False)

            payload = {
                "state": "deferred",
                # expected format is now in serde serialized format
                "trigger_kwargs": {
                    "key": "value",
                    "moment": {
                        "__classname__": "datetime.datetime",
                        "__version__": 2,
                        "__data__": {
                            "timestamp": 1734480001.0,
                            "tz": {
                                "__classname__": "builtins.tuple",
                                "__version__": 1,
                                "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                            },
                        },
                    },
                },
                "trigger_timeout": "P1D",  # 1 day
                "queue": "default" if queues_enabled else None,
                "classpath": "my-classpath",
                "next_method": "execute_callback",
                # expected format is now in serde serialized format
                "next_kwargs": {
                    "foo": {
                        "__classname__": "datetime.datetime",
                        "__version__": 2,
                        "__data__": {
                            "timestamp": 1734480000.0,
                            "tz": {
                                "__classname__": "builtins.tuple",
                                "__version__": 1,
                                "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                            },
                        },
                    },
                    "bar": "abc",
                },
            }

            response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)

            assert response.status_code == 204
            assert response.text == ""

            session.expire_all()

            tis = session.scalars(select(TaskInstance)).all()
            assert len(tis) == 1

            assert tis[0].state == TaskInstanceState.DEFERRED
            assert tis[0].next_method == "execute_callback"

            assert tis[0].next_kwargs == {
                "foo": {
                    "__classname__": "datetime.datetime",
                    "__version__": 2,
                    "__data__": {
                        "timestamp": 1734480000.0,
                        "tz": {
                            "__classname__": "builtins.tuple",
                            "__version__": 1,
                            "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
                        },
                    },
                },
                "bar": "abc",
            }
            assert tis[0].trigger_timeout == timezone.make_aware(
                datetime(2024, 11, 23), timezone=timezone.utc
            )

            t = session.scalars(select(Trigger)).all()
            assert len(t) == 1
            assert t[0].created_date == instant
            assert t[0].classpath == "my-classpath"
            assert t[0].kwargs == {
                "key": "value",
                "moment": datetime(2024, 12, 18, 00, 00, 1, tzinfo=timezone.utc),
            }
            if queues_enabled:
                assert t[0].queue == "default"
            else:
                assert t[0].queue is None

    @staticmethod
    def _defer_ti_in_team_bundle(client, session, create_task_instance):
        """Map the TI's Dag to a bundle/team, then defer it via the Execution API."""
        from airflow.models.dagbundle import DagBundleModel
        from airflow.models.team import Team

        ti = create_task_instance(
            task_id="test_ti_deferred_team",
            state=State.RUNNING,
            session=session,
        )

        bundle_name = "bundle_deferred_team_test"
        team_name = "team_deferred_test"
        bundle = session.get(DagBundleModel, bundle_name) or DagBundleModel(name=bundle_name)
        team = session.get(Team, team_name) or Team(name=team_name)
        if team not in bundle.teams:
            bundle.teams.append(team)
        session.add(bundle)
        session.flush()
        session.execute(update(DagModel).where(DagModel.dag_id == ti.dag_id).values(bundle_name=bundle_name))
        session.commit()

        payload = {
            "state": "deferred",
            "trigger_kwargs": {"key": "value"},
            "classpath": "my-classpath",
            "next_method": "execute_callback",
        }
        response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
        assert response.status_code == 204
        return team_name

    @conf_vars({("core", "multi_team"): "True"})
    def test_ti_update_state_to_deferred_populates_trigger_team_name(
        self, client, session, create_task_instance, time_machine
    ):
        """Trigger created on deferral gets team_name from the TI's bundle."""
        from airflow.models.trigger import Trigger

        time_machine.move_to(timezone.datetime(2024, 11, 22), tick=False)

        team_name = self._defer_ti_in_team_bundle(client, session, create_task_instance)

        session.expire_all()
        trigger = session.scalars(select(Trigger)).one()
        assert trigger.team_name == team_name

    @conf_vars({("core", "multi_team"): "False"})
    def test_ti_update_state_to_deferred_skips_trigger_team_name_when_multi_team_disabled(
        self, client, session, create_task_instance, time_machine
    ):
        """When multi_team is disabled, the trigger team_name stays NULL."""
        from airflow.models.trigger import Trigger

        time_machine.move_to(timezone.datetime(2024, 11, 22), tick=False)

        self._defer_ti_in_team_bundle(client, session, create_task_instance)

        session.expire_all()
        trigger = session.scalars(select(Trigger)).one()
        assert trigger.team_name is None

    def test_ti_update_state_to_reschedule(self, client, session, create_task_instance, time_machine):
        """
        Test that tests if the transition to reschedule state is handled correctly.
        """
        instant = timezone.datetime(2024, 10, 30)
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_ti_update_state_to_reschedule",
            state=State.RUNNING,
            session=session,
        )
        ti.start_date = instant
        session.commit()

        payload = {
            "state": "up_for_reschedule",
            "reschedule_date": "2024-10-31T11:03:00+00:00",
            "end_date": DEFAULT_END_DATE.isoformat(),
        }

        response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()

        tis = session.scalars(select(TaskInstance)).all()
        assert len(tis) == 1
        assert tis[0].state == TaskInstanceState.UP_FOR_RESCHEDULE
        assert tis[0].next_method is None
        assert tis[0].next_kwargs is None
        assert tis[0].duration == 129600

        trs = session.scalars(select(TaskReschedule)).all()
        assert len(trs) == 1
        assert trs[0].task_instance.dag_id == "dag"
        assert trs[0].task_instance.task_id == "test_ti_update_state_to_reschedule"
        assert trs[0].task_instance.run_id == "test"
        assert trs[0].ti_id == tis[0].id
        assert trs[0].start_date == instant
        assert trs[0].end_date == DEFAULT_END_DATE
        assert trs[0].reschedule_date == timezone.parse("2024-10-31T11:03:00+00:00")
        assert trs[0].task_instance.map_index == -1
        assert trs[0].duration == 129600

    @pytest.mark.backend("mysql")
    def test_ti_update_state_reschedule_mysql_limit(
        self, client, session, create_task_instance, time_machine
    ):
        """Test that the reschedule date is validated against MySQL's TIMESTAMP limit."""
        instant = timezone.datetime(2024, 10, 30)
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_ti_update_state_reschedule_mysql_limit",
            state=State.RUNNING,
            session=session,
        )
        ti.start_date = instant
        session.commit()

        # Date beyond MySQL's TIMESTAMP limit (2038-01-19 03:14:07)
        future_date = timezone.datetime(2038, 1, 19, 3, 14, 8)

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": TaskInstanceState.UP_FOR_RESCHEDULE,
                "reschedule_date": future_date.isoformat(),
                "end_date": DEFAULT_END_DATE.isoformat(),
            },
        )

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()

        ti = session.get(TaskInstance, ti.id)
        assert ti.state == State.FAILED

    # DataError handling: an oversized field the DB rejects must surface as 422 (not an opaque 500
    # or a silent mark-FAILED). Only MySQL/Postgres enforce varchar limits; SQLite ignores them.
    @pytest.mark.backend("mysql", "postgres")
    def test_ti_run_oversized_field_returns_422(self, client, session, create_task_instance):
        """ti_run: a DataError on the running-state UPDATE surfaces as 422, not 500."""
        ti = create_task_instance(
            task_id="test_ti_run_dataerror",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            dag_id=str(uuid4()),
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "h" * 1100,  # > hostname varchar(1000)
                "unixname": "u",
                "pid": 1,
                "start_date": "2024-09-30T12:00:00Z",
            },
        )
        assert response.status_code == 422
        assert response.json()["detail"]["reason"] == "Value rejected by database"

    @pytest.mark.backend("mysql", "postgres")
    def test_ti_update_state_oversized_field_returns_422(self, client, session, create_task_instance):
        """ti_update_state: a DataError on the state UPDATE surfaces as 422, not 500."""
        ti = create_task_instance(
            task_id="test_ti_update_dataerror",
            state=State.RUNNING,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            dag_id=str(uuid4()),
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "up_for_retry",
                "end_date": "2024-09-30T12:00:00Z",
                "rendered_map_index": "x" * 300,  # > rendered_map_index varchar(250)
            },
        )
        assert response.status_code == 422
        assert response.json()["detail"]["reason"] == "Value rejected by database"

    @pytest.mark.backend("mysql", "postgres")
    def test_ti_update_state_dataerror_does_not_silently_fail(self, client, session, create_task_instance):
        """A DataError raised while building the state update must not silently mark the TI FAILED."""
        ti = create_task_instance(
            task_id="test_ti_defer_dataerror",
            state=State.RUNNING,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            dag_id=str(uuid4()),
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "deferred",
                "classpath": "c" * 1100,  # > trigger.classpath varchar(1000)
                "next_method": "execute_complete",
                "trigger_kwargs": {},
                "next_kwargs": {},
            },
        )
        assert response.status_code == 422
        assert response.json()["detail"]["reason"] == "Value rejected by database"

        # The bare `except Exception` would have marked this FAILED behind the worker's back.
        session.expire_all()
        assert session.get(TaskInstance, ti.id).state == State.RUNNING

    def test_ti_update_state_handle_retry(self, client, session, create_task_instance):
        ti = create_task_instance(
            task_id="test_ti_update_state_to_retry",
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": State.UP_FOR_RETRY,
                "end_date": DEFAULT_END_DATE.isoformat(),
            },
        )

        assert response.status_code == 204
        assert response.text == ""

        ti = session.scalar(
            select(TaskInstance).filter_by(task_id=ti.task_id, run_id=ti.run_id, dag_id=ti.dag_id)
        )
        # ti = session.get(TaskInstance, ti.id)
        assert ti.state == State.UP_FOR_RETRY
        assert ti.next_method is None
        assert ti.next_kwargs is None

        tih = session.scalars(
            select(TaskInstanceHistory).where(
                TaskInstanceHistory.task_id == ti.task_id, TaskInstanceHistory.run_id == ti.run_id
            )
        ).one()
        assert tih.task_instance_id
        assert tih.task_instance_id != ti.id

    def test_ti_update_state_retry_with_policy_overrides(self, client, session, create_task_instance):
        """Test that retry_delay_seconds and retry_reason from a RetryPolicy are stored on the TI."""
        ti = create_task_instance(
            task_id="test_retry_policy_override",
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": State.UP_FOR_RETRY,
                "end_date": DEFAULT_END_DATE.isoformat(),
                "retry_delay_seconds": 42.5,
                "retry_reason": "Rate limit: backing off",
            },
        )

        assert response.status_code == 204

        ti = session.scalar(
            select(TaskInstance).filter_by(task_id=ti.task_id, run_id=ti.run_id, dag_id=ti.dag_id)
        )
        assert ti.state == State.UP_FOR_RETRY
        assert ti.retry_delay_override == 42.5
        assert ti.retry_reason == "Rate limit: backing off"

    def test_ti_update_state_retry_policy_overrides_persisted_in_history(
        self, client, session, create_task_instance
    ):
        """The finished try's values must be archived to task_instance_history.

        record_ti() snapshots columns off the TI object, so the overrides, end_date,
        and rendered_map_index must be set on the TI before prepare_db_for_next_try()
        archives it; the live-row UPDATE is not visible to it.
        """
        ti = create_task_instance(
            task_id="test_retry_policy_override_history",
            state=State.RUNNING,
        )
        ti.start_date = DEFAULT_START_DATE
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": State.UP_FOR_RETRY,
                "end_date": DEFAULT_END_DATE.isoformat(),
                "rendered_map_index": DEFAULT_RENDERED_MAP_INDEX,
                "retry_delay_seconds": 42.5,
                "retry_reason": "Rate limit: backing off",
            },
        )

        assert response.status_code == 204

        tih = session.scalars(
            select(TaskInstanceHistory).where(
                TaskInstanceHistory.dag_id == ti.dag_id,
                TaskInstanceHistory.task_id == ti.task_id,
                TaskInstanceHistory.run_id == ti.run_id,
            )
        ).one()
        assert tih.retry_delay_override == 42.5
        assert tih.retry_reason == "Rate limit: backing off"
        assert tih.end_date == DEFAULT_END_DATE
        assert tih.duration == (DEFAULT_END_DATE - DEFAULT_START_DATE).total_seconds()
        assert tih.rendered_map_index == DEFAULT_RENDERED_MAP_INDEX

    def test_ti_update_state_retry_clears_rendered_map_index_in_history(
        self, client, session, create_task_instance
    ):
        """An explicit ``rendered_map_index: null`` must clear the value archived to history too.

        The worker sends ``rendered_map_index`` on every retry, even when this try never
        (re-)computed it (e.g. it failed before rendering). That must null out the archived
        row along with the live one, not leave the history row snapshotting a stale value
        left over from an earlier try.
        """
        ti = create_task_instance(
            task_id="test_retry_clears_rendered_map_index_history",
            state=State.RUNNING,
        )
        ti.start_date = DEFAULT_START_DATE
        ti._rendered_map_index = DEFAULT_RENDERED_MAP_INDEX
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": State.UP_FOR_RETRY,
                "end_date": DEFAULT_END_DATE.isoformat(),
                "rendered_map_index": None,
            },
        )

        assert response.status_code == 204

        ti = session.scalars(
            select(TaskInstance).filter_by(task_id=ti.task_id, run_id=ti.run_id, dag_id=ti.dag_id)
        ).one()
        assert ti.rendered_map_index is None

        tih = session.scalars(
            select(TaskInstanceHistory).where(
                TaskInstanceHistory.dag_id == ti.dag_id,
                TaskInstanceHistory.task_id == ti.task_id,
                TaskInstanceHistory.run_id == ti.run_id,
            )
        ).one()
        assert tih.rendered_map_index is None

    def test_ti_update_state_retry_without_policy_overrides(self, client, session, create_task_instance):
        """Without retry policy fields, the columns remain NULL."""
        ti = create_task_instance(
            task_id="test_retry_no_policy",
            state=State.RUNNING,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": State.UP_FOR_RETRY,
                "end_date": DEFAULT_END_DATE.isoformat(),
            },
        )

        assert response.status_code == 204

        ti = session.scalar(
            select(TaskInstance).filter_by(task_id=ti.task_id, run_id=ti.run_id, dag_id=ti.dag_id)
        )
        assert ti.state == State.UP_FOR_RETRY
        assert ti.retry_delay_override is None
        assert ti.retry_reason is None

        tih = session.scalars(
            select(TaskInstanceHistory).where(
                TaskInstanceHistory.dag_id == ti.dag_id,
                TaskInstanceHistory.task_id == ti.task_id,
                TaskInstanceHistory.run_id == ti.run_id,
            )
        ).one()
        assert tih.retry_delay_override is None
        assert tih.retry_reason is None

    def test_ti_run_clears_retry_policy_overrides(self, client, session, create_task_instance):
        """When a task enters RUNNING, retry policy overrides from the previous attempt are cleared."""
        ti = create_task_instance(
            task_id="test_retry_cleared_on_run",
            state=State.QUEUED,
        )
        # Simulate having retry overrides from a previous attempt
        ti.retry_delay_override = 60.0
        ti.retry_reason = "previous attempt rate limit"
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "test-host",
                "unixname": "test-user",
                "pid": 12345,
                "start_date": DEFAULT_END_DATE.isoformat(),
            },
        )

        assert response.status_code == 200

        session.expire_all()
        ti = session.scalar(
            select(TaskInstance).filter_by(task_id=ti.task_id, run_id=ti.run_id, dag_id=ti.dag_id)
        )
        assert ti.state == State.RUNNING
        assert ti.retry_delay_override is None
        assert ti.retry_reason is None

    @pytest.mark.parametrize(
        "target_state",
        [
            State.UP_FOR_RETRY,
            TerminalTIState.FAILED,
            TerminalTIState.SUCCESS,
        ],
    )
    def test_ti_update_state_clears_deferred_fields(
        self, client, session, create_task_instance, target_state
    ):
        """Test that next_method and next_kwargs are cleared when transitioning to terminal/retry states."""
        ti = create_task_instance(
            task_id="test_ti_update_state_clears_deferred_fields",
            state=State.RUNNING,
        )
        # Simulate a task that resumed from deferred state with next_method/next_kwargs set
        ti.next_method = "execute_complete"
        ti.next_kwargs = {"event": "test_event", "data": "test_data"}
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": target_state,
                "end_date": DEFAULT_END_DATE.isoformat(),
            },
        )

        assert response.status_code == 204

        if target_state == State.UP_FOR_RETRY:
            # Retry creates a new TI ID, so we need to fetch by unique key
            ti = session.scalar(
                select(TaskInstance).filter_by(task_id=ti.task_id, run_id=ti.run_id, dag_id=ti.dag_id)
            )
        else:
            session.expire_all()
            ti = session.get(TaskInstance, ti.id)

        assert ti.state == target_state
        assert ti.next_method is None
        assert ti.next_kwargs is None

    def test_ti_update_state_to_failed_table_check(self, client, session, create_task_instance):
        # we just want to fail in this test, no need to retry
        ti = create_task_instance(
            task_id="test_ti_update_state_to_failed_table_check",
            state=State.RUNNING,
        )
        ti.start_date = DEFAULT_START_DATE
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": TerminalTIState.FAILED,
                "end_date": DEFAULT_END_DATE.isoformat(),
            },
        )

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()

        ti = session.get(TaskInstance, ti.id)
        assert ti.state == State.FAILED
        assert ti.next_method is None
        assert ti.next_kwargs is None
        assert ti.duration == 3600.00

    def test_ti_update_state_not_running(self, client, session, create_task_instance):
        """Test that a 409 error is returned when attempting to update a TI that is not in RUNNING state."""
        ti = create_task_instance(
            task_id="test_ti_update_state_not_running",
            state=State.SUCCESS,
            session=session,
            start_date=DEFAULT_START_DATE,
        )
        session.commit()

        payload = {
            "state": "failed",
            "end_date": DEFAULT_END_DATE.isoformat(),
        }

        response = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
        assert response.status_code == 409
        assert response.json()["detail"] == {
            "reason": "invalid_state",
            "message": "TI was not in the running state so it cannot be updated",
            "previous_state": State.SUCCESS,
        }

        # Verify the task instance state hasn't changed
        session.refresh(ti)
        assert ti.state == State.SUCCESS

    def test_ti_update_state_same_state_is_idempotent(self, client, session, create_task_instance):
        """Test that setting a TI to its current state is treated as an idempotent no-op."""
        ti = create_task_instance(
            task_id="test_ti_update_state_same_state_is_idempotent",
            state=State.SUCCESS,
            session=session,
            start_date=DEFAULT_START_DATE,
        )
        ti.end_date = DEFAULT_END_DATE
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "success",
                "end_date": timezone.parse("2024-10-31T13:00:00Z").isoformat(),
            },
        )
        assert response.status_code == 200
        assert response.content == b""

        session.refresh(ti)
        assert ti.state == State.SUCCESS
        assert ti.end_date == DEFAULT_END_DATE

    def test_ti_update_state_terminal_state_mismatch_returns_conflict(
        self, client, session, create_task_instance
    ):
        """A completed TI cannot be updated to a different state."""
        ti = create_task_instance(
            task_id="test_ti_update_state_terminal_state_mismatch_returns_conflict",
            state=State.SUCCESS,
            session=session,
            start_date=DEFAULT_START_DATE,
        )
        ti.end_date = DEFAULT_END_DATE
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={
                "state": "failed",
                "end_date": timezone.parse("2024-10-31T13:00:00Z").isoformat(),
            },
        )
        assert response.status_code == 409
        assert response.json()["detail"] == {
            "reason": "invalid_state",
            "message": "TI was not in the running state so it cannot be updated",
            "previous_state": State.SUCCESS,
        }

        session.refresh(ti)
        assert ti.state == State.SUCCESS
        assert ti.end_date == DEFAULT_END_DATE

    def test_ti_update_state_to_failed_without_fail_fast(self, client, session, dag_maker):
        """Test that SerializedDAG is NOT loaded when fail_fast=False (default)."""
        with dag_maker(dag_id="test_dag_no_fail_fast", serialized=True):
            EmptyOperator(task_id="task1")
            EmptyOperator(task_id="task2")

        dr = dag_maker.create_dagrun()
        ti = dr.get_task_instance(task_id="task1", session=session)
        ti.state = State.RUNNING
        ti.start_date = DEFAULT_START_DATE
        session.commit()
        session.refresh(ti)

        with mock.patch("airflow.models.dagbag.DBDagBag.get_dag_for_run", autospec=True) as mock_get_dag:
            response = client.patch(
                f"/execution/task-instances/{ti.id}/state",
                json={
                    "state": TerminalTIState.FAILED,
                    "end_date": DEFAULT_END_DATE.isoformat(),
                },
            )

            assert response.status_code == 204
            # Verify SerializedDAG was NOT loaded (memory optimization)
            mock_get_dag.assert_not_called()

        session.expire_all()
        ti = session.get(TaskInstance, ti.id)
        assert ti.state == State.FAILED

    def test_ti_update_state_to_failed_with_fail_fast(self, client, session, dag_maker):
        """Test that SerializedDAG IS loaded and other tasks stopped when fail_fast=True."""
        with dag_maker(dag_id="test_dag_with_fail_fast", fail_fast=True, serialized=True):
            EmptyOperator(task_id="task1")
            EmptyOperator(task_id="task2")

        dr = dag_maker.create_dagrun()
        ti1 = dr.get_task_instance(task_id="task1", session=session)
        ti1.state = State.RUNNING
        ti1.start_date = DEFAULT_START_DATE

        ti2 = dr.get_task_instance(task_id="task2", session=session)
        ti2.state = State.QUEUED
        session.commit()
        session.refresh(ti1)
        session.refresh(ti2)

        with mock.patch(
            "airflow.api_fastapi.execution_api.routes.task_instances._stop_remaining_tasks", autospec=True
        ) as mock_stop:
            response = client.patch(
                f"/execution/task-instances/{ti1.id}/state",
                json={
                    "state": TerminalTIState.FAILED,
                    "end_date": DEFAULT_END_DATE.isoformat(),
                },
            )

            assert response.status_code == 204
            mock_stop.assert_called_once()

        session.expire_all()
        ti1 = session.get(TaskInstance, ti1.id)
        assert ti1.state == State.FAILED

    def test_ti_update_state_reschedule_mysql_limit_triggers_fail_fast(
        self, client, session, dag_maker, time_machine
    ):
        """
        When a reschedule date exceeds MySQL's TIMESTAMP limit and the DAG has fail_fast=True,
        sibling tasks must still be stopped. The MySQL-limit branch routes through a different
        FAILED transition than the regular fail path -- both must honor fail_fast.
        """
        instant = timezone.datetime(2024, 10, 30)
        time_machine.move_to(instant, tick=False)

        with dag_maker(dag_id="test_dag_with_fail_fast_mysql_reschedule", fail_fast=True, serialized=True):
            EmptyOperator(task_id="task1")
            EmptyOperator(task_id="task2")

        dr = dag_maker.create_dagrun()
        ti1 = dr.get_task_instance(task_id="task1", session=session)
        ti1.state = State.RUNNING
        ti1.start_date = instant

        ti2 = dr.get_task_instance(task_id="task2", session=session)
        ti2.state = State.QUEUED
        session.commit()
        session.refresh(ti1)
        session.refresh(ti2)

        # Date beyond MySQL's TIMESTAMP limit (2038-01-19 03:14:07).
        future_date = timezone.datetime(2038, 1, 19, 3, 14, 8)

        with (
            mock.patch(
                "airflow.api_fastapi.execution_api.routes.task_instances.get_dialect_name",
                return_value="mysql",
            ),
            mock.patch(
                "airflow.api_fastapi.execution_api.routes.task_instances._stop_remaining_tasks",
                autospec=True,
            ) as mock_stop,
        ):
            response = client.patch(
                f"/execution/task-instances/{ti1.id}/state",
                json={
                    "state": TaskInstanceState.UP_FOR_RESCHEDULE,
                    "reschedule_date": future_date.isoformat(),
                    "end_date": DEFAULT_END_DATE.isoformat(),
                },
            )

            assert response.status_code == 204
            mock_stop.assert_called_once()

        session.expire_all()
        ti1 = session.get(TaskInstance, ti1.id)
        assert ti1.state == State.FAILED

    @pytest.mark.db_test
    @conf_vars({("state_store", "clear_on_success"): "True"})
    def test_ti_update_state_to_success_clears_task_state(self, client, session, create_task_instance):
        """When clear_on_success=True, task_state rows are deleted after TI transitions to SUCCESS."""
        ti = create_task_instance(
            task_id="test_clear_on_success",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        backend = MetastoreBackend()
        scope = TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
        backend.set(scope, "job_id", "app_1234", session=session)
        backend.set(scope, "checkpoint", "step_3", session=session)
        session.commit()

        assert session.scalars(
            select(TaskStateStoreModel).where(TaskStateStoreModel.task_id == ti.task_id)
        ).all()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
        )

        assert response.status_code == 204
        session.expire_all()
        assert not session.scalars(
            select(TaskStateStoreModel).where(TaskStateStoreModel.task_id == ti.task_id)
        ).all()

    @pytest.mark.db_test
    @conf_vars({("state_store", "clear_on_success"): "True"})
    def test_ti_update_state_to_failed_does_not_clear_task_state(self, client, session, create_task_instance):
        """Task state rows are preserved when a TI transitions to FAILED."""
        ti = create_task_instance(
            task_id="test_no_clear_on_failed",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        backend = MetastoreBackend()
        scope = TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
        backend.set(scope, "job_id", "app_1234", session=session)
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={"state": "failed", "end_date": DEFAULT_END_DATE.isoformat()},
        )

        assert response.status_code == 204
        session.expire_all()
        assert session.scalars(
            select(TaskStateStoreModel).where(TaskStateStoreModel.task_id == ti.task_id)
        ).all()

    @pytest.mark.db_test
    @conf_vars({("state_store", "clear_on_success"): "False"})
    def test_ti_update_state_to_success_skips_clear_when_config_disabled(
        self, client, session, create_task_instance
    ):
        """Task state rows are preserved on SUCCESS when clear_on_success=False."""
        ti = create_task_instance(
            task_id="test_clear_disabled",
            start_date=DEFAULT_START_DATE,
            state=State.RUNNING,
        )
        session.commit()

        backend = MetastoreBackend()
        scope = TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index)
        backend.set(scope, "job_id", "app_1234", session=session)
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/state",
            json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
        )

        assert response.status_code == 204
        session.expire_all()
        assert session.scalars(
            select(TaskStateStoreModel).where(TaskStateStoreModel.task_id == ti.task_id)
        ).all()


class TestTISkipDownstream:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    @pytest.mark.parametrize("_json", (({"tasks": ["t1"]}), ({"tasks": [("t1", -1)]})))
    def test_ti_skip_downstream(self, client, session, create_task_instance, dag_maker, _json):
        with dag_maker("skip_downstream_dag", session=session):
            t0 = EmptyOperator(task_id="t0")
            t1 = EmptyOperator(task_id="t1")
            t0 >> t1
        dr = dag_maker.create_dagrun(run_id="run")

        ti0 = dr.get_task_instance("t0")
        ti0.set_state(State.SUCCESS)

        response = client.patch(
            f"/execution/task-instances/{ti0.id}/skip-downstream",
            json=_json,
        )
        ti1 = dr.get_task_instance("t1")

        assert response.status_code == 204
        assert ti1.state == State.SKIPPED


class TestTISkipDownstreamRaceCondition:
    """Regression tests for #59378: state guard in ti_skip_downstream()."""

    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    @pytest.mark.parametrize(
        "initial_state",
        [
            State.RUNNING,
            State.SUCCESS,
            State.FAILED,
        ],
    )
    def test_skip_downstream_does_not_overwrite_terminal_or_running_ti(
        self, client, session, dag_maker, initial_state
    ):
        with dag_maker(f"skip_race_dag_{initial_state}", session=session):
            branch = EmptyOperator(task_id="branch")
            downstream = EmptyOperator(task_id="downstream")
            branch >> downstream
        dr = dag_maker.create_dagrun(run_id="run")

        ti_branch = dr.get_task_instance("branch")
        ti_branch.set_state(State.SUCCESS)

        ti_downstream = dr.get_task_instance("downstream")
        ti_downstream.set_state(initial_state)
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti_branch.id}/skip-downstream",
            json={"tasks": ["downstream"]},
        )
        assert response.status_code == 204

        session.expire_all()
        ti_downstream = dr.get_task_instance("downstream")
        assert ti_downstream.state == initial_state

    def test_skip_downstream_does_skip_queued_ti(self, client, session, dag_maker):
        with dag_maker("skip_race_dag_queued", session=session):
            branch = EmptyOperator(task_id="branch")
            downstream = EmptyOperator(task_id="downstream")
            branch >> downstream
        dr = dag_maker.create_dagrun(run_id="run")

        ti_branch = dr.get_task_instance("branch")
        ti_branch.set_state(State.SUCCESS)

        ti_downstream = dr.get_task_instance("downstream")
        ti_downstream.set_state(TaskInstanceState.QUEUED)
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti_branch.id}/skip-downstream",
            json={"tasks": ["downstream"]},
        )
        assert response.status_code == 204

        session.expire_all()
        ti_downstream = dr.get_task_instance("downstream")
        assert ti_downstream.state == State.SKIPPED

    def test_skip_downstream_still_skips_none_state_ti(self, client, session, dag_maker):
        with dag_maker("skip_race_dag_normal", session=session):
            branch = EmptyOperator(task_id="branch")
            downstream = EmptyOperator(task_id="downstream")
            branch >> downstream
        dr = dag_maker.create_dagrun(run_id="run")

        ti_branch = dr.get_task_instance("branch")
        ti_branch.set_state(State.SUCCESS)
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti_branch.id}/skip-downstream",
            json={"tasks": ["downstream"]},
        )
        assert response.status_code == 204

        session.expire_all()
        ti_downstream = dr.get_task_instance("downstream")
        assert ti_downstream.state == State.SKIPPED


class TestTIHealthEndpoint:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    @pytest.mark.parametrize(
        ("hostname", "pid", "expected_status_code", "expected_detail"),
        [
            # Case: Successful heartbeat
            ("random-hostname", 1789, 204, None),
            # Case: Conflict due to hostname mismatch
            (
                "wrong-hostname",
                1789,
                409,
                {
                    "reason": "running_elsewhere",
                    "message": "TI is already running elsewhere",
                    "current_hostname": "random-hostname",
                    "current_pid": 1789,
                },
            ),
            # Case: Conflict due to pid mismatch
            (
                "random-hostname",
                1054,
                409,
                {
                    "reason": "running_elsewhere",
                    "message": "TI is already running elsewhere",
                    "current_hostname": "random-hostname",
                    "current_pid": 1789,
                },
            ),
        ],
    )
    def test_ti_heartbeat(
        self,
        client,
        session,
        create_task_instance,
        hostname,
        pid,
        expected_status_code,
        expected_detail,
        time_machine,
    ):
        """Test the TI heartbeat endpoint for various scenarios including conflicts."""
        time_now = timezone.parse("2024-10-31T12:00:00Z")

        # Freeze time to a specific time
        time_machine.move_to(time_now, tick=False)

        ti = create_task_instance(
            task_id="test_ti_heartbeat",
            state=State.RUNNING,
            hostname="random-hostname",
            pid=1789,
            session=session,
        )
        session.commit()
        task_instance_id = ti.id

        # Pre-condition: TI heartbeat is NONE
        assert ti.last_heartbeat_at is None

        response = client.put(
            f"/execution/task-instances/{task_instance_id}/heartbeat",
            json={"hostname": hostname, "pid": pid},
        )

        assert response.status_code == expected_status_code

        if expected_status_code == 204:
            # If successful, ensure last_heartbeat_at is updated
            session.refresh(ti)
            assert ti.last_heartbeat_at == time_now
            assert response.text == ""
        else:
            # If there's an error, check the error detail
            assert response.json()["detail"] == expected_detail

    def test_ti_heartbeat_non_existent_task(self, client, session, create_task_instance):
        """Test that a 404 error is returned when the Task Instance does not exist."""
        task_instance_id = UUID("0182e924-0f1e-77e6-ab50-e977118bc139")

        # Pre-condition: the Task Instance does not exist
        assert session.get(TaskInstance, task_instance_id) is None

        response = client.put(
            f"/execution/task-instances/{task_instance_id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task Instance not found",
        }

    def test_ti_heartbeat_cleared_task_returns_410(self, client, session, create_task_instance):
        """Test that a 410 error is returned when a TI was cleared and moved to TIH."""
        ti = create_task_instance(
            task_id="test_ti_heartbeat_cleared",
            state=State.RUNNING,
            hostname="random-hostname",
            pid=1547,
            session=session,
        )
        session.commit()
        old_ti_id = ti.id

        # Simulate task being cleared: this archives the current try to TIH
        # and assigns a new UUID to the TI, mirroring prepare_db_for_next_try().
        ti.prepare_db_for_next_try(session)
        session.commit()

        assert session.get(TaskInstance, old_ti_id) is None
        tih = session.scalar(
            select(TaskInstanceHistory).where(TaskInstanceHistory.task_instance_id == old_ti_id)
        )
        assert tih is not None

        response = client.put(
            f"/execution/task-instances/{old_ti_id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 410
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task Instance not found, it may have been moved to the Task Instance History table",
        }

    @pytest.mark.parametrize(
        "ti_state",
        [State.SUCCESS, State.FAILED],
    )
    def test_ti_heartbeat_when_task_not_running(self, client, session, create_task_instance, ti_state):
        """Test that a 409 error is returned when the Task Instance is not in RUNNING state."""
        ti = create_task_instance(
            task_id="test_ti_heartbeat_when_task_not_running",
            state=ti_state,
            hostname="random-hostname",
            pid=1547,
            session=session,
        )
        session.commit()
        task_instance_id = ti.id

        response = client.put(
            f"/execution/task-instances/{task_instance_id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 409
        assert response.json()["detail"] == {
            "reason": "not_running",
            "message": "TI is no longer in the running state and task should terminate",
            "current_state": ti_state,
        }

    def test_ti_heartbeat_update(self, client, session, create_task_instance, time_machine):
        """Test that the Task Instance heartbeat is updated when the Task Instance is running."""
        # Set initial time for the test
        time_now = timezone.parse("2024-10-31T12:00:00Z")
        time_machine.move_to(time_now, tick=False)

        ti = create_task_instance(
            task_id="test_ti_heartbeat_update",
            state=State.RUNNING,
            hostname="random-hostname",
            pid=1547,
            last_heartbeat_at=time_now,
            session=session,
        )
        session.commit()
        task_instance_id = ti.id

        # Pre-condition: TI heartbeat is set
        assert ti.last_heartbeat_at == time_now, "Initial last_heartbeat_at should match time_now"

        # Move time forward by 10 minutes
        new_time = time_now.add(minutes=10)
        time_machine.move_to(new_time, tick=False)

        response = client.put(
            f"/execution/task-instances/{task_instance_id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 204

        # If successful, ensure last_heartbeat_at is updated
        session.refresh(ti)
        assert ti.last_heartbeat_at == time_now.add(minutes=10)

    def test_ti_heartbeat_fast_path_skips_fallback(
        self, client, session, create_task_instance, monkeypatch, time_machine
    ):
        """When the fast-path UPDATE succeeds, the fallback path does not run."""
        time_now = timezone.parse("2024-10-31T12:00:00Z")
        time_machine.move_to(time_now, tick=False)

        ti = create_task_instance(
            task_id="test_ti_heartbeat_fast_path_skips_fallback",
            state=State.RUNNING,
            hostname="random-hostname",
            pid=1547,
            last_heartbeat_at=time_now,
            session=session,
        )
        session.commit()

        new_time = time_now.add(minutes=10)
        time_machine.move_to(new_time, tick=False)

        original_execute = Session.execute
        task_instance_updates = []
        for_update_selects = []

        def counting_execute(session_obj, statement, *args, **kwargs):
            if _is_task_instance_update(statement):
                task_instance_updates.append(statement)
            if _is_select_for_update(statement):
                for_update_selects.append(statement)
            return original_execute(session_obj, statement, *args, **kwargs)

        monkeypatch.setattr(Session, "execute", counting_execute)

        response = client.put(
            f"/execution/task-instances/{ti.id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 204
        assert len(task_instance_updates) == 1
        assert _where_column_keys(task_instance_updates[0]) == {"id", "state", "hostname", "pid"}
        assert len(for_update_selects) == 0
        session.refresh(ti)
        assert ti.last_heartbeat_at == new_time

    def test_ti_heartbeat_fallback_updates_on_fast_path_miss(
        self, client, session, create_task_instance, monkeypatch, time_machine
    ):
        """When the fast-path UPDATE returns rowcount=0 the fallback path should
        still update last_heartbeat_at."""
        time_now = timezone.parse("2024-10-31T12:00:00Z")
        time_machine.move_to(time_now, tick=False)

        ti = create_task_instance(
            task_id="test_ti_heartbeat_fallback_updates_on_fast_path_miss",
            state=State.RUNNING,
            hostname="random-hostname",
            pid=1547,
            last_heartbeat_at=time_now,
            session=session,
        )
        session.commit()

        new_time = time_now.add(minutes=10)
        time_machine.move_to(new_time, tick=False)

        original_execute = Session.execute
        fast_path_intercepted = False

        def execute_with_fast_path_miss(session_obj, statement, *args, **kwargs):
            nonlocal fast_path_intercepted
            if (
                not fast_path_intercepted
                and getattr(statement, "is_update", False)
                and statement.table.name == TaskInstance.__table__.name
            ):
                fast_path_intercepted = True
                return mock.MagicMock(rowcount=0)
            return original_execute(session_obj, statement, *args, **kwargs)

        monkeypatch.setattr(Session, "execute", execute_with_fast_path_miss)

        response = client.put(
            f"/execution/task-instances/{ti.id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 204
        assert fast_path_intercepted
        session.refresh(ti)
        assert ti.last_heartbeat_at == new_time

    def test_ti_heartbeat_fallback_updates_on_unknown_fast_path_rowcount(
        self, client, session, create_task_instance, monkeypatch, time_machine
    ):
        """A truthy-but-unknown rowcount must not be treated as fast-path success."""
        time_now = timezone.parse("2024-10-31T12:00:00Z")
        time_machine.move_to(time_now, tick=False)

        ti = create_task_instance(
            task_id="test_ti_heartbeat_fallback_updates_on_unknown_fast_path_rowcount",
            state=State.RUNNING,
            hostname="random-hostname",
            pid=1547,
            last_heartbeat_at=time_now,
            session=session,
        )
        session.commit()

        new_time = time_now.add(minutes=10)
        time_machine.move_to(new_time, tick=False)

        original_execute = Session.execute
        fast_path_intercepted = False

        def execute_with_unknown_fast_path_rowcount(session_obj, statement, *args, **kwargs):
            nonlocal fast_path_intercepted
            if (
                not fast_path_intercepted
                and getattr(statement, "is_update", False)
                and statement.table.name == TaskInstance.__table__.name
            ):
                fast_path_intercepted = True
                return mock.MagicMock(rowcount=-1)
            return original_execute(session_obj, statement, *args, **kwargs)

        monkeypatch.setattr(Session, "execute", execute_with_unknown_fast_path_rowcount)

        response = client.put(
            f"/execution/task-instances/{ti.id}/heartbeat",
            json={"hostname": "random-hostname", "pid": 1547},
        )

        assert response.status_code == 204
        assert fast_path_intercepted
        session.refresh(ti)
        assert ti.last_heartbeat_at == new_time


class TestTIPutRTIF:
    def setup_method(self):
        clear_db_runs()
        clear_rendered_ti_fields()

    def teardown_method(self):
        clear_db_runs()
        clear_rendered_ti_fields()

    @pytest.mark.parametrize(
        "payload",
        [
            # string value
            {"field1": "string_value", "field2": "another_string"},
            # dictionary value
            {"field1": {"nested_key": "nested_value"}},
            # string lists value
            {"field1": ["123"], "field2": ["a", "b", "c"]},
            # list of JSON values
            {"field1": [1, "string", 3.14, True, None, {"nested": "dict"}]},
            # nested dictionary with mixed types in lists
            {
                "field1": {"nested_dict": {"key1": 123, "key2": "value"}},
                "field2": [3.14, {"sub_key": "sub_value"}, [1, 2]],
            },
        ],
    )
    def test_ti_put_rtif_success(self, client, session, create_task_instance, payload):
        ti = create_task_instance(
            task_id="test_ti_put_rtif_success",
            state=State.RUNNING,
            session=session,
        )
        session.commit()
        response = client.put(f"/execution/task-instances/{ti.id}/rtif", json=payload)
        assert response.status_code == 201
        assert response.json() == {"message": "Rendered task instance fields successfully set"}

        session.expire_all()

        rtifs = session.scalars(select(RenderedTaskInstanceFields)).all()
        assert len(rtifs) == 1

        assert rtifs[0].dag_id == "dag"
        assert rtifs[0].run_id == "test"
        assert rtifs[0].task_id == "test_ti_put_rtif_success"
        assert rtifs[0].map_index == -1
        assert rtifs[0].rendered_fields == payload

    def test_ti_put_rtif_missing_ti(self, client, session, create_task_instance):
        create_task_instance(
            task_id="test_ti_put_rtif_missing_ti",
            state=State.RUNNING,
            session=session,
        )
        session.commit()

        payload = {"field1": "rendered_value1", "field2": "rendered_value2"}

        random_id = uuid6.uuid7()
        response = client.put(f"/execution/task-instances/{random_id}/rtif", json=payload)
        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task Instance not found",
        }

    def test_ti_put_rtif_archived_ti_returns_410(self, client, session, create_task_instance):
        ti = create_task_instance(
            task_id="test_ti_put_rtif_archived",
            state=State.RUNNING,
            session=session,
        )
        session.commit()
        old_ti_id = ti.id

        # Archive the current try to TIH and assign a new UUID, mirroring prepare_db_for_next_try().
        ti.prepare_db_for_next_try(session)
        session.commit()

        assert session.get(TaskInstance, old_ti_id) is None
        assert session.get(TaskInstanceHistory, old_ti_id) is not None

        response = client.put(
            f"/execution/task-instances/{old_ti_id}/rtif",
            json={"field1": "rendered_value1"},
        )

        assert response.status_code == 410
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task Instance not found, it may have been moved to the Task Instance History table",
        }


class TestPreviousDagRun:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    def test_ti_previous_dag_run(self, client, session, create_task_instance, dag_maker):
        """Test that the previous dag run is returned correctly for a task instance."""
        ti = create_task_instance(
            task_id="test_ti_previous_dag_run",
            dag_id="test_dag",
            logical_date=timezone.datetime(2025, 1, 19),
            state=State.RUNNING,
            start_date=timezone.datetime(2024, 1, 17),
            session=session,
        )
        session.commit()

        # Create 2 DagRuns for the same DAG to verify that the correct DagRun (last) is returned
        dr1 = dag_maker.create_dagrun(
            run_id="test_run_id_1",
            logical_date=timezone.datetime(2025, 1, 17),
            run_type="scheduled",
            state=State.SUCCESS,
            session=session,
        )
        dr1.end_date = timezone.datetime(2025, 1, 17, 1, 0, 0)

        dr2 = dag_maker.create_dagrun(
            run_id="test_run_id_2",
            logical_date=timezone.datetime(2025, 1, 18),
            run_type="scheduled",
            state=State.SUCCESS,
            session=session,
        )

        dr2.end_date = timezone.datetime(2025, 1, 18, 1, 0, 0)

        session.commit()

        response = client.get(f"/execution/task-instances/{ti.id}/previous-successful-dagrun")
        assert response.status_code == 200
        assert response.json() == {
            "data_interval_start": "2025-01-18T00:00:00Z",
            "data_interval_end": "2025-01-19T00:00:00Z",
            "start_date": "2024-01-17T00:00:00Z",
            "end_date": "2025-01-18T01:00:00Z",
        }

    def test_ti_previous_dag_run_not_found(self, client, session):
        ti_id = UUID("0182e924-0f1e-77e6-ab50-e977118bc139")

        assert session.get(TaskInstance, ti_id) is None

        response = client.get(f"/execution/task-instances/{ti_id}/previous-successful-dagrun")
        assert response.status_code == 200
        assert response.json() == {
            "data_interval_start": None,
            "data_interval_end": None,
            "start_date": None,
            "end_date": None,
        }

    def test_ti_with_none_as_logical_date(self, client, session, create_task_instance, dag_maker):
        ti = create_task_instance(
            task_id="test_ti_with_none_as_logical_date",
            dag_id="test_dag",
            logical_date=None,
            state=State.RUNNING,
            start_date=timezone.datetime(2024, 1, 17),
            session=session,
        )
        session.commit()

        assert ti.logical_date is None

        dr1 = dag_maker.create_dagrun(
            run_id="test_ti_with_none_as_logical_date",
            logical_date=timezone.datetime(2025, 1, 17),
            run_type="scheduled",
            state=State.SUCCESS,
            session=session,
        )
        dr1.end_date = timezone.datetime(2025, 1, 17, 1, 0, 0)

        session.commit()

        response = client.get(f"/execution/task-instances/{ti.id}/previous-successful-dagrun")
        assert response.status_code == 200
        assert response.json() == {
            "data_interval_start": None,
            "data_interval_end": None,
            "start_date": None,
            "end_date": None,
        }


class TestGetRescheduleStartDate:
    def test_get_start_date(self, client, session, create_task_instance):
        ti = create_task_instance(
            task_id="test_ti_update_state_reschedule_mysql_limit",
            state=State.RUNNING,
            start_date=timezone.datetime(2024, 1, 1),
            session=session,
        )
        tr = TaskReschedule(
            ti_id=ti.id,
            start_date=timezone.datetime(2024, 1, 1),
            end_date=timezone.datetime(2024, 1, 1, 1),
            reschedule_date=timezone.datetime(2024, 1, 1, 2),
        )
        session.add(tr)
        session.commit()

        response = client.get(f"/execution/task-reschedules/{ti.id}/start_date")
        assert response.status_code == 200
        assert response.json() == "2024-01-01T00:00:00Z"

    def test_get_start_date_not_found(self, client):
        ti_id = "0182e924-0f1e-77e6-ab50-e977118bc139"
        response = client.get(f"/execution/task-reschedules/{ti_id}/start_date")
        assert response.json() is None


class TestGetCount:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    def test_get_count_basic(self, client, session, create_task_instance):
        create_task_instance(task_id="test_task", state=State.SUCCESS)
        session.commit()

        response = client.get("/execution/task-instances/count", params={"dag_id": "dag"})
        assert response.status_code == 200
        assert response.json() == 1

    def test_get_count_with_task_ids(self, client, session, create_task_instance):
        for i in range(3):
            create_task_instance(
                task_id=f"task{i}",
                state=State.SUCCESS,
                dag_id="test_get_count_with_task_ids",
                run_id=f"test_run_id{i}",
            )
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "test_get_count_with_task_ids", "task_ids": ["task1", "task2"]},
        )
        assert response.status_code == 200
        assert response.json() == 2

    def test_get_count_with_states(self, client, session, dag_maker):
        """Test counting tasks in specific states."""
        with dag_maker("test_get_count_with_states"):
            EmptyOperator(task_id="task1")
            EmptyOperator(task_id="task2")
            EmptyOperator(task_id="task3")

        dr = dag_maker.create_dagrun()

        tis = dr.get_task_instances()

        # Set different states for the task instances
        for ti, state in zip(tis, [State.SUCCESS, State.FAILED, State.SKIPPED]):
            ti.state = state
            session.merge(ti)
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "test_get_count_with_states", "states": [State.SUCCESS, State.FAILED]},
        )
        assert response.status_code == 200
        assert response.json() == 2

    def test_get_count_with_logical_dates(self, client, session, dag_maker):
        with dag_maker("test_get_count_with_logical_dates"):
            EmptyOperator(task_id="task1")

        date1 = timezone.datetime(2025, 1, 1)
        date2 = timezone.datetime(2025, 1, 2)

        dag_maker.create_dagrun(run_id="test_run_id1", logical_date=date1)
        dag_maker.create_dagrun(run_id="test_run_id2", logical_date=date2)

        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={
                "dag_id": "test_get_count_with_logical_dates",
                "logical_dates": [date1.isoformat(), date2.isoformat()],
            },
        )
        assert response.status_code == 200
        assert response.json() == 2

    def test_get_count_with_run_ids(self, client, session, dag_maker):
        with dag_maker("test_get_count_with_run_ids"):
            EmptyOperator(task_id="task1")

        dag_maker.create_dagrun(run_id="run1", logical_date=timezone.datetime(2025, 1, 1))
        dag_maker.create_dagrun(run_id="run2", logical_date=timezone.datetime(2025, 1, 2))

        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "test_get_count_with_run_ids", "run_ids": ["run1", "run2"]},
        )
        assert response.status_code == 200
        assert response.json() == 2

    def test_get_count_with_task_group(self, client, session, dag_maker):
        with dag_maker(dag_id="test_dag", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")
                EmptyOperator(task_id="task2")

            with TaskGroup("group2"):
                EmptyOperator(task_id="task3")

        dag_maker.create_dagrun(session=session)
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "test_dag", "task_group_id": "group1"},
        )
        assert response.status_code == 200
        assert response.json() == 2

    def test_get_count_task_group_not_found(self, client, session, dag_maker):
        with dag_maker(dag_id="test_get_count_task_group_not_found", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")
        dag_maker.create_dagrun(session=session)

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "test_get_count_task_group_not_found", "task_group_id": "non_existent_group"},
        )
        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task group non_existent_group not found in DAG test_get_count_task_group_not_found",
        }

    def test_get_count_dag_not_found(self, client, session):
        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "non_existent_dag", "task_group_id": "group1"},
        )
        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "The Dag with ID: `non_existent_dag` was not found",
        }

    def test_get_count_with_none_state(self, client, session, create_task_instance):
        create_task_instance(task_id="task1", dag_id="get_count_with_none", state=None)
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "get_count_with_none", "states": ["null"]},
        )
        assert response.status_code == 200
        assert response.json() == 1

    def test_get_count_with_mixed_states(self, client, session, create_task_instance):
        create_task_instance(task_id="task1", state=State.SUCCESS, run_id="runid1", dag_id="mixed_states")
        create_task_instance(task_id="task2", state=None, run_id="runid2", dag_id="mixed_states")
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "mixed_states", "states": [State.SUCCESS, "null"]},
        )
        assert response.status_code == 200
        assert response.json() == 2

    def test_get_count_with_map_index_less_than_zero(self, client, session, create_task_instance):
        create_task_instance(task_id="task1", state=State.SUCCESS, run_id="runid1", dag_id="map_index_test")
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": "map_index_test", "states": [State.SUCCESS], "map_index": -1},
        )
        assert response.status_code == 200
        assert response.json() == 1

    def test_get_count_with_multiple_tasks_and_map_index_less_than_zero(
        self, dag_maker, client, session, create_task_instance
    ):
        with dag_maker("test_get_count_with_multiple_tasks_and_map_index_less_than_zero"):
            EmptyOperator(task_id="task1")
            EmptyOperator(task_id="task2")
            EmptyOperator(task_id="task3")

        dr = dag_maker.create_dagrun()

        tis = dr.get_task_instances()

        # Set different states for the task instances
        for ti, state in zip(tis, [State.SUCCESS, State.FAILED, State.SKIPPED]):
            ti.state = state
            session.merge(ti)
        session.commit()

        response = client.get(
            "/execution/task-instances/count",
            params={
                "dag_id": "test_get_count_with_multiple_tasks_and_map_index_less_than_zero",
                "map_index": -1,
            },
        )
        assert response.status_code == 200
        assert response.json() == 3

    @pytest.mark.parametrize(
        ("map_index", "dynamic_task_args", "expected_count"),
        (
            pytest.param(None, [1, 2, 3], 4, id="use-default-map-index"),
            pytest.param(-1, [1, 2, 3], 1, id="map-index-(-1)"),
            pytest.param(0, [1, 2, 3], 1, id="map-index-0"),
            pytest.param(1, [1, 2, 3], 1, id="map-index-1"),
            pytest.param(2, [1, 2, 3], 1, id="map-index-2"),
        ),
    )
    def test_get_count_for_dynamic_task_mapping(
        self, dag_maker, client, session, map_index, dynamic_task_args, expected_count
    ):
        """
        case 1: map_index is None, it should fetch all the tasks
        other cases: when map index is provided, it should return the count of tasks that falls under the map index
        """
        with dag_maker(session=session) as dag:
            EmptyOperator(task_id="op1")

            @dag.task()
            def add_one(x):
                return [x + 1]

            add_one.expand(x=dynamic_task_args)

        dr = dag_maker.create_dagrun()

        tis = dr.get_task_instances()

        for ti in tis:
            ti.state = State.SUCCESS
            session.merge(ti)
        session.commit()

        map_index = {} if map_index is None else {"map_index": map_index}

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": dr.dag_id, "run_ids": [dr.run_id], **map_index},
        )
        assert response.status_code == 200
        assert response.json() == expected_count

    @pytest.mark.parametrize(
        ("map_index", "dynamic_task_args", "task_ids", "task_group_name", "expected_count"),
        (
            pytest.param(None, [1, 2, 3], None, None, 5, id="use-default-map-index-None"),
            pytest.param(0, [1, 2, 3], None, None, 1, id="with-map-index-0-no-task-group"),
            pytest.param(-1, [1, 2, 3], ["task1"], None, 1, id="with-task-ids-and-map-index-(-1)"),
            pytest.param(None, [1, 2, 3], None, "group1", 4, id="with-task-group-id-and-map-index-None"),
            pytest.param(0, [1, 2, 3], None, "group1", 1, id="with-task-group-id-and-map-index-0"),
            pytest.param(-1, [1, 2, 3], None, "group1", 1, id="with-task-group-id-and-map-index-(-1)"),
        ),
    )
    def test_get_count_mix_of_task_and_task_group_dynamic_task_mapping(
        self,
        dag_maker,
        client,
        session,
        map_index,
        dynamic_task_args,
        task_ids,
        task_group_name,
        expected_count,
    ):
        """
        case 1: map_index is None, task_ids is None, task_group_name is None, it should fetch all the tasks
        case 2: when map index -1 and provided task_ids, it should return the count of task_ids
        case 3: when map index is None and provided task_group_id, it should return the count of tasks under the task group
        case 4: when map index is 0 and provided task_group_id, it should return the count of tasks under the task group that falls map index =0
        case 5: when map index is -1 and provided task_group_id, it should return the count of tasks under the task group that falls map index =-1 i.e this task is not mapped
        """

        with dag_maker(session=session, serialized=True) as dag:
            EmptyOperator(task_id="task1")

            with TaskGroup("group1"):

                @dag.task()
                def add_one(x):
                    return [x + 1]

                add_one.expand(x=dynamic_task_args)

                EmptyOperator(task_id="task2")

        dr = dag_maker.create_dagrun(session=session)

        session.commit()
        params = {}

        if task_ids:
            params["task_ids"] = task_ids
        if task_group_name:
            params["task_group_id"] = task_group_name
        if map_index is not None:
            params["map_index"] = map_index

        response = client.get(
            "/execution/task-instances/count",
            params={"dag_id": dr.dag_id, "run_ids": [dr.run_id], **params},
        )
        assert response.status_code == 200
        assert response.json() == expected_count


class TestGetPreviousTI:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    def test_get_previous_ti_basic(self, client, session, create_task_instance):
        """Test basic get_previous_ti without filters."""
        # Create TIs with different logical dates
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 1),
            run_id="run1",
        )
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 2),
            run_id="run2",
        )
        session.commit()

        response = client.get(
            "/execution/task-instances/previous/dag/test_task",
            params={
                "logical_date": "2025-01-02T00:00:00Z",
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert data["task_id"] == "test_task"
        assert data["dag_id"] == "dag"
        assert data["run_id"] == "run1"
        assert data["state"] == State.SUCCESS

    def test_get_previous_ti_with_state_filter(self, client, session, create_task_instance):
        """Test get_previous_ti with state filter."""
        # Create TIs with different states
        create_task_instance(
            task_id="test_task",
            state=State.FAILED,
            logical_date=timezone.datetime(2025, 1, 1),
            run_id="run1",
        )
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 2),
            run_id="run2",
        )
        create_task_instance(
            task_id="test_task",
            state=State.FAILED,
            logical_date=timezone.datetime(2025, 1, 3),
            run_id="run3",
        )
        session.commit()

        # Query for previous successful TI before 2025-01-03
        response = client.get(
            "/execution/task-instances/previous/dag/test_task",
            params={
                "logical_date": "2025-01-03T00:00:00Z",
                "state": State.SUCCESS,
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert data["task_id"] == "test_task"
        assert data["run_id"] == "run2"
        assert data["state"] == State.SUCCESS

    def test_get_previous_ti_with_map_index_filter(self, client, session, create_task_instance):
        """Test get_previous_ti with map_index filter for mapped tasks."""
        # Create TIs with different map_index values
        # map_index=0 in run1
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 1),
            run_id="run1",
            map_index=0,
        )
        # map_index=1 in run1_alt (different logical date to avoid constraint)
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 1, 12, 0, 0),
            run_id="run1_alt",
            map_index=1,
        )
        # map_index=0 in run2
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 2),
            run_id="run2",
            map_index=0,
        )
        session.commit()

        # Query for previous TI with map_index=0 before 2025-01-02
        response = client.get(
            "/execution/task-instances/previous/dag/test_task",
            params={
                "logical_date": "2025-01-02T00:00:00Z",
                "map_index": 0,
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert data["run_id"] == "run1"
        assert data["map_index"] == 0

    def test_get_previous_ti_not_found(self, client, session):
        """Test get_previous_ti when no previous TI exists."""
        response = client.get(
            "/execution/task-instances/previous/dag/test_task",
        )
        assert response.status_code == 200
        assert response.json() is None

    def test_get_previous_ti_returns_most_recent(self, client, session, create_task_instance):
        """Test that get_previous_ti returns the most recent matching TI."""
        # Create multiple TIs
        for i in range(5):
            create_task_instance(
                task_id="test_task",
                state=State.SUCCESS,
                logical_date=timezone.datetime(2025, 1, i + 1),
                run_id=f"run{i + 1}",
            )
        session.commit()

        # Query for TI before 2025-01-05
        response = client.get(
            "/execution/task-instances/previous/dag/test_task",
            params={
                "logical_date": "2025-01-05T00:00:00Z",
            },
        )
        assert response.status_code == 200
        data = response.json()
        # Should return the most recent one (2025-01-04)
        assert data["run_id"] == "run4"

    def test_get_previous_ti_with_all_filters(self, client, session, create_task_instance):
        """Test get_previous_ti with all filters combined."""
        # Create TIs with different states and map_index values
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 1),
            run_id="target_run_1",
            map_index=0,
        )
        create_task_instance(
            task_id="test_task",
            state=State.FAILED,
            logical_date=timezone.datetime(2025, 1, 2),
            run_id="target_run_2",
            map_index=0,
        )
        create_task_instance(
            task_id="test_task",
            state=State.SUCCESS,
            logical_date=timezone.datetime(2025, 1, 3),
            run_id="target_run_3",
            map_index=1,
        )
        session.commit()

        # Query for previous successful TI before 2025-01-03 with map_index=0
        response = client.get(
            "/execution/task-instances/previous/dag/test_task",
            params={
                "logical_date": "2025-01-03T00:00:00Z",
                "state": State.SUCCESS,
                "map_index": 0,
            },
        )
        assert response.status_code == 200
        data = response.json()
        assert data["run_id"] == "target_run_1"
        assert data["state"] == State.SUCCESS


class TestGetTaskStates:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    def test_get_task_states_basic(self, client, session, create_task_instance):
        create_task_instance(task_id="test_task", state=State.SUCCESS)
        session.commit()

        response = client.get("/execution/task-instances/states", params={"dag_id": "dag"})
        assert response.status_code == 200
        assert response.json() == {"task_states": {"test": {"test_task": "success"}}}

    def test_get_task_states_group_id_basic(self, client, dag_maker, session):
        with dag_maker(dag_id="test_dag", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")

        dag_maker.create_dagrun(session=session)
        session.commit()

        response = client.get(
            "/execution/task-instances/states",
            params={"dag_id": "test_dag", "task_group_id": "group1"},
        )
        assert response.status_code == 200
        assert response.json() == {
            "task_states": {
                "test": {
                    "group1.task1": None,
                },
            },
        }

    def test_get_task_states_with_task_group_id_and_task_id(self, client, session, dag_maker):
        with dag_maker("test_get_task_group_states_with_multiple_task_tasks", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")
            EmptyOperator(task_id="task2")

        dr = dag_maker.create_dagrun()

        tis = dr.get_task_instances()

        # Set different states for the task instances
        for ti, state in zip(tis, [State.SUCCESS, State.FAILED]):
            ti.state = state
            session.merge(ti)
        session.commit()

        response = client.get(
            "/execution/task-instances/states",
            params={
                "dag_id": "test_get_task_group_states_with_multiple_task_tasks",
                "task_group_id": "group1",
            },
        )
        assert response.status_code == 200
        assert response.json() == {
            "task_states": {
                "test": {
                    "group1.task1": "success",
                },
            },
        }

    def test_get_task_group_states_with_multiple_task(self, client, session, dag_maker):
        with dag_maker("test_get_task_group_states_with_multiple_task_tasks", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")
                EmptyOperator(task_id="task2")
                EmptyOperator(task_id="task3")

        dr = dag_maker.create_dagrun()

        tis = dr.get_task_instances()

        # Set different states for the task instances
        for ti, state in zip(tis, [State.SUCCESS, State.FAILED, State.SKIPPED]):
            ti.state = state
            session.merge(ti)
        session.commit()

        response = client.get(
            "/execution/task-instances/states",
            params={
                "dag_id": "test_get_task_group_states_with_multiple_task_tasks",
                "task_group_id": "group1",
            },
        )
        assert response.status_code == 200
        assert response.json() == {
            "task_states": {
                "test": {
                    "group1.task1": "success",
                    "group1.task2": "failed",
                    "group1.task3": "skipped",
                },
            },
        }

    def test_get_task_group_states_with_logical_dates(self, client, session, dag_maker):
        with dag_maker("test_get_task_group_states_with_logical_dates", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")

        date1 = timezone.datetime(2025, 1, 1)
        date2 = timezone.datetime(2025, 1, 2)

        dag_maker.create_dagrun(run_id="test_run_id1", logical_date=date1)
        dag_maker.create_dagrun(run_id="test_run_id2", logical_date=date2)

        session.commit()

        response = client.get(
            "/execution/task-instances/states",
            params={
                "dag_id": "test_get_task_group_states_with_logical_dates",
                "logical_dates": [date1.isoformat(), date2.isoformat()],
                "task_group_id": "group1",
            },
        )
        assert response.status_code == 200
        assert response.json() == {
            "task_states": {
                "test_run_id1": {
                    "group1.task1": None,
                },
                "test_run_id2": {
                    "group1.task1": None,
                },
            },
        }

    def test_get_task_group_states_with_run_ids(self, client, session, dag_maker):
        with dag_maker("test_get_task_group_states_with_run_ids", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")

        dag_maker.create_dagrun(run_id="run1", logical_date=timezone.datetime(2025, 1, 1))
        dag_maker.create_dagrun(run_id="run2", logical_date=timezone.datetime(2025, 1, 2))

        session.commit()

        response = client.get(
            "/execution/task-instances/states",
            params={
                "dag_id": "test_get_task_group_states_with_run_ids",
                "run_ids": ["run1", "run2"],
                "task_group_id": "group1",
            },
        )
        assert response.status_code == 200
        assert response.json() == {
            "task_states": {
                "run1": {
                    "group1.task1": None,
                },
                "run2": {
                    "group1.task1": None,
                },
            },
        }

    def test_get_task_states_task_group_not_found(self, client, session, dag_maker):
        with dag_maker(dag_id="test_get_task_states_task_group_not_found", serialized=True):
            with TaskGroup("group1"):
                EmptyOperator(task_id="task1")
        dag_maker.create_dagrun(session=session)

        response = client.get(
            "/execution/task-instances/states",
            params={
                "dag_id": "test_get_task_states_task_group_not_found",
                "task_group_id": "non_existent_group",
            },
        )
        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "Task group non_existent_group not found in DAG test_get_task_states_task_group_not_found",
        }

    def test_get_task_states_dag_not_found(self, client, session):
        response = client.get(
            "/execution/task-instances/states",
            params={"dag_id": "non_existent_dag", "task_group_id": "group1"},
        )
        assert response.status_code == 404
        assert response.json()["detail"] == {
            "reason": "not_found",
            "message": "The Dag with ID: `non_existent_dag` was not found",
        }

    @pytest.mark.parametrize(
        ("map_index", "dynamic_task_args", "states", "expected"),
        (
            pytest.param(
                None,
                [1, 2, 3],
                {"-1": State.SUCCESS, "0": State.SUCCESS, "1": State.SUCCESS, "2": State.SUCCESS},
                {"task1": "success", "add_one_0": "success", "add_one_1": "success", "add_one_2": "success"},
                id="with-default-map-index-None",
            ),
            pytest.param(
                0,
                [1, 2, 3],
                {"-1": State.SUCCESS, "0": State.FAILED, "1": State.SUCCESS, "2": State.SUCCESS},
                {"add_one_0": "failed"},
                id="with-map-index-0",
            ),
            pytest.param(
                1,
                [1, 2, 3],
                {"-1": State.SUCCESS, "0": State.SUCCESS, "1": State.FAILED, "2": State.SUCCESS},
                {"add_one_1": "failed"},
                id="with-map-index-1",
            ),
        ),
    )
    def test_get_task_states_for_dynamic_task_mapping(
        self, dag_maker, client, session, map_index, dynamic_task_args, states, expected
    ):
        """
        case 1: map_index is None, it should fetch all the tasks
        other cases: when map index is provided, it should return the count of tasks that falls under the map index
        """
        with dag_maker(session=session, serialized=True) as dag:
            EmptyOperator(task_id="task1")

            @dag.task()
            def add_one(x):
                return [x + 1]

            add_one.expand(x=dynamic_task_args)

        dr = dag_maker.create_dagrun()

        tis = dr.get_task_instances()
        for ti in tis:
            ti.state = states.get(str(ti.map_index))
            session.merge(ti)
        session.commit()

        map_index = {} if map_index is None else {"map_index": map_index}

        response = client.get("/execution/task-instances/states", params={"dag_id": dr.dag_id, **map_index})
        assert response.status_code == 200
        assert response.json() == {"task_states": {dr.run_id: expected}}

    @pytest.mark.parametrize(
        ("map_index", "dynamic_task_args", "task_ids", "task_group_name", "states", "expected"),
        (
            pytest.param(
                None,
                [1, 2, 3],
                None,
                None,
                {"-1": State.SUCCESS, "0": State.SUCCESS, "1": State.SUCCESS, "2": State.SUCCESS},
                {
                    "group1.add_one_0": "success",
                    "group1.add_one_1": "success",
                    "group1.add_one_2": "success",
                    "group1.task2": "success",
                    "task1": "success",
                },
                id="with-default-map-index-None",
            ),
            pytest.param(
                0,
                [1, 2, 3],
                None,
                None,
                {"-1": State.SUCCESS, "0": State.FAILED, "1": State.SUCCESS, "2": State.SUCCESS},
                {"group1.add_one_0": "failed"},
                id="with-map-index-0-no-task-group",
            ),
            pytest.param(
                -1,
                [1, 2, 3],
                ["task1"],
                None,
                {"-1": State.SUCCESS, "0": State.SUCCESS, "1": State.SUCCESS, "2": State.SUCCESS},
                {"task1": "success"},
                id="with-task-ids-map-index-(-1)",
            ),
            pytest.param(
                None,
                [1, 2, 3],
                None,
                "group1",
                {"-1": State.SUCCESS, "0": State.SUCCESS, "1": State.SUCCESS, "2": State.SUCCESS},
                {
                    "group1.task2": "success",
                    "group1.add_one_0": "success",
                    "group1.add_one_1": "success",
                    "group1.add_one_2": "success",
                },
                id="with-task-group-id-and-map-index-None",
            ),
            pytest.param(
                0,
                [1, 2, 3],
                None,
                "group1",
                {"-1": State.SUCCESS, "0": State.FAILED, "1": State.SUCCESS, "2": State.SUCCESS},
                {"group1.add_one_0": "failed"},
                id="with-task-group-id-and-map-index-0",
            ),
            pytest.param(
                -1,
                [1, 2, 3],
                ["task1"],
                "group1",
                {"-1": State.SUCCESS, "0": State.SUCCESS, "1": State.SUCCESS, "2": State.SUCCESS},
                {"task1": "success", "group1.task2": "success"},
                id="with-task-id-and-task-group-map-index-(-1)",
            ),
        ),
    )
    def test_get_task_states_mix_of_task_and_task_group_dynamic_task_mapping(
        self,
        dag_maker,
        client,
        session,
        map_index,
        dynamic_task_args,
        task_ids,
        task_group_name,
        states,
        expected,
    ):
        """
        case1: map_index is None, task_ids is None, task_group_name is None, it should fetch all the task states
        case2: when map index -1 and provided task_ids, it should return the task states of task_ids
        case3: when map index is None and provided task_group_id, it should return the task states of tasks under the task group and normal task states under task group
        case4: when map index is 0 and provided task_group_id, it should return the task states of tasks under the task group that falls under map index = 0
        case5: when map index is -1 and provided both task_id and task_group_id, it should return the task states of tasks under the task group that falls under map index = -1 and normal task_ids states
        """

        with dag_maker(session=session, serialized=True) as dag:
            EmptyOperator(task_id="task1")

            with TaskGroup("group1"):

                @dag.task()
                def add_one(x):
                    return [x + 1]

                add_one.expand(x=dynamic_task_args)

                EmptyOperator(task_id="task2")

        dr = dag_maker.create_dagrun(session=session)

        tis = dr.get_task_instances()
        for ti in tis:
            ti.state = states.get(str(ti.map_index))
            session.merge(ti)
        session.commit()
        params = {}

        if task_ids:
            params["task_ids"] = task_ids
        if task_group_name:
            params["task_group_id"] = task_group_name
        if map_index is not None:
            params["map_index"] = map_index

        response = client.get("/execution/task-instances/states", params={"dag_id": dr.dag_id, **params})
        assert response.status_code == 200
        assert response.json() == {"task_states": {dr.run_id: expected}}


class TestGetTaskInstanceBreadcrumbs:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    @pytest.fixture(autouse=True)
    def dag_run(self, dag_maker, session):
        with dag_maker(session=session):
            for name in TaskInstanceState._member_names_:
                EmptyOperator(task_id=name)
        return dag_maker.create_dagrun(state="running")

    @pytest.fixture(autouse=True)
    def task_instances(self, dag_run, session):
        tis = {ti.task_id: ti for ti in dag_run.task_instances}
        for name, value in TaskInstanceState._member_map_.items():
            tis[name].state = value
        session.commit()
        return tis

    def test_get_breadcrumbs(self, client, dag_run):
        response = client.get(
            "/execution/task-instances/breadcrumbs",
            params={"dag_id": dag_run.dag_id, "run_id": dag_run.run_id},
        )
        assert response.status_code == 200
        assert response.json() == {  # Should find tis with terminal states.
            "breadcrumbs": [
                {
                    "duration": None,
                    "map_index": -1,
                    "operator": "EmptyOperator",
                    "state": "failed",
                    "task_id": "FAILED",
                },
                {
                    "duration": None,
                    "map_index": -1,
                    "operator": "EmptyOperator",
                    "state": "removed",
                    "task_id": "REMOVED",
                },
                {
                    "duration": None,
                    "map_index": -1,
                    "operator": "EmptyOperator",
                    "state": "skipped",
                    "task_id": "SKIPPED",
                },
                {
                    "duration": None,
                    "map_index": -1,
                    "operator": "EmptyOperator",
                    "state": "success",
                    "task_id": "SUCCESS",
                },
                {
                    "duration": None,
                    "map_index": -1,
                    "operator": "EmptyOperator",
                    "state": "upstream_failed",
                    "task_id": "UPSTREAM_FAILED",
                },
            ]
        }


class TestInvactiveInletsAndOutlets:
    @pytest.mark.parametrize(
        "logical_date",
        [
            datetime(2025, 6, 6, tzinfo=timezone.utc),
            None,
        ],
    )
    def test_ti_inactive_inlets_and_outlets(self, logical_date, client, dag_maker):
        """Test the inactive assets in inlets and outlets can be found."""
        with dag_maker("test_inlets_and_outlets"):
            EmptyOperator(
                task_id="task1",
                inlets=[Asset(name="inlet-name"), Asset(name="inlet-name", uri="but-different-uri")],
                outlets=[
                    Asset(name="outlet-name", uri="uri"),
                    Asset(name="outlet-name", uri="second-different-uri"),
                ],
            )

        dr = dag_maker.create_dagrun(logical_date=logical_date)

        task1_ti = dr.get_task_instance("task1")
        response = client.get(f"/execution/task-instances/{task1_ti.id}/validate-inlets-and-outlets")
        assert response.status_code == 200
        inactive_assets = response.json()["inactive_assets"]
        expected_inactive_assets = (
            {
                "name": "inlet-name",
                "type": "Asset",
                "uri": "but-different-uri",
            },
            {
                "name": "outlet-name",
                "type": "Asset",
                "uri": "second-different-uri",
            },
        )
        for asset in expected_inactive_assets:
            assert asset in inactive_assets

    @pytest.mark.parametrize(
        "logical_date",
        [
            datetime(2025, 6, 6, tzinfo=timezone.utc),
            None,
        ],
    )
    def test_ti_inactive_inlets_and_outlets_without_inactive_assets(self, logical_date, client, dag_maker):
        """Test the task without inactive assets in its inlets or outlets returns empty list."""
        with dag_maker("test_inlets_and_outlets_inactive"):
            EmptyOperator(
                task_id="inactive_task1",
                inlets=[Asset(name="inlet-name")],
                outlets=[Asset(name="outlet-name", uri="uri")],
            )

        dr = dag_maker.create_dagrun(logical_date=logical_date)

        task1_ti = dr.get_task_instance("inactive_task1")
        response = client.get(f"/execution/task-instances/{task1_ti.id}/validate-inlets-and-outlets")
        assert response.status_code == 200
        assert response.json() == {"inactive_assets": []}

    def test_ti_run_with_null_conf(self, client, session, create_task_instance):
        """Test that task instances can start when dag_run.conf is NULL."""
        ti = create_task_instance(
            task_id="test_ti_run_with_null_conf",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
        )
        # Set conf to NULL to simulate Airflow 2.x upgrade or offline migration
        ti.dag_run.conf = None
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "pid": 100,
                "hostname": "test-hostname",
                "unixname": "test-user",
                "start_date": timezone.utcnow().isoformat(),
            },
        )

        assert response.status_code == 200, f"Response: {response.text}"
        context = response.json()
        assert context["dag_run"]["conf"] is None


class TestTIPatchRenderedMapIndex:
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    def test_ti_patch_rendered_map_index(self, client, session, create_task_instance):
        """Test updating rendered_map_index for a task instance."""
        ti = create_task_instance(
            task_id="test_ti_patch_rendered_map_index",
            state=State.RUNNING,
            session=session,
        )
        session.commit()

        rendered_map_index = "custom_label_123"
        response = client.patch(
            f"/execution/task-instances/{ti.id}/rendered-map-index",
            json=rendered_map_index,
        )

        assert response.status_code == 204
        assert response.text == ""

        session.expire_all()
        ti = session.get(TaskInstance, ti.id)
        assert ti.rendered_map_index == rendered_map_index

    def test_ti_patch_rendered_map_index_not_found(self, client, session):
        """Test 404 error when task instance does not exist."""
        fake_id = str(uuid4())
        response = client.patch(
            f"/execution/task-instances/{fake_id}/rendered-map-index",
            json="test",
        )

        assert response.status_code == 404

    def test_ti_patch_rendered_map_index_empty_string(self, client, session, create_task_instance):
        """Test that empty string is accepted (clears the rendered_map_index)."""
        ti = create_task_instance(
            task_id="test_ti_patch_rendered_map_index_empty",
            state=State.RUNNING,
            session=session,
        )
        session.commit()

        response = client.patch(
            f"/execution/task-instances/{ti.id}/rendered-map-index",
            json="",
        )

        assert response.status_code == 422


@pytest.mark.usefixtures("_use_real_jwt_bearer")
class TestTokenTypeValidation:
    """Test token scope enforcement (workload vs execution)."""

    def _register_scoped_validator(self, ti_id, scope):
        """Register a JWTValidator mock returning claims with the given scope."""
        validator = mock.AsyncMock(spec=JWTValidator)
        claims = {"sub": str(ti_id), "exp": 9999999999, "iat": 1000000000, "nbf": 1000000000}
        if scope is not None:
            claims["scope"] = scope
        validator.avalidated_claims.side_effect = lambda cred, validators: claims
        lifespan.registry.register_value(JWTValidator, validator)

    def test_workload_scope_rejected_on_heartbeat_endpoint(self, client, session, create_task_instance):
        """Workload scoped tokens should be rejected on /heartbeat."""
        ti = create_task_instance(task_id="test_ti_run_heartbeat", state=State.RUNNING)
        session.commit()

        self._register_scoped_validator(ti.id, "workload")

        payload = {"hostname": "test-host", "pid": 100}
        resp = client.put(f"/execution/task-instances/{ti.id}/heartbeat", json=payload)
        assert resp.status_code == 403
        assert "Token type 'workload' not allowed" in resp.json()["detail"]

    def test_workload_scope_rejected_on_state_endpoint(self, client, session, create_task_instance):
        """Workload scoped tokens should be rejected on PATCH /state."""
        ti = create_task_instance(task_id="test_workload_state", state=State.RUNNING)
        session.commit()

        self._register_scoped_validator(ti.id, "workload")

        payload = {"state": "success", "end_date": "2024-10-31T13:00:00Z"}
        resp = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
        assert resp.status_code == 403
        assert "Token type 'workload' not allowed" in resp.json()["detail"]

    def test_workload_scope_rejected_on_connections_endpoint(self, client, session, create_task_instance):
        """Workload scoped tokens should be rejected on GET /connections (different router)."""
        ti = create_task_instance(task_id="test_workload_conn", state=State.RUNNING)
        session.commit()

        self._register_scoped_validator(ti.id, "workload")

        resp = client.get("/execution/connections/test_conn")
        assert resp.status_code == 403
        assert "Token type 'workload' not allowed" in resp.json()["detail"]

    def test_execution_scope_accepted_on_all_endpoints(self, client, session, create_task_instance):
        """Execution scoped tokens should be accepted on all endpoints."""
        ti = create_task_instance(task_id="test_ti_star", state=State.RUNNING)
        session.commit()

        self._register_scoped_validator(ti.id, "execution")

        payload = {"state": "success", "end_date": "2024-10-31T13:00:00Z"}
        resp = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
        assert resp.status_code in [200, 204]

    def test_invalid_scope_value_rejected(self, client, session, create_task_instance):
        """Tokens with unrecognized scope values should be rejected."""
        ti = create_task_instance(task_id="test_invalid_scope", state=State.QUEUED)
        session.commit()

        self._register_scoped_validator(ti.id, "bogus:scope")

        payload = {
            "state": "running",
            "hostname": "test-host",
            "unixname": "test-user",
            "pid": 100,
            "start_date": "2024-10-31T12:00:00Z",
        }

        resp = client.patch(f"/execution/task-instances/{ti.id}/run", json=payload)
        assert resp.status_code == 403
        assert "Invalid auth token" in resp.json()["detail"]

    def test_workload_scope_accepted_on_run_endpoint(
        self, client, session, create_task_instance, time_machine
    ):
        """Workload scoped tokens should be accepted on the /run endpoint."""
        instant = timezone.parse("2024-10-31T12:00:00Z")
        time_machine.move_to(instant, tick=False)

        ti = create_task_instance(
            task_id="test_workload_run",
            state=State.QUEUED,
            dagrun_state=DagRunState.RUNNING,
            session=session,
            start_date=instant,
            dag_id=str(uuid4()),
        )
        session.commit()

        self._register_scoped_validator(ti.id, "workload")

        resp = client.patch(
            f"/execution/task-instances/{ti.id}/run",
            json={
                "state": "running",
                "hostname": "test-host",
                "unixname": "test-user",
                "pid": 100,
                "start_date": "2024-10-31T12:00:00Z",
            },
        )
        assert resp.status_code == 200

    def test_no_scope_defaults_to_execution(self, client, session, create_task_instance):
        """Tokens without scope claim should default to 'execution'."""
        ti = create_task_instance(task_id="test_no_scope", state=State.RUNNING)
        session.commit()

        self._register_scoped_validator(ti.id, None)

        payload = {"state": "success", "end_date": "2024-10-31T13:00:00Z"}
        resp = client.patch(f"/execution/task-instances/{ti.id}/state", json=payload)
        assert resp.status_code in [200, 204]


class TestEmitTaskSpan:
    """Tests for the _emit_task_span function in the execution API task-instance route."""

    @pytest.fixture(autouse=True)
    def sdk_tracer_provider(self):
        self.exporter = InMemorySpanExporter()
        provider = TracerProvider(id_generator=OverrideableRandomIdGenerator())
        provider.add_span_processor(SimpleSpanProcessor(self.exporter))
        test_tracer = provider.get_tracer("test")
        with mock.patch("airflow.api_fastapi.execution_api.routes.task_instances.tracer", test_tracer):
            yield

    def _make_carriers(self):
        """Return a (dr_carrier, ti_carrier) pair built with a real SDK provider."""
        p = TracerProvider()
        t = p.get_tracer("setup")
        dr_span = t.start_span("dr")
        dr_ctx = otel_trace.set_span_in_context(dr_span)
        dr_carrier: dict = {}
        TraceContextTextMapPropagator().inject(dr_carrier, context=dr_ctx)
        ti_span = t.start_span("ti", context=dr_ctx)
        ti_ctx = otel_trace.set_span_in_context(ti_span)
        ti_carrier: dict = {}
        TraceContextTextMapPropagator().inject(ti_carrier, context=ti_ctx)
        return dr_carrier, ti_carrier

    def _make_ti(self, task_id="my_task", map_index=-1, queued_dttm=None, start_date=None):
        dr_carrier, ti_carrier = self._make_carriers()
        ti = mock.MagicMock()
        ti.id = UUID("0182e924-0f1e-77e6-ab50-e977118bc139")
        ti.dag_id = "test_dag"
        ti.task_id = task_id
        ti.run_id = "test_run"
        ti.try_number = 1
        ti.map_index = map_index
        ti.queued_dttm = queued_dttm
        ti.start_date = start_date or DEFAULT_START_DATE
        ti.dag_run.context_carrier = dr_carrier
        ti.context_carrier = ti_carrier
        return ti

    def test_emit_task_span_success_sets_ok_status(self):
        _emit_task_span(self._make_ti(), TaskInstanceState.SUCCESS)

        spans = self.exporter.get_finished_spans()
        assert len(spans) == 1
        assert spans[0].status.status_code == StatusCode.OK

    def test_emit_task_span_failed_sets_error_status(self):
        _emit_task_span(self._make_ti(), TaskInstanceState.FAILED)

        spans = self.exporter.get_finished_spans()
        assert len(spans) == 1
        assert spans[0].status.status_code == StatusCode.ERROR

    def test_emit_task_span_sets_attributes(self):
        ti = self._make_ti(task_id="my_task", map_index=2)
        _emit_task_span(ti, TaskInstanceState.SUCCESS)

        attrs = self.exporter.get_finished_spans()[0].attributes
        assert attrs["airflow.dag_id"] == "test_dag"
        assert attrs["airflow.task_id"] == "my_task"
        assert attrs["airflow.dag_run.run_id"] == "test_run"
        assert attrs["airflow.task_instance.try_number"] == 1
        assert attrs["airflow.task_instance.map_index"] == 2
        assert attrs["airflow.task_instance.state"] == TaskInstanceState.SUCCESS
        # The OTEL SDK only accepts str/bytes/int/float/bool attribute values,
        # so the UUID must be stringified before being set on the span.
        assert attrs["airflow.task_instance.id"] == str(ti.id)
        assert isinstance(attrs["airflow.task_instance.id"], str)

    def test_emit_task_span_name_unmapped(self):
        _emit_task_span(self._make_ti(task_id="my_task", map_index=-1), TaskInstanceState.SUCCESS)
        assert self.exporter.get_finished_spans()[0].name == "task_run.my_task"

    def test_emit_task_span_name_mapped(self):
        _emit_task_span(self._make_ti(task_id="my_task", map_index=3), TaskInstanceState.SUCCESS)
        assert self.exporter.get_finished_spans()[0].name == "task_run.my_task[3]"

    def test_emit_task_span_start_time_uses_queued_dttm(self):
        queued_dttm = timezone.parse("2024-01-01T10:00:00Z")
        start_date = timezone.parse("2024-01-01T10:05:00Z")
        ti = self._make_ti(queued_dttm=queued_dttm, start_date=start_date)
        _emit_task_span(ti, TaskInstanceState.SUCCESS)

        assert self.exporter.get_finished_spans()[0].start_time == int(queued_dttm.timestamp() * 1e9)

    def test_emit_task_span_start_time_falls_back_to_start_date(self):
        start_date = timezone.parse("2024-01-01T10:05:00Z")
        ti = self._make_ti(queued_dttm=None, start_date=start_date)
        _emit_task_span(ti, TaskInstanceState.SUCCESS)

        assert self.exporter.get_finished_spans()[0].start_time == int(start_date.timestamp() * 1e9)

    def test_emit_task_span_skips_if_no_ti_carrier(self):
        ti = mock.MagicMock()
        ti.dag_run.context_carrier = {
            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
        }
        ti.context_carrier = None

        _emit_task_span(ti, TaskInstanceState.SUCCESS)
        assert len(self.exporter.get_finished_spans()) == 0

    def test_emit_task_span_skips_if_no_dagrun_carrier(self):
        ti = mock.MagicMock()
        ti.dag_run.context_carrier = None
        ti.context_carrier = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}

        _emit_task_span(ti, TaskInstanceState.SUCCESS)
        assert len(self.exporter.get_finished_spans()) == 0

    @pytest.mark.parametrize(
        ("trace_flag", "expected_spans"),
        [
            pytest.param("01", 1, id="sampled-carrier-emits"),
            pytest.param("00", 0, id="unsampled-carrier-skips"),
        ],
    )
    def test_emit_task_span_honors_dagrun_carrier_sampling(self, trace_flag, expected_spans):
        """A SAMPLED dag_run carrier (flag 01) emits the task span; an unsampled one (flag 00) is head-sampled out."""
        ti = self._make_ti()
        ti.dag_run.context_carrier = {
            "traceparent": f"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-{trace_flag}"
        }
        _emit_task_span(ti, TaskInstanceState.SUCCESS)
        assert len(self.exporter.get_finished_spans()) == expected_spans
