This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5798f1f9ee0 Allow asking for human input from inside a Python task
(#75358)
5798f1f9ee0 is described below
commit 5798f1f9ee01a1f2b645203090cf5fb70500dd06
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Oct 22 22:10:17 2026 +0100
Allow asking for human input from inside a Python task (#75358)
request_input() writes a Required Actions request from inside a running
@task or PythonOperator callable and parks the task in awaiting_input
without holding a worker slot. When the request is answered, the
callable runs again and get_input_response() returns the answer.
HITLOperator and request_input() now share one helper that writes the
request and parks, and one set of request and answer checks. The answer
is validated against the options and params the reviewer was shown.
Needs Airflow 3.4+: earlier versions keep the first request's content
when a task asks again or retries.
* Add tests for the shared Human-in-the-loop checks
Cover utils/hitl.py directly and fix the type errors MyPy reported in
the request_input() tests.
* Drop Dag arguments that match the defaults from the request_input()
example
* Document how to keep a side effect from running twice after a resume
* Return the answer from request_input() itself
request_input() now pauses the task the first time a question is asked and
returns the answer when the task runs again, so a callable reads top to
bottom without checking whether it is the first run. get_input_response()
is removed.
Answers travel with the paused task, keyed by key= (default: the subject).
Each key is asked at most once per try, and a run that never reads the
answer it resumed with fails, so an answer cannot reach the wrong question
and a subject that changes between runs cannot ask forever.
* Leave the params_input cases to the shared-check tests
* Say only what is known about request_input() in subprocess tasks
Also word the unread-answer rule so it is clear that asking a new question
after reading the resumed answer is fine.
* Keep a param without a default from failing the resume
A param with no default serialized its value as NOTSET into the paused
task, and the resume could not read it back, so the task failed after the
reviewer had answered. Also document asking one question at a time.
---
providers/standard/docs/operators/hitl.rst | 71 ++
.../example_dags/example_hitl_request_input.py | 43 ++
.../airflow/providers/standard/operators/hitl.py | 261 +++++--
.../airflow/providers/standard/operators/python.py | 46 ++
.../src/airflow/providers/standard/utils/hitl.py | 153 ++++
.../airflow/providers/standard/version_compat.py | 2 +
.../standard/operators/test_hitl_request_input.py | 795 +++++++++++++++++++++
.../tests/unit/standard/utils/test_hitl.py | 97 +++
8 files changed, 1407 insertions(+), 61 deletions(-)
diff --git a/providers/standard/docs/operators/hitl.rst
b/providers/standard/docs/operators/hitl.rst
index 837dca55d11..7db9da38fe7 100644
--- a/providers/standard/docs/operators/hitl.rst
+++ b/providers/standard/docs/operators/hitl.rst
@@ -42,3 +42,74 @@ For an end-to-end guide on using HITL operators, including
examples and
recommended patterns, see the HITL tutorial:
* `HITL tutorial
<https://airflow.apache.org/docs/apache-airflow/stable/tutorial/hitl.html>`_
+
+.. _howto/operator:request_input:
+
+Asking for input from inside a Python task
+------------------------------------------
+
+Sometimes the question only comes up halfway through a task. An agent built
with the OpenAI
+Agents SDK, Strands or Google ADK, for example, stops in front of a tool call
and needs a person
+to approve it.
:func:`~airflow.providers.standard.operators.hitl.request_input` asks the
question
+and returns the answer:
+
+.. exampleinclude::
/../../standard/src/airflow/providers/standard/example_dags/example_hitl_request_input.py
+ :language: python
+ :start-after: [START howto_hitl_request_input]
+ :end-before: [END howto_hitl_request_input]
+
+The first time the task reaches ``request_input()``, the question appears on
the Required Actions
+page and the task stops there. It waits in the ``awaiting_input`` state
without holding a worker
+slot. When the reviewer answers, the task runs again from the start of its
callable, and this
+time the same call returns the answer: a dict with ``chosen_options``,
``params_input``,
+``responded_at`` (a ``datetime``), ``responded_by_user`` and ``timedout``.
+
+``request_input()`` takes the same request fields as
+:class:`~airflow.providers.standard.operators.hitl.HITLOperator` (``subject``,
``options``, ``body``,
+``defaults``, ``multiple``, ``params``, ``notifiers``, ``assigned_users``), a
``timeout``, and a
+``key``. A task can ask several questions; each run returns every answer it
already has and stops
+at the first question that has none.
+
+Things to know:
+
+* Code before ``request_input()`` runs again on every run. Keep anything
expensive, such as an
+ agent framework's saved run state, in ``task_state_store`` so the next run
restores it instead of
+ redoing it.
+* Each answer goes back to the call with the same ``key``, which defaults to
``subject``, and each
+ key is asked at most once per try. Pass ``key`` when the subject can change
between runs (for
+ example when it includes a value computed fresh each run), or when the task
can ask the same
+ subject more than once, such as one approval per agent tool call: key on the
call's id.
+* The task fails instead of guessing when the keys don't line up: when it asks
a key whose answer
+ it has already read, and when a run that resumed with an answer asks a new
question, or finishes,
+ before reading that answer.
+* Every pause carries the answers the task has received so far in this try,
including their
+ ``params_input``, so keep the number of questions per try small (a handful,
not hundreds) and the
+ form fields short.
+* A side effect between two questions, such as issuing a refund before asking
whether to email the
+ customer, runs again when the task resumes after the second answer. Record
it in
+ ``task_state_store`` and skip it when the record is there:
+
+ .. code-block:: python
+
+ if task_state_store.get("refund_id") is None:
+ task_state_store.set("refund_id", issue_refund(order_id))
+
+* A retry asks every question again, but ``task_state_store`` keeps its values
across retries, so
+ the retry still sees what the failed try saved. Delete those values when you
no longer need them.
+* Without ``defaults``, the task fails with
+ :class:`~airflow.providers.standard.exceptions.HITLTimeoutError` when
``timeout`` passes. With
+ ``defaults``, it resumes with ``timedout`` set to ``True`` and
``responded_by_user`` set to
+ ``None``.
+* It works in ``@task``, ``@task.branch``, ``@task.short_circuit`` and
``PythonOperator`` callables,
+ sync or async. It can't pause tasks that run their callable in another
process
+ (``@task.virtualenv``, ``@task.external_python``, ``@task.docker``,
``@task.kubernetes``).
+* Ask one question at a time, from the task's own code. A task instance has a
single Required
+ Actions request, so two questions asked at once (``asyncio.gather``,
``asyncio.TaskGroup``, or
+ threads) overwrite each other, and a call from a worker thread doesn't see
the answers the task
+ resumed with. Don't call it inside code that turns exceptions into a result,
such as an agent
+ framework's tool-error handler: the errors above only fail the task if they
reach it.
+* It needs Airflow 3.4+.
+* The Required Actions page refuses an answer to a request that the task has
since replaced. A
+ client answering through the REST API gets the same protection only if it
sends the request's
+ ``created_at`` (from the GET) as ``request_created_at``; without it, the
answer is applied to
+ whatever request is waiting.
diff --git
a/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_request_input.py
b/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_request_input.py
new file mode 100644
index 00000000000..b0b13c9e7a2
--- /dev/null
+++
b/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_request_input.py
@@ -0,0 +1,43 @@
+#
+# 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.
+"""Example Dag asking a reviewer for input from inside a ``@task`` with
``request_input()``."""
+
+from __future__ import annotations
+
+import datetime
+
+from airflow.providers.standard.operators.hitl import request_input
+from airflow.sdk import DAG, Param, task
+
+with DAG(dag_id="example_hitl_request_input", tags=["example", "HITL"]):
+ # [START howto_hitl_request_input]
+ @task
+ def refund(order_id: str) -> str:
+ answer = request_input(
+ subject=f"Refund order {order_id}?",
+ options=["Approve", "Reject"],
+ params={"reason": Param("", type=["string", "null"])},
+ timeout=datetime.timedelta(hours=4),
+ )
+ if answer["chosen_options"] == ["Approve"]:
+ return f"Refunded order {order_id}"
+ return f"Refund for {order_id} rejected:
{answer['params_input'].get('reason')}"
+
+ # [END howto_hitl_request_input]
+
+ refund("o-1")
diff --git
a/providers/standard/src/airflow/providers/standard/operators/hitl.py
b/providers/standard/src/airflow/providers/standard/operators/hitl.py
index 3ad7639ad98..d5519e7ad69 100644
--- a/providers/standard/src/airflow/providers/standard/operators/hitl.py
+++ b/providers/standard/src/airflow/providers/standard/operators/hitl.py
@@ -25,6 +25,7 @@ from airflow.providers.standard.version_compat import (
AIRFLOW_V_3_1_3_PLUS,
AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS,
+ AIRFLOW_V_3_4_PLUS,
)
if not AIRFLOW_V_3_1_PLUS:
@@ -32,13 +33,28 @@ if not AIRFLOW_V_3_1_PLUS:
from collections.abc import Collection, Mapping, Sequence
from datetime import timedelta
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NoReturn
from urllib.parse import ParseResult, urlencode, urlparse, urlunparse
-from airflow.providers.common.compat.sdk import BranchMixIn, SkipMixin, conf
-from airflow.providers.standard.exceptions import HITLRejectException,
HITLTimeoutError, HITLTriggerEventError
+from airflow.providers.common.compat.sdk import BranchMixIn, SkipMixin, conf,
get_current_context
+from airflow.providers.standard.exceptions import ( # noqa: F401 - kept
importable from this module
+ HITLRejectException,
+ HITLTimeoutError,
+ HITLTriggerEventError,
+)
+from airflow.providers.standard.operators.python import PythonOperator,
_BasePythonVirtualenvOperator
from airflow.providers.standard.triggers.hitl import HITLTrigger,
HITLTriggerEventSuccessPayload
-from airflow.providers.standard.version_compat import BaseOperator
+from airflow.providers.standard.utils.hitl import (
+ InputResponse,
+ answer_log_var,
+ raise_for_error_event,
+ validate_chosen_options,
+ validate_defaults,
+ validate_options,
+ validate_params,
+ validate_params_input,
+)
+from airflow.providers.standard.version_compat import BaseOperator, is_arg_set
from airflow.sdk.bases.notifier import BaseNotifier
from airflow.sdk.definitions.param import ParamsDict
from airflow.sdk.execution_time.hitl import upsert_hitl_detail
@@ -50,11 +66,172 @@ if AIRFLOW_V_3_3_PLUS:
from airflow.sdk.exceptions import TaskAwaitingInput
if TYPE_CHECKING:
+ from datetime import datetime
+
from airflow.providers.common.compat.sdk import Context
from airflow.sdk.execution_time.hitl import HITLUser
from airflow.sdk.types import RuntimeTaskInstanceProtocol
+def _request_input_and_wait(
+ task: BaseOperator,
+ context: Context,
+ *,
+ options: list[str],
+ subject: str,
+ body: str | None,
+ defaults: list[str] | None,
+ multiple: bool,
+ params: dict[str, Any],
+ assigned_users: list[HITLUser] | None,
+ notifiers: Sequence[BaseNotifier],
+ timeout: timedelta | None,
+ timeout_datetime: datetime | None,
+ method_name: str,
+ kwargs: dict[str, Any] | None = None,
+) -> NoReturn:
+ """Write the Human-in-the-loop request, notify, and stop the task until a
response arrives."""
+ ti_id = context["task_instance"].id
+ upsert_hitl_detail(
+ ti_id=ti_id,
+ options=options,
+ subject=subject,
+ body=body,
+ defaults=defaults,
+ multiple=multiple,
+ params=params,
+ assigned_users=assigned_users,
+ )
+
+ task.log.info("Waiting for response")
+ for notifier in notifiers:
+ notifier(context)
+
+ if AIRFLOW_V_3_3_PLUS:
+ # Airflow 3.3+: park the task in AWAITING_INPUT -- no trigger, no
triggerer. The task
+ # is resumed by the Core API response handler or the scheduler timeout
sweep, so the
+ # triggerer no longer needs to run for Human-in-the-loop tasks to make
progress.
+ raise TaskAwaitingInput(method_name=method_name, kwargs=kwargs,
timeout=timeout)
+
+ # Fallback for Airflow versions < 3.3: defer the response check to
HITLTrigger on the triggerer.
+ task.defer(
+ trigger=HITLTrigger(
+ ti_id=ti_id,
+ options=options,
+ defaults=defaults,
+ params=params,
+ multiple=multiple,
+ timeout_datetime=timeout_datetime,
+ ),
+ method_name=method_name,
+ kwargs=kwargs,
+ )
+
+
+def request_input(
+ *,
+ subject: str,
+ options: list[str],
+ body: str | None = None,
+ defaults: str | list[str] | None = None,
+ multiple: bool = False,
+ params: ParamsDict | dict[str, Any] | None = None,
+ notifiers: Sequence[BaseNotifier] | BaseNotifier | None = None,
+ assigned_users: HITLUser | list[HITLUser] | None = None,
+ timeout: timedelta | None = None,
+ key: str | None = None,
+) -> InputResponse:
+ """
+ Ask a person for input from inside a running Python task, and return their
answer.
+
+ The first time a question is asked, the request appears on the Required
Actions page, the same
+ as one from :class:`HITLOperator`, and the task stops there: it waits in
the ``awaiting_input``
+ state without holding a worker slot. Once the request is answered, the
task runs again from the
+ start of its callable, and this call returns the answer instead of asking
again. Save anything
+ the second run should not redo (for example an agent framework's run
state) to
+ ``task_state_store`` before calling this.
+
+ It works in ``@task`` and ``PythonOperator`` callables (sync or async),
including ``@task.branch``
+ and ``@task.short_circuit``. It can't pause tasks that run their callable
in another process, such
+ as ``@task.virtualenv``, ``@task.external_python``, ``@task.docker`` or
``@task.kubernetes``.
+ Ask one question at a time from the task's own code, not from worker
threads or from several
+ ``async`` calls running at once: a task instance has a single Required
Actions request.
+
+ Needs Airflow 3.4+. Earlier versions keep the first request's subject,
body and options when a
+ task asks again, after a resume or on a retry, so the reviewer would
answer a question they were
+ not shown.
+
+ :param subject: Headline shown to the reviewer.
+ :param options: Options the reviewer chooses from.
+ :param body: Details the reviewer needs to decide (Markdown).
+ :param defaults: Options applied if ``timeout`` passes without an answer.
Without defaults, the
+ task fails with
:class:`~airflow.providers.standard.exceptions.HITLTimeoutError` at the timeout.
+ :param multiple: Whether the reviewer can choose more than one option.
+ :param params: Form fields for the reviewer to fill in, defined like Dag
params.
+ :param notifiers: Notifiers called with the task context once the request
is written.
+ :param assigned_users: Users allowed to answer. Anyone with access can
answer if not set.
+ :param timeout: How long to wait for an answer. Waits indefinitely if not
set.
+ :param key: Which question this is, so a resumed task gets each answer
back at the right call.
+ Defaults to ``subject``. Each key is asked at most once per try. Set
it when the subject can
+ change between runs, or when the task can ask the same subject more
than once, for example
+ one approval per agent tool call (use the call's id).
+ """
+ if not AIRFLOW_V_3_4_PLUS:
+ raise AirflowOptionalProviderFeatureException("request_input() needs
Airflow 3.4+.")
+ context = get_current_context()
+ task = context["task"]
+ if not isinstance(task, PythonOperator) or isinstance(task,
_BasePythonVirtualenvOperator):
+ raise TypeError(
+ f"request_input() was called from {type(task).__name__}, which
cannot resume it. It works "
+ "only in tasks that run their Python callable in the task's own
process, such as @task "
+ "and PythonOperator."
+ )
+ defaults = [defaults] if isinstance(defaults, str) else defaults
+ params = params if isinstance(params, ParamsDict) else ParamsDict(params
or {})
+ validate_options(options)
+ validate_params(params)
+ validate_defaults(defaults, options, multiple)
+ key = subject if key is None else key
+ answer_log = answer_log_var.get()
+ if answer_log is not None and (answer := answer_log.take(key)) is not None:
+ # The code may have changed since the reviewer answered; never hand
back an option it no
+ # longer offers.
+ validate_chosen_options(answer["chosen_options"], options)
+ return answer
+ serialized_params = {name: params.get_param(name).serialize() for name in
params}
+
+ _request_input_and_wait(
+ task,
+ context,
+ options=options,
+ subject=subject,
+ body=body,
+ defaults=defaults,
+ multiple=multiple,
+ params=serialized_params,
+ assigned_users=[assigned_users] if isinstance(assigned_users, dict)
else assigned_users,
+ notifiers=[notifiers] if isinstance(notifiers, BaseNotifier) else
notifiers or [],
+ timeout=timeout,
+ timeout_datetime=None,
+ method_name="resume_with_input",
+ kwargs={
+ "key": key,
+ # What the resume checks the answer against.
+ "options": options,
+ # A param without a default serializes its value as NOTSET, which
can't be stored with
+ # the paused task. The resume only checks the answer against each
schema, so None will do.
+ "params": {
+ name: {**spec, "value": spec["value"] if
is_arg_set(spec["value"]) else None}
+ for name, spec in serialized_params.items()
+ },
+ # Earlier answers travel with the task so each run can replay
every one of them.
+ "answered": [[answered_key, earlier] for answered_key, earlier in
answer_log.answers.items()]
+ if answer_log is not None
+ else [],
+ },
+ )
+
+
class HITLOperator(BaseOperator):
"""
Base class for all Human-in-the-loop Operators to inherit from.
@@ -163,8 +340,7 @@ class HITLOperator(BaseOperator):
Raises:
ValueError: If `options` is empty.
"""
- if not self.options:
- raise ValueError('"options" cannot be empty.')
+ validate_options(self.options)
def validate_params(self) -> None:
"""
@@ -179,8 +355,7 @@ class HITLOperator(BaseOperator):
Raises:
ValueError: If ``"_options"`` key is present in ``params``, which
is not allowed.
"""
- if "_options" in self.params:
- raise ValueError('"_options" is not allowed in params')
+ validate_params(self.params)
def validate_defaults(self) -> None:
"""
@@ -189,12 +364,7 @@ class HITLOperator(BaseOperator):
1. Default options should be the subset of options.
2. When multiple is False, there should only be one option.
"""
- if self.defaults is not None:
- if not set(self.defaults).issubset(self.options):
- raise ValueError(f'defaults "{self.defaults}" should be a
subset of options "{self.options}"')
-
- if self.multiple is False and len(self.defaults) > 1:
- raise ValueError('More than one defaults given when "multiple"
is set to False.')
+ validate_defaults(self.defaults, self.options, self.multiple)
def execute(self, context: Context):
"""
@@ -204,19 +374,6 @@ class HITLOperator(BaseOperator):
involved; on older versions it defers to :class:`HITLTrigger`. Either
way it resumes in
``execute_complete`` once a response (or timeout default) arrives.
"""
- ti_id = context["task_instance"].id
- # Write Human-in-the-loop input request to DB
- upsert_hitl_detail(
- ti_id=ti_id,
- options=self.options,
- subject=self.subject,
- body=self.body,
- defaults=self.defaults,
- multiple=self.multiple,
- params=self.serialized_params,
- assigned_users=self.assigned_users,
- )
-
if self.response_timeout:
timeout_datetime = utcnow() + self.response_timeout
else:
@@ -227,29 +384,19 @@ class HITLOperator(BaseOperator):
timeout_datetime.isoformat() if timeout_datetime else None
)
- self.log.info("Waiting for response")
- for notifier in self.notifiers:
- notifier(context)
-
- if AIRFLOW_V_3_3_PLUS:
- # Airflow 3.3+: park the task in AWAITING_INPUT -- no trigger, no
triggerer. The task
- # is resumed by the Core API response handler or the scheduler
timeout sweep, so the
- # triggerer no longer needs to run for Human-in-the-loop tasks to
make progress.
- raise TaskAwaitingInput(
- method_name="execute_complete",
- timeout=self.response_timeout,
- )
-
- # Fallback for Airflow versions < 3.3: defer the response check to
HITLTrigger on the triggerer.
- self.defer(
- trigger=HITLTrigger(
- ti_id=ti_id,
- options=self.options,
- defaults=self.defaults,
- params=self.serialized_params,
- multiple=self.multiple,
- timeout_datetime=timeout_datetime,
- ),
+ _request_input_and_wait(
+ self,
+ context,
+ options=self.options,
+ subject=self.subject,
+ body=self.body,
+ defaults=self.defaults,
+ multiple=self.multiple,
+ params=self.serialized_params,
+ assigned_users=self.assigned_users,
+ notifiers=self.notifiers,
+ timeout=self.response_timeout,
+ timeout_datetime=timeout_datetime,
method_name="execute_complete",
)
@@ -286,23 +433,15 @@ class HITLOperator(BaseOperator):
)
def process_trigger_event_error(self, event: dict[str, Any]) -> None:
- if event["error_type"] == "timeout":
- raise HITLTimeoutError(event)
-
- raise HITLTriggerEventError(event)
+ raise_for_error_event(event)
def validate_chosen_options(self, chosen_options: list[str]) -> None:
"""Check whether user provide valid response."""
- if diff := set(chosen_options) - set(self.options):
- raise ValueError(f"Responses {diff} not in {self.options}")
+ validate_chosen_options(chosen_options, self.options)
def validate_params_input(self, params_input: Mapping) -> None:
"""Check whether user provide valid params input."""
- if self.params and params_input and set(self.serialized_params.keys())
^ set(params_input):
- raise ValueError(f"params_input {params_input} does not match
params {self.params}")
-
- for key, value in params_input.items():
- self.params[key] = value
+ validate_params_input(params_input, self.params)
def generate_link_to_ui(
self,
diff --git
a/providers/standard/src/airflow/providers/standard/operators/python.py
b/providers/standard/src/airflow/providers/standard/operators/python.py
index 8698ddc403b..f83c3c17923 100644
--- a/providers/standard/src/airflow/providers/standard/operators/python.py
+++ b/providers/standard/src/airflow/providers/standard/operators/python.py
@@ -61,6 +61,15 @@ from airflow.providers.common.compat.standard.operators
import (
is_async_callable,
)
from airflow.providers.standard.hooks.package_index import PackageIndexHook
+from airflow.providers.standard.utils.hitl import (
+ AnswerLog,
+ InputResponse,
+ answer_log_var,
+ params_from_serialized,
+ raise_for_error_event,
+ validate_chosen_options,
+ validate_params_input,
+)
from airflow.providers.standard.utils.python_virtualenv import (
_execute_in_subprocess,
prepare_virtualenv,
@@ -236,6 +245,43 @@ class PythonOperator(BaseAsyncOperator):
return return_value
+ def resume_with_input(
+ self,
+ context: Context,
+ event: dict[str, Any],
+ key: str,
+ options: list[str],
+ params: dict[str, dict[str, Any]],
+ answered: list[list[Any]],
+ ) -> Any:
+ """
+ Run the task again once the input it asked for with
``request_input()`` arrives.
+
+ The callable runs from the top. Each ``request_input()`` call that was
answered, in this or
+ an earlier pause, returns its answer instead of asking again, and the
run fails if it never
+ reads the answer it resumed with. A request that timed out without
defaults raises
+ :class:`~airflow.providers.standard.exceptions.HITLTimeoutError`.
+ """
+ if "error" in event:
+ raise_for_error_event(event)
+ validate_chosen_options(event["chosen_options"], options)
+ validate_params_input(event["params_input"],
params_from_serialized(params))
+ answer = InputResponse(
+ chosen_options=event["chosen_options"],
+ params_input=event["params_input"],
+ responded_at=event["responded_at"],
+ responded_by_user=event["responded_by_user"],
+ timedout=event["timedout"],
+ )
+ answer_log = AnswerLog({**dict(answered), key: answer}, latest=key)
+ token = answer_log_var.set(answer_log)
+ try:
+ result = self.execute(context)
+ answer_log.check_latest_read()
+ return result
+ finally:
+ answer_log_var.reset(token)
+
def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str,
Any]:
return KeywordParameters.determine(self.python_callable, self.op_args,
context).unpacking()
diff --git a/providers/standard/src/airflow/providers/standard/utils/hitl.py
b/providers/standard/src/airflow/providers/standard/utils/hitl.py
new file mode 100644
index 00000000000..4994b9c950f
--- /dev/null
+++ b/providers/standard/src/airflow/providers/standard/utils/hitl.py
@@ -0,0 +1,153 @@
+# 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.
+"""
+Human-in-the-loop checks shared by ``HITLOperator`` and ``request_input()``,
and the answers a resumed task replays.
+
+Kept apart from :mod:`airflow.providers.standard.operators.hitl`, which cannot
be imported before
+Airflow 3.1, so that ``PythonOperator`` can import it on every supported
Airflow version.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any, NoReturn, TypedDict
+
+from airflow.providers.standard.exceptions import HITLTimeoutError,
HITLTriggerEventError
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+ from airflow.sdk.definitions.param import Param, ParamsDict
+else:
+ from airflow.models.param import Param, ParamsDict # type:
ignore[no-redef]
+
+if TYPE_CHECKING:
+ from datetime import datetime
+
+ from airflow.sdk.execution_time.hitl import HITLUser
+
+# The helpers below are shared by the HITL operators and PythonOperator, not
public API.
+__all__ = ["InputResponse"]
+
+
+class InputResponse(TypedDict):
+ """The answer to a
:func:`~airflow.providers.standard.operators.hitl.request_input` call."""
+
+ chosen_options: list[str]
+ params_input: dict[str, Any]
+ responded_at: datetime
+ #: ``None`` when nobody answered and the request's defaults were applied
at its timeout.
+ responded_by_user: HITLUser | None
+ timedout: bool
+
+
+class AnswerLog:
+ """
+ The answers a resumed task has received in this try, by the key of the
request each one answered.
+
+ Each key is asked at most once per try, so an answer can never be handed
to a later question that
+ happens to share its key. The run that resumed with an answer (``latest``)
must read it before it
+ asks anything new or finishes; otherwise the key changed between runs, and
asking again would
+ never end.
+ """
+
+ def __init__(self, answers: dict[str, InputResponse], latest: str) -> None:
+ self.answers = answers
+ self.latest = latest
+ self._read: set[str] = set()
+
+ def take(self, key: str) -> InputResponse | None:
+ """Return the answer to ``key``, or ``None`` when ``key`` has not been
asked yet."""
+ if key in self._read:
+ raise ValueError(
+ f"request_input() was called again with the key {key!r} after
its answer was read. Each "
+ "question gets one answer per try, so pass a key that is
unique to each question, for "
+ "example an agent's tool call id."
+ )
+ answer = self.answers.get(key)
+ if answer is not None:
+ self._read.add(key)
+ return answer
+ self.check_latest_read()
+ return None
+
+ def check_latest_read(self) -> None:
+ """Refuse to go on when the answer this run resumed with was never
read."""
+ if self.latest not in self._read:
+ raise RuntimeError(
+ f"The answer to {self.latest!r} was never read: no
request_input() call with that key "
+ "ran after the task resumed. If the subject changes between
runs, pass a fixed key=."
+ )
+
+
+answer_log_var: ContextVar[AnswerLog | None] = ContextVar("hitl_answer_log",
default=None)
+
+
+def validate_options(options: list[str]) -> None:
+ """Refuse a request with no options."""
+ if not options:
+ raise ValueError('"options" cannot be empty.')
+
+
+def validate_params(params: Mapping[str, Any]) -> None:
+ """Refuse params that collide with the reserved ``_options`` key."""
+ if "_options" in params:
+ raise ValueError('"_options" is not allowed in params')
+
+
+def validate_defaults(defaults: list[str] | None, options: list[str],
multiple: bool) -> None:
+ """Refuse defaults that are not options, or more than one default for a
single-choice request."""
+ if defaults is not None:
+ if not set(defaults).issubset(options):
+ raise ValueError(f'defaults "{defaults}" should be a subset of
options "{options}"')
+
+ if multiple is False and len(defaults) > 1:
+ raise ValueError('More than one defaults given when "multiple" is
set to False.')
+
+
+def raise_for_error_event(event: Mapping[str, Any]) -> NoReturn:
+ """Raise the error a resume event carries instead of an answer, such as a
timeout without defaults."""
+ if event["error_type"] == "timeout":
+ raise HITLTimeoutError(event)
+
+ raise HITLTriggerEventError(event)
+
+
+def validate_chosen_options(chosen_options: list[str], options: list[str]) ->
None:
+ """Refuse an answer that chose something the request did not offer."""
+ if diff := set(chosen_options) - set(options):
+ raise ValueError(f"Responses {diff} not in {options}")
+
+
+def validate_params_input(params_input: Mapping[str, Any], params: ParamsDict)
-> None:
+ """
+ Refuse form input that does not match the request's params, and check each
value against its schema.
+
+ The values are written into ``params``, which raises if one fails its
param's schema.
+ """
+ if params and params_input and set(params.keys()) ^ set(params_input):
+ raise ValueError(f"params_input {params_input} does not match params
{params}")
+
+ for key, value in params_input.items():
+ params[key] = value
+
+
+def params_from_serialized(serialized: Mapping[str, Mapping[str, Any]]) ->
ParamsDict:
+ """Rebuild the params of a request from the form they were written to the
Human-in-the-loop detail in."""
+ return ParamsDict(
+ {key: Param.deserialize(dict(value), Param.__version__) for key, value
in serialized.items()}
+ )
diff --git
a/providers/standard/src/airflow/providers/standard/version_compat.py
b/providers/standard/src/airflow/providers/standard/version_compat.py
index f99a5589378..146a5828fa4 100644
--- a/providers/standard/src/airflow/providers/standard/version_compat.py
+++ b/providers/standard/src/airflow/providers/standard/version_compat.py
@@ -37,6 +37,7 @@ AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple()
>= (3, 1, 0)
AIRFLOW_V_3_1_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 3)
AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0)
AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS: bool = get_base_airflow_version_tuple() >= (3, 4, 0)
# BaseOperator: Use 3.1+ due to xcom_push method missing in SDK BaseOperator
3.0.x
# This is needed for DecoratedOperator compatibility
@@ -60,6 +61,7 @@ __all__ = [
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_2_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
"ArgNotSet",
"BaseOperator",
"is_arg_set",
diff --git
a/providers/standard/tests/unit/standard/operators/test_hitl_request_input.py
b/providers/standard/tests/unit/standard/operators/test_hitl_request_input.py
new file mode 100644
index 00000000000..44130a11a10
--- /dev/null
+++
b/providers/standard/tests/unit/standard/operators/test_hitl_request_input.py
@@ -0,0 +1,795 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import datetime
+from typing import TYPE_CHECKING, Any
+
+import pytest
+import structlog
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_4_PLUS
+
+if not AIRFLOW_V_3_4_PLUS:
+ pytest.skip("request_input() needs Airflow 3.4+", allow_module_level=True)
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, ParamValidationError
+from airflow.providers.standard.exceptions import HITLTimeoutError
+from airflow.providers.standard.operators.bash import BashOperator
+from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.providers.standard.operators.hitl import HITLOperator,
request_input
+from airflow.providers.standard.operators.python import PythonOperator,
PythonVirtualenvOperator
+from airflow.providers.standard.utils.hitl import answer_log_var
+from airflow.sdk import DAG, BaseNotifier, Param, task
+from airflow.sdk.exceptions import TaskAwaitingInput
+from airflow.sdk.execution_time.comms import (
+ AwaitInputTask,
+ CreateHITLDetailPayload,
+ SetXCom,
+ SkipDownstreamTasks,
+)
+from airflow.sdk.execution_time.context import set_current_context
+from airflow.sdk.execution_time.task_runner import run
+from airflow.sdk.serde import deserialize, serialize
+from airflow.utils.state import TaskInstanceState
+
+if TYPE_CHECKING:
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+RESPONDED_AT = datetime.datetime(2026, 10, 21, 12, 0, tzinfo=datetime.UTC)
+REVIEWER = {"id": "alice", "name": "Alice"}
+
+
+def _answer(chosen: str, *, reason: str = "", timedout: bool = False) ->
dict[str, Any]:
+ """The event the Core API response handler puts in next_kwargs
(HITLDetail.as_resume_event_payload)."""
+ return {
+ "chosen_options": [chosen],
+ "params_input": {"reason": reason},
+ "responded_at": RESPONDED_AT,
+ "responded_by_user": None if timedout else REVIEWER,
+ "timedout": timedout,
+ }
+
+
+def _plain_answer(chosen: str) -> dict[str, Any]:
+ """An answer to a request that has no params."""
+ return {**_answer(chosen), "params_input": {}}
+
+
+def _resume(ti: RuntimeTaskInstance, park_msg: AwaitInputTask, event:
dict[str, Any]) -> None:
+ """Set up ``ti`` the way the server hands back a task whose request was
answered."""
+ server_context = ti._ti_context_from_server
+ parked_kwargs = deserialize(park_msg.next_kwargs)
+ assert server_context is not None
+ assert isinstance(parked_kwargs, dict)
+ server_context.next_method = park_msg.next_method
+ # handle_event_submit keeps the parked kwargs and adds the response as
``event``.
+ server_context.next_kwargs = serialize({**parked_kwargs, "event": event})
+
+
+def _sent(mock_supervisor_comms, msg_type: type) -> list[Any]:
+ # Some callers pass the message positionally, others as msg=.
+ messages = (
+ c.kwargs.get("msg", c.args[0] if c.args else None) for c in
mock_supervisor_comms.send.call_args_list
+ )
+ return [msg for msg in messages if isinstance(msg, msg_type)]
+
+
+def _hitl_requests(mock_supervisor_comms) -> list[CreateHITLDetailPayload]:
+ return _sent(mock_supervisor_comms, CreateHITLDetailPayload)
+
+
+def _skipped_task_ids(mock_supervisor_comms) -> list[str]:
+ return [
+ t if isinstance(t, str) else t[0]
+ for msg in _sent(mock_supervisor_comms, SkipDownstreamTasks)
+ for t in msg.tasks
+ ]
+
+
+def _return_values(mock_supervisor_comms) -> list[Any]:
+ return [x.value for x in _sent(mock_supervisor_comms, SetXCom) if x.key ==
"return_value"]
+
+
+def _noop():
+ pass
+
+
+def _run(ti: RuntimeTaskInstance):
+ return run(ti, context=ti.get_template_context(),
log=structlog.get_logger())
+
+
+def _answer_and_run(create_runtime_ti, op, ti_id, park_msg: AwaitInputTask,
event: dict[str, Any]):
+ """Resume the task with ``event`` as the answer to the request it parked
on, and run it."""
+ resumed = create_runtime_ti(task=op, ti_id=ti_id)
+ _resume(resumed, park_msg, event)
+ return _run(resumed)
+
+
+class _RecordingNotifier(BaseNotifier):
+ def __init__(self) -> None:
+ super().__init__()
+ self.task_ids: list[str] = []
+
+ def notify(self, context) -> None:
+ self.task_ids.append(context["ti"].task_id)
+
+
+class TestRequestInputFromTask:
+ @pytest.fixture
+ def seen(self) -> list[Any]:
+ return []
+
+ @pytest.fixture
+ def refund_op(self, seen):
+ with DAG(dag_id="refund_dag"):
+
+ @task
+ def refund(order_id: str):
+ answer = request_input(
+ subject=f"Refund {order_id}?",
+ body="The agent wants to call `refund_order`.",
+ options=["Approve", "Reject"],
+ params={"reason": Param("", type=["string", "null"])},
+ timeout=datetime.timedelta(hours=4),
+ )
+ seen.append(answer)
+ if answer["chosen_options"] == ["Approve"]:
+ return f"refunded {order_id}"
+ return f"refused {order_id}:
{answer['params_input']['reason']}"
+
+ xcom_arg = refund("o-1")
+ return xcom_arg.operator
+
+ def test_first_run_writes_the_request_and_waits_in_awaiting_input(
+ self, refund_op, seen, create_runtime_ti, mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=refund_op)
+
+ state, msg, err = _run(ti)
+
+ assert err is None
+ assert state == TaskInstanceState.AWAITING_INPUT
+ assert isinstance(msg, AwaitInputTask)
+ assert msg.next_method == "resume_with_input"
+ (request,) = _hitl_requests(mock_supervisor_comms)
+ # The resume checks the answer against the options and params the
reviewer was shown.
+ assert deserialize(msg.next_kwargs) == {
+ "key": "Refund o-1?",
+ "options": ["Approve", "Reject"],
+ "params": request.params,
+ "answered": [],
+ }
+ assert msg.timeout == datetime.timedelta(hours=4)
+ assert request.ti_id == ti.id
+ assert request.subject == "Refund o-1?"
+ assert request.options == ["Approve", "Reject"]
+ assert request.body == "The agent wants to call `refund_order`."
+ assert set(request.params) == {"reason"}
+ assert seen == []
+
+ @pytest.mark.parametrize(
+ ("chosen", "reason", "expected"),
+ [
+ pytest.param("Approve", "", "refunded o-1", id="approve"),
+ pytest.param("Reject", "over limit", "refused o-1: over limit",
id="reject"),
+ ],
+ )
+ def test_resumed_run_returns_the_answer_from_the_same_call(
+ self, refund_op, seen, create_runtime_ti, mock_supervisor_comms,
chosen, reason, expected
+ ):
+ ti = create_runtime_ti(task=refund_op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(
+ create_runtime_ti, refund_op, ti.id, park_msg, _answer(chosen,
reason=reason)
+ )
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert seen == [
+ {
+ "chosen_options": [chosen],
+ "params_input": {"reason": reason},
+ "responded_at": RESPONDED_AT,
+ "responded_by_user": REVIEWER,
+ "timedout": False,
+ }
+ ]
+ assert _return_values(mock_supervisor_comms) == [expected]
+ assert len(_hitl_requests(mock_supervisor_comms)) == 1
+
+ def test_timeout_without_defaults_fails_the_task(
+ self, refund_op, seen, create_runtime_ti, mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=refund_op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(
+ create_runtime_ti,
+ refund_op,
+ ti.id,
+ park_msg,
+ {
+ "error": "The Human-in-the-loop response timeout has passed
without a response.",
+ "error_type": "timeout",
+ },
+ )
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, HITLTimeoutError)
+ assert seen == []
+
+ def test_answer_outside_options_fails(self, refund_op, create_runtime_ti,
mock_supervisor_comms):
+ ti = create_runtime_ti(task=refund_op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, refund_op, ti.id,
park_msg, _answer("Maybe"))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, ValueError)
+
+ def test_params_input_is_checked_against_the_request_params(
+ self, refund_op, seen, create_runtime_ti, mock_supervisor_comms
+ ):
+ ti = create_runtime_ti(task=refund_op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(
+ create_runtime_ti,
+ refund_op,
+ ti.id,
+ park_msg,
+ {**_answer("Approve"), "params_input": {"reason": 5}},
+ )
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, ParamValidationError)
+ assert "Invalid input for param reason" in str(err)
+ assert seen == []
+
+
+class TestRequestFields:
+ def test_every_argument_reaches_the_request_and_defaults_apply_at_timeout(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ notifier = _RecordingNotifier()
+ seen: list[Any] = []
+ with DAG(dag_id="fields_dag"):
+
+ @task
+ def ask():
+ answer = request_input(
+ subject="Pick sizes",
+ body="**Which** sizes?",
+ options=["S", "M", "L"],
+ defaults="M",
+ multiple=True,
+ params={"note": Param("", type=["string", "null"])},
+ notifiers=notifier,
+ assigned_users={"id": "alice", "name": "Alice"},
+ timeout=datetime.timedelta(minutes=7),
+ )
+ seen.append(answer)
+ return answer["chosen_options"]
+
+ op = ask().operator
+ ti = create_runtime_ti(task=op)
+
+ _, park_msg, _ = _run(ti)
+
+ (request,) = _hitl_requests(mock_supervisor_comms)
+ assert request.subject == "Pick sizes"
+ assert request.body == "**Which** sizes?"
+ assert request.options == ["S", "M", "L"]
+ assert request.defaults == ["M"]
+ assert request.multiple is True
+ assert set(request.params) == {"note"}
+ assert [(u.id, u.name) for u in request.assigned_users] == [("alice",
"Alice")]
+ assert park_msg.timeout == datetime.timedelta(minutes=7)
+ assert notifier.task_ids == ["ask"]
+
+ # The scheduler's timeout sweep resumes a request that has defaults
with them applied.
+ state, _, err = _answer_and_run(
+ create_runtime_ti,
+ op,
+ ti.id,
+ park_msg,
+ {
+ "chosen_options": ["M"],
+ "params_input": {"note": ""},
+ "responded_at": RESPONDED_AT,
+ "responded_by_user": None,
+ "timedout": True,
+ },
+ )
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert seen[0]["timedout"] is True
+ assert seen[0]["responded_by_user"] is None
+ assert seen[0]["responded_at"] == RESPONDED_AT
+
+ def test_a_required_param_without_a_default_survives_the_pause(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ with DAG(dag_id="required_param_dag"):
+
+ @task
+ def ask():
+ answer = request_input(
+ subject="Why?",
+ options=["Approve", "Reject"],
+ params={"reason": Param(type="string", minLength=1)},
+ )
+ return answer["params_input"]["reason"]
+
+ op = ask().operator
+ ti = create_runtime_ti(task=op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(
+ create_runtime_ti, op, ti.id, park_msg, {**_answer("Approve"),
"params_input": {"reason": "ok"}}
+ )
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _return_values(mock_supervisor_comms) == ["ok"]
+
+ def test_hitl_operator_parks_with_its_own_resume_method_and_timeout(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ op = HITLOperator(
+ task_id="hitl",
+ subject="Go?",
+ options=["Yes", "No"],
+ response_timeout=datetime.timedelta(minutes=3),
+ )
+ context = create_runtime_ti(task=op).get_template_context()
+
+ with set_current_context(context):
+ with pytest.raises(TaskAwaitingInput) as exc_info:
+ op.execute(context)
+
+ assert exc_info.value.method_name == "execute_complete"
+ assert exc_info.value.timeout == datetime.timedelta(minutes=3)
+ assert exc_info.value.kwargs is None
+
+
+class TestSeveralQuestions:
+ def test_each_run_replays_earlier_answers_and_stops_at_the_next_question(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ refund_answers: list[Any] = []
+ with DAG(dag_id="two_questions_dag"):
+
+ @task
+ def refund_then_email():
+ refund = request_input(subject="Refund the order?",
options=["Approve", "Reject"])
+ refund_answers.append(refund)
+ email = request_input(subject="Email the customer?",
options=["Yes", "No"])
+ return [refund["chosen_options"][0],
email["chosen_options"][0]]
+
+ op = refund_then_email().operator
+ ti = create_runtime_ti(task=op)
+ _, first_park, _ = _run(ti)
+
+ state, second_park, err = _answer_and_run(
+ create_runtime_ti, op, ti.id, first_park, _plain_answer("Approve")
+ )
+
+ assert err is None
+ assert state == TaskInstanceState.AWAITING_INPUT
+ assert isinstance(second_park, AwaitInputTask)
+ parked = deserialize(second_park.next_kwargs)
+ assert parked["key"] == "Email the customer?"
+ assert [answered_key for answered_key, _ in parked["answered"]] ==
["Refund the order?"]
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id,
second_park, _plain_answer("No"))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _return_values(mock_supervisor_comms) == [["Approve", "No"]]
+ # The third run replayed the first answer instead of asking again.
+ assert [r.subject for r in _hitl_requests(mock_supervisor_comms)] == [
+ "Refund the order?",
+ "Email the customer?",
+ ]
+ # The first answer went through the parked kwargs twice and is still
the same answer.
+ assert refund_answers[-1] == refund_answers[0]
+ assert type(refund_answers[-1]["responded_at"]) is datetime.datetime
+
+ def
test_a_task_that_skips_an_answered_question_gets_the_answer_to_the_one_it_asks(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ # Stands in for state the task restores from task_state_store, the way
an agent restores
+ # its saved run and goes straight to the question it is waiting on.
+ progress: list[str] = []
+ with DAG(dag_id="restoring_dag"):
+
+ @task
+ def agent():
+ if not progress:
+ first = request_input(subject="Call refund_order?",
options=["Yes", "No"])
+ progress.append(first["chosen_options"][0])
+ second = request_input(subject="Call email_customer?",
options=["Yes", "No"])
+ return [*progress, second["chosen_options"][0]]
+
+ op = agent().operator
+ ti = create_runtime_ti(task=op)
+ _, first_park, _ = _run(ti)
+ _, second_park, _ = _answer_and_run(create_runtime_ti, op, ti.id,
first_park, _plain_answer("Yes"))
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id,
second_park, _plain_answer("No"))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _return_values(mock_supervisor_comms) == [["Yes", "No"]]
+
+ def test_asking_the_same_key_twice_in_one_run_fails(self,
create_runtime_ti, mock_supervisor_comms):
+ with DAG(dag_id="loop_dag"):
+
+ @task
+ def rounds():
+ count = 0
+ # Bounded so that replaying the same answer forever fails the
test instead of hanging it.
+ while count < 5 and request_input(subject="Another round?",
options=["Yes", "No"])[
+ "chosen_options"
+ ] == ["Yes"]:
+ count += 1
+ return count
+
+ op = rounds().operator
+ ti = create_runtime_ti(task=op)
+ _, park, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id, park,
_plain_answer("Yes"))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, ValueError)
+ assert "'Another round?'" in str(err)
+ assert len(_hitl_requests(mock_supervisor_comms)) == 1
+
+ def
test_a_restored_task_asking_an_answered_subject_again_fails_instead_of_reusing_the_answer(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ # An agent approving two calls to the same tool: each run restores the
saved state and asks
+ # about the call it is waiting on, under the same subject both times.
+ progress: list[str] = []
+ with DAG(dag_id="repeated_subject_dag"):
+
+ @task
+ def agent():
+ if not progress:
+ first = request_input(subject="Approve refund_order?",
options=["Yes", "No"])
+ progress.append(first["chosen_options"][0])
+ second = request_input(subject="Approve refund_order?",
options=["Yes", "No"])
+ return [*progress, second["chosen_options"][0]]
+
+ op = agent().operator
+ ti = create_runtime_ti(task=op)
+ _, park, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id, park,
_plain_answer("Yes"))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, ValueError)
+ assert "pass a key that is unique to each question" in str(err)
+ assert len(_hitl_requests(mock_supervisor_comms)) == 1
+
+ def test_a_subject_that_changes_between_runs_fails_instead_of_asking_again(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ runs: list[int] = []
+ with DAG(dag_id="drifting_subject_dag"):
+
+ @task
+ def refund():
+ runs.append(len(runs) + 1)
+ answer = request_input(subject=f"Refund (asked on run
{runs[-1]})?", options=["Yes", "No"])
+ return answer["chosen_options"][0]
+
+ op = refund().operator
+ ti = create_runtime_ti(task=op)
+ _, park, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id, park,
_plain_answer("Yes"))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, RuntimeError)
+ assert "'Refund (asked on run 1)?' was never read" in str(err)
+ assert len(_hitl_requests(mock_supervisor_comms)) == 1
+
+ def test_finishing_without_reading_the_answer_fails(self,
create_runtime_ti, mock_supervisor_comms):
+ runs: list[int] = []
+ with DAG(dag_id="unread_answer_dag"):
+
+ @task
+ def refund():
+ runs.append(len(runs) + 1)
+ if runs[-1] == 1:
+ request_input(subject="Refund?", options=["Yes", "No"])
+ return "done"
+
+ op = refund().operator
+ ti = create_runtime_ti(task=op)
+ _, park, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id, park,
_plain_answer("Yes"))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, RuntimeError)
+ assert "'Refund?' was never read" in str(err)
+
+ def test_key_matches_an_answer_when_the_subject_changes(self,
create_runtime_ti, mock_supervisor_comms):
+ runs: list[int] = []
+ with DAG(dag_id="changing_subject_dag"):
+
+ @task
+ def refund():
+ runs.append(len(runs) + 1)
+ answer = request_input(
+ subject=f"Refund (asked on run {runs[-1]})?",
options=["Approve", "Reject"], key="refund"
+ )
+ return answer["chosen_options"][0]
+
+ op = refund().operator
+ ti = create_runtime_ti(task=op)
+ _, park, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id, park,
_plain_answer("Approve"))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _return_values(mock_supervisor_comms) == ["Approve"]
+
+ def test_answer_is_refused_when_the_options_changed_between_runs(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ options = [["Approve", "Reject"], ["Approve", "Hold"]]
+ with DAG(dag_id="changing_options_dag"):
+
+ @task
+ def refund():
+ return request_input(subject="Refund?",
options=options.pop(0))["chosen_options"]
+
+ op = refund().operator
+ ti = create_runtime_ti(task=op)
+ _, park, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id, park,
_plain_answer("Reject"))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(err, ValueError)
+ assert "not in ['Approve', 'Hold']" in str(err)
+
+
+class TestTaskflowVariants:
+ def test_async_task_parks_and_resumes(self, create_runtime_ti,
mock_supervisor_comms):
+ with DAG(dag_id="async_dag"):
+
+ @task
+ async def ask():
+ return request_input(subject="Go?", options=["Yes",
"No"])["chosen_options"][0]
+
+ op = ask().operator
+ ti = create_runtime_ti(task=op)
+ state, park_msg, err = _run(ti)
+ assert err is None
+ assert state == TaskInstanceState.AWAITING_INPUT
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id,
park_msg, _answer("Yes"))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _return_values(mock_supervisor_comms) == ["Yes"]
+
+ @pytest.mark.parametrize(
+ ("chosen", "skipped"),
+ [
+ pytest.param("Approve", "reject_path", id="approve"),
+ pytest.param("Reject", "approve_path", id="reject"),
+ ],
+ )
+ def test_branch_applies_the_decision_made_after_resume(
+ self, create_runtime_ti, mock_supervisor_comms, chosen, skipped
+ ):
+ with DAG(dag_id="branch_dag"):
+
+ @task.branch
+ def decide():
+ answer = request_input(subject="Which way?",
options=["Approve", "Reject"])
+ return "approve_path" if answer["chosen_options"] ==
["Approve"] else "reject_path"
+
+ branch = decide()
+ branch >> [EmptyOperator(task_id="approve_path"),
EmptyOperator(task_id="reject_path")]
+ op = branch.operator
+ ti = create_runtime_ti(task=op)
+ state, park_msg, _ = _run(ti)
+ assert state == TaskInstanceState.AWAITING_INPUT
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id,
park_msg, _answer(chosen))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _skipped_task_ids(mock_supervisor_comms) == [skipped]
+
+ def test_short_circuit_applies_the_condition_returned_after_resume(
+ self, create_runtime_ti, mock_supervisor_comms
+ ):
+ with DAG(dag_id="short_circuit_dag"):
+
+ @task.short_circuit
+ def go_on():
+ return request_input(subject="Continue?", options=["Yes",
"No"])["chosen_options"] == ["Yes"]
+
+ gate = go_on()
+ gate >> EmptyOperator(task_id="downstream")
+ op = gate.operator
+ ti = create_runtime_ti(task=op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id,
park_msg, _answer("No"))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _skipped_task_ids(mock_supervisor_comms) == ["downstream"]
+
+ def test_classic_python_operator_parks_and_resumes(self,
create_runtime_ti, mock_supervisor_comms):
+ def ask():
+ return request_input(subject="Go?", options=["Yes",
"No"])["chosen_options"][0]
+
+ op = PythonOperator(task_id="classic", python_callable=ask)
+ ti = create_runtime_ti(task=op)
+ _, park_msg, _ = _run(ti)
+
+ state, _, err = _answer_and_run(create_runtime_ti, op, ti.id,
park_msg, _answer("No"))
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+ assert _return_values(mock_supervisor_comms) == ["No"]
+
+ @pytest.mark.parametrize("map_index", [0, 1])
+ def test_mapped_task_parks_and_resumes_its_own_index(
+ self, create_runtime_ti, mock_supervisor_comms, map_index
+ ):
+ with DAG(dag_id="mapped_dag"):
+
+ @task
+ def ask(order_id: str):
+ request_input(subject=f"Refund {order_id}?", options=["Yes",
"No"])
+ return order_id
+
+ # The worker runs one unmapped copy of a mapped task per index,
with that index's arguments.
+ unmapped = ask(f"o-{map_index}").operator
+ ti = create_runtime_ti(task=unmapped, map_index=map_index)
+ state, park_msg, _ = _run(ti)
+ assert state == TaskInstanceState.AWAITING_INPUT
+ (request,) = _hitl_requests(mock_supervisor_comms)
+ assert request.ti_id == ti.id
+ assert request.subject == f"Refund o-{map_index}?"
+
+ resumed = create_runtime_ti(task=unmapped, ti_id=ti.id,
map_index=map_index)
+ _resume(resumed, park_msg, _answer("Yes"))
+ state, _, err = _run(resumed)
+
+ assert err is None
+ assert state == TaskInstanceState.SUCCESS
+
+
+class TestAnswerLogOutsideTheTaskRunner:
+ """
+ ``run()`` executes each task in a copied context, so a leaked answer log
could not outlive it there.
+
+ These call ``resume_with_input`` directly, the way a caller outside the
task runner would.
+ """
+
+ @staticmethod
+ def _resume_directly(op: PythonOperator, create_runtime_ti):
+ context = create_runtime_ti(task=op).get_template_context()
+ with set_current_context(context):
+ return op.resume_with_input(
+ context,
+ event=_plain_answer("Approve"),
+ key="Go?",
+ options=["Approve", "Reject"],
+ params={},
+ answered=[],
+ )
+
+ def test_answers_are_cleared_after_the_callable_returns(self,
create_runtime_ti, mock_supervisor_comms):
+ op = PythonOperator(
+ task_id="t",
+ python_callable=lambda: request_input(subject="Go?",
options=["Approve", "Reject"])[
+ "chosen_options"
+ ],
+ )
+
+ result = self._resume_directly(op, create_runtime_ti)
+
+ assert result == ["Approve"]
+ assert answer_log_var.get() is None
+
+ def test_answers_are_cleared_after_the_callable_raises(self,
create_runtime_ti, mock_supervisor_comms):
+ def fail_after_reading():
+ answer = request_input(subject="Go?", options=["Approve",
"Reject"])
+ raise RuntimeError(f"failed after {answer['chosen_options']}")
+
+ op = PythonOperator(task_id="t", python_callable=fail_after_reading)
+
+ with pytest.raises(RuntimeError, match=r"failed after \['Approve'\]"):
+ self._resume_directly(op, create_runtime_ti)
+
+ assert answer_log_var.get() is None
+
+
+class TestRequestInputRefusals:
+ @pytest.mark.parametrize(
+ "make_task",
+ [
+ pytest.param(
+ lambda: PythonVirtualenvOperator(task_id="venv",
python_callable=_noop),
+ id="virtualenv",
+ ),
+ pytest.param(lambda: BashOperator(task_id="bash",
bash_command="true"), id="bash"),
+ pytest.param(lambda: EmptyOperator(task_id="empty"),
id="not-python"),
+ ],
+ )
+ def test_refuses_tasks_that_cannot_resume_before_writing_a_request(
+ self, create_runtime_ti, mock_supervisor_comms, make_task
+ ):
+ ti = create_runtime_ti(task=make_task())
+
+ with set_current_context(ti.get_template_context()):
+ with pytest.raises(TypeError, match="cannot resume it"):
+ request_input(subject="Go?", options=["Yes", "No"])
+
+ assert _hitl_requests(mock_supervisor_comms) == []
+
+ def test_refuses_before_airflow_3_4(self, create_runtime_ti,
mock_supervisor_comms, monkeypatch):
+
monkeypatch.setattr("airflow.providers.standard.operators.hitl.AIRFLOW_V_3_4_PLUS",
False)
+ ti = create_runtime_ti(task=PythonOperator(task_id="t",
python_callable=lambda: None))
+
+ with set_current_context(ti.get_template_context()):
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match="needs Airflow 3.4"):
+ request_input(subject="Go?", options=["Yes", "No"])
+
+ assert _hitl_requests(mock_supervisor_comms) == []
+
+ @pytest.mark.parametrize(
+ ("kwargs", "match"),
+ [
+ pytest.param({"options": []}, "cannot be empty", id="no-options"),
+ pytest.param({"options": ["Yes"], "defaults": "No"}, "subset",
id="default-not-an-option"),
+ pytest.param(
+ {"options": ["Yes", "No"], "defaults": ["Yes", "No"]},
+ "More than one defaults",
+ id="two-defaults-single-choice",
+ ),
+ pytest.param(
+ {"options": ["Yes"], "params": {"_options": 1}}, "_options",
id="reserved-param-name"
+ ),
+ ],
+ )
+ def test_refuses_invalid_requests(self, create_runtime_ti,
mock_supervisor_comms, kwargs, match):
+ ti = create_runtime_ti(task=PythonOperator(task_id="t",
python_callable=lambda: None))
+
+ with set_current_context(ti.get_template_context()):
+ with pytest.raises(ValueError, match=match):
+ request_input(subject="Go?", **kwargs)
+
+ assert _hitl_requests(mock_supervisor_comms) == []
diff --git a/providers/standard/tests/unit/standard/utils/test_hitl.py
b/providers/standard/tests/unit/standard/utils/test_hitl.py
new file mode 100644
index 00000000000..a59a2baa827
--- /dev/null
+++ b/providers/standard/tests/unit/standard/utils/test_hitl.py
@@ -0,0 +1,97 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import pytest
+
+from airflow.providers.common.compat.sdk import Param, ParamValidationError
+from airflow.providers.standard.exceptions import HITLTimeoutError,
HITLTriggerEventError
+from airflow.providers.standard.utils.hitl import (
+ params_from_serialized,
+ raise_for_error_event,
+ validate_params_input,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_3_PLUS
+
+SERIALIZED_PARAMS = {
+ "reason": {"value": "", "description": "Why", "schema": {"type":
["string", "null"]}, "source": "task"},
+ "amount": {"value": 1, "description": None, "schema": {"type": "integer",
"minimum": 1}, "source": None},
+}
+
+
[email protected](not AIRFLOW_V_3_1_3_PLUS, reason="Param.serialize() form
is 3.1.3+")
+def test_params_from_serialized_round_trips_param_serialize():
+ original = {"reason": Param("", type=["string", "null"],
description="Why")}
+
+ rebuilt = params_from_serialized({key: param.serialize() for key, param in
original.items()})
+
+ assert rebuilt.get_param("reason").serialize() ==
original["reason"].serialize()
+
+
+def test_params_from_serialized_keeps_each_schema():
+ params = params_from_serialized(SERIALIZED_PARAMS)
+
+ assert params.get_param("amount").schema == {"type": "integer", "minimum":
1}
+ assert params.get_param("reason").description == "Why"
+ assert params["reason"] == ""
+
+
+def test_validate_params_input_writes_valid_values():
+ params = params_from_serialized(SERIALIZED_PARAMS)
+
+ validate_params_input({"reason": "over limit", "amount": 5}, params)
+
+ assert params["reason"] == "over limit"
+ assert params["amount"] == 5
+
+
[email protected](
+ ("params_input", "error", "match"),
+ [
+ pytest.param({"reason": "x"}, ValueError, "does not match params",
id="missing-field"),
+ pytest.param(
+ {"reason": "x", "amount": 1, "note": "y"}, ValueError, "does not
match params", id="unknown-field"
+ ),
+ pytest.param(
+ {"reason": "x", "amount": 0},
+ ParamValidationError,
+ "Invalid input for param amount",
+ id="fails-schema",
+ ),
+ ],
+)
+def test_validate_params_input_refuses_input_that_does_not_fit(params_input,
error, match):
+ params = params_from_serialized(SERIALIZED_PARAMS)
+
+ with pytest.raises(error, match=match):
+ validate_params_input(params_input, params)
+
+
[email protected](
+ ("event", "error"),
+ [
+ pytest.param({"error": "no answer", "error_type": "timeout"},
HITLTimeoutError, id="timeout"),
+ pytest.param({"error": "boom", "error_type": "unknown"},
HITLTriggerEventError, id="other"),
+ ],
+)
+def test_raise_for_error_event_maps_the_error_type(event, error):
+ with pytest.raises(error) as exc_info:
+ raise_for_error_event(event)
+
+ # HITLTimeoutError subclasses HITLTriggerEventError, so check the exact
class.
+ assert type(exc_info.value) is error
