fix(hitl): handle hitl details when task instance is retried (apache#53824)

Author: Lee-W

airflow-core/src/airflow/api_fastapi/execution_api/routes/hitl.py

@@ -40,32 +40,49 @@
     "/{task_instance_id}",
     status_code=status.HTTP_201_CREATED,
 )
-def add_hitl_detail(
+def upsert_hitl_detail(
     task_instance_id: UUID,
     payload: HITLDetailRequest,
     session: SessionDep,
 ) -> HITLDetailRequest:
-    """Get Human-in-the-loop detail for a specific Task Instance."""
+    """
+    Create a Human-in-the-loop detail for a specific Task Instance.
+
+    There're 3 cases handled here.
+
+    1. If a HITLOperator task instance does not have a HITLDetail,
+       a new HITLDetail is created without a response section.
+    2. If a HITLOperator task instance has a HITLDetail but lacks a response,
+       the existing HITLDetail is returned.
+       This situation occurs when a task instance is cleared before a response is received.
+    3. If a HITLOperator task instance has both a HITLDetail and a response section,
+       the existing response is removed, and the HITLDetail is returned.
+       This happens when a task instance is cleared after a response has been received.
+       This design ensures that each task instance has only one HITLDetail.
+    """
     ti_id_str = str(task_instance_id)
     hitl_detail_model = session.scalar(select(HITLDetail).where(HITLDetail.ti_id == ti_id_str))
-    if hitl_detail_model:
-        raise HTTPException(
-            status.HTTP_409_CONFLICT,
-            f"Human-in-the-loop detail for Task Instance with id {ti_id_str} already exists.",
+    if not hitl_detail_model:
+        hitl_detail_model = HITLDetail(
+            ti_id=ti_id_str,
+            options=payload.options,
+            subject=payload.subject,
+            body=payload.body,
+            defaults=payload.defaults,
+            multiple=payload.multiple,
+            params=payload.params,
         )
+        session.add(hitl_detail_model)
+    elif hitl_detail_model.response_received:
+        # Cleanup the response part of HITLDetail as we only store one response for one task instance.
+        # It normally happens after retry, we keep only the latest response.
+        hitl_detail_model.user_id = None
+        hitl_detail_model.response_at = None
+        hitl_detail_model.chosen_options = None
+        hitl_detail_model.params_input = {}
+        session.add(hitl_detail_model)
 
-    hitl_detail = HITLDetail(
-        ti_id=ti_id_str,
-        options=payload.options,
-        subject=payload.subject,
-        body=payload.body,
-        defaults=payload.defaults,
-        multiple=payload.multiple,
-        params=payload.params,
-    )
-    session.add(hitl_detail)
-    session.commit()
-    return HITLDetailRequest.model_validate(hitl_detail)
+    return HITLDetailRequest.model_validate(hitl_detail_model)
 
 
 @router.patch("/{task_instance_id}")

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py

@@ -16,43 +16,60 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
-
 import pytest
-import time_machine
-from uuid6 import uuid7
+from httpx import Client
 
 from tests_common.test_utils.db import AIRFLOW_V_3_1_PLUS
 
 if not AIRFLOW_V_3_1_PLUS:
     pytest.skip("Human in the loop public API compatible with Airflow >= 3.0.1", allow_module_level=True)
 
+from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
+import time_machine
+
+from airflow._shared.timezones.timezone import convert_to_utc
 from airflow.models.hitl import HITLDetail
 
 if TYPE_CHECKING:
+    from fastapi.testclient import TestClient
+    from sqlalchemy.orm import Session
+
     from airflow.models.taskinstance import TaskInstance
 
+    from tests_common.pytest_plugin import CreateTaskInstance
+
 pytestmark = pytest.mark.db_test
-TI_ID = uuid7()
+
+default_hitl_detail_request_kwargs: dict[str, Any] = {
+    # ti_id decided at a later stage
+    "subject": "This is subject",
+    "body": "this is body",
+    "options": ["Approve", "Reject"],
+    "defaults": ["Approve"],
+    "multiple": False,
+    "params": {"input_1": 1},
+}
+expected_empty_hitl_detail_response_part: dict[str, Any] = {
+    "response_at": None,
+    "chosen_options": None,
+    "user_id": None,
+    "params_input": {},
+    "response_received": False,
+}
 
 
 @pytest.fixture
-def sample_ti(create_task_instance) -> TaskInstance:
+def sample_ti(create_task_instance: CreateTaskInstance) -> TaskInstance:
     return create_task_instance()
 
 
 @pytest.fixture
-def sample_hitl_detail(session, sample_ti) -> HITLDetail:
+def sample_hitl_detail(session: Session, sample_ti: TaskInstance) -> HITLDetail:
     hitl_detail_model = HITLDetail(
         ti_id=sample_ti.id,
-        options=["Approve", "Reject"],
-        subject="This is subject",
-        body="this is body",
-        defaults=["Approve"],
-        multiple=False,
-        params={"input_1": 1},
+        **default_hitl_detail_request_kwargs,
     )
     session.add(hitl_detail_model)
     session.commit()
@@ -61,54 +78,65 @@ def sample_hitl_detail(session, sample_ti) -> HITLDetail:
 
 
 @pytest.fixture
-def expected_sample_hitl_detail_dict(sample_ti) -> dict[str, Any]:
+def expected_sample_hitl_detail_dict(sample_ti: TaskInstance) -> dict[str, Any]:
     return {
-        "body": "this is body",
-        "defaults": ["Approve"],
-        "multiple": False,
-        "options": ["Approve", "Reject"],
-        "params": {"input_1": 1},
-        "params_input": {},
-        "response_at": None,
-        "chosen_options": None,
-        "response_received": False,
-        "subject": "This is subject",
         "ti_id": sample_ti.id,
-        "user_id": None,
+        **default_hitl_detail_request_kwargs,
+        **expected_empty_hitl_detail_response_part,
     }
 
 
-def test_add_hitl_detail(client, create_task_instance, session) -> None:
+@pytest.mark.parametrize(
+    "existing_hitl_detail_args",
+    [
+        None,
+        default_hitl_detail_request_kwargs,
+        {
+            **default_hitl_detail_request_kwargs,
+            **{
+                "params_input": {"input_1": 2},
+                "response_at": convert_to_utc(datetime(2025, 7, 3, 0, 0, 0)),
+                "chosen_options": ["Reject"],
+                "user_id": "Fallback to defaults",
+            },
+        },
+    ],
+    ids=[
+        "no existing hitl detail",
+        "existing hitl detail without response",
+        "existing hitl detail with response",
+    ],
+)
+def test_upsert_hitl_detail(
+    client: TestClient,
+    create_task_instance: CreateTaskInstance,
+    session: Session,
+    existing_hitl_detail_args: dict[str, Any],
+) -> None:
     ti = create_task_instance()
     session.commit()
 
+    if existing_hitl_detail_args:
+        session.add(HITLDetail(ti_id=ti.id, **existing_hitl_detail_args))
+        session.commit()
+
     response = client.post(
         f"/execution/hitl-details/{ti.id}",
         json={
             "ti_id": ti.id,
-            "options": ["Approve", "Reject"],
-            "subject": "This is subject",
-            "body": "this is body",
-            "defaults": ["Approve"],
-            "multiple": False,
-            "params": {"input_1": 1},
+            **default_hitl_detail_request_kwargs,
         },
     )
     assert response.status_code == 201
     assert response.json() == {
         "ti_id": ti.id,
-        "options": ["Approve", "Reject"],
-        "subject": "This is subject",
-        "body": "this is body",
-        "defaults": ["Approve"],
-        "multiple": False,
-        "params": {"input_1": 1},
+        **default_hitl_detail_request_kwargs,
     }
 
 
 @time_machine.travel(datetime(2025, 7, 3, 0, 0, 0), tick=False)
 @pytest.mark.usefixtures("sample_hitl_detail")
-def test_update_hitl_detail(client, sample_ti) -> None:
+def test_update_hitl_detail(client: Client, sample_ti: TaskInstance) -> None:
     response = client.patch(
         f"/execution/hitl-details/{sample_ti.id}",
         json={
@@ -128,13 +156,7 @@ def test_update_hitl_detail(client, sample_ti) -> None:
 
 
 @pytest.mark.usefixtures("sample_hitl_detail")
-def test_get_hitl_detail(client, sample_ti) -> None:
+def test_get_hitl_detail(client: Client, sample_ti: TaskInstance) -> None:
     response = client.get(f"/execution/hitl-details/{sample_ti.id}")
     assert response.status_code == 200
-    assert response.json() == {
-        "params_input": {},
-        "response_at": None,
-        "chosen_options": None,
-        "response_received": False,
-        "user_id": None,
-    }
+    assert response.json() == expected_empty_hitl_detail_response_part

providers/standard/src/airflow/providers/standard/operators/hitl.py

@@ -34,7 +34,7 @@
 from airflow.providers.standard.triggers.hitl import HITLTrigger, HITLTriggerEventSuccessPayload
 from airflow.providers.standard.utils.skipmixin import SkipMixin
 from airflow.sdk.definitions.param import ParamsDict
-from airflow.sdk.execution_time.hitl import add_hitl_detail
+from airflow.sdk.execution_time.hitl import upsert_hitl_detail
 
 if TYPE_CHECKING:
     from airflow.sdk.definitions.context import Context
@@ -98,7 +98,7 @@ def execute(self, context: Context):
         """Add a Human-in-the-loop Response and then defer to HITLTrigger and wait for user input."""
         ti_id = context["task_instance"].id
         # Write Human-in-the-loop input request to DB
-        add_hitl_detail(
+        upsert_hitl_detail(
             ti_id=ti_id,
             options=self.options,
             subject=self.subject,

task-sdk/src/airflow/sdk/execution_time/hitl.py

@@ -30,7 +30,7 @@
     from airflow.api_fastapi.execution_api.datamodels.hitl import HITLDetailResponse
 
 
-def add_hitl_detail(
+def upsert_hitl_detail(
     ti_id: UUID,
     options: list[str],
     subject: str,

task-sdk/tests/task_sdk/execution_time/test_hitl.py

@@ -23,16 +23,16 @@
 from airflow.sdk.api.datamodels._generated import HITLDetailResponse
 from airflow.sdk.execution_time.comms import CreateHITLDetailPayload
 from airflow.sdk.execution_time.hitl import (
-    add_hitl_detail,
     get_hitl_detail_content_detail,
     update_htil_detail_response,
+    upsert_hitl_detail,
 )
 
 TI_ID = uuid7()
 
 
-def test_add_hitl_detail(mock_supervisor_comms) -> None:
-    add_hitl_detail(
+def test_upsert_hitl_detail(mock_supervisor_comms) -> None:
+    upsert_hitl_detail(
         ti_id=TI_ID,
         options=["Approve", "Reject"],
         subject="Subject",