This is an automated email from the ASF dual-hosted git repository.
guan404ming 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 2a1ba5e960f Support timeout defaults in LLM approval reviews (#72155)
2a1ba5e960f is described below
commit 2a1ba5e960f555483a31e8779e0dae4be3e9805c
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Wed Sep 16 22:47:08 2026 +0800
Support timeout defaults in LLM approval reviews (#72155)
* Support timeout defaults in LLM approval reviews
The approval review always passed no HITL defaults, so approval_timeout
could only ever fail the task. An unattended pipeline had no way to say
"nobody reviewed this in time, treat it as approved" (or rejected), which
is what a deadline on a review is usually for.
* Address review on LLM approval timeout defaults
Dropping the defer timeout keeps the trigger deadline as the single
authority on pre-3.3 cores, so the scheduler sweep cannot fail the task
before the timeout default is applied. A timed-out review has no
reviewer, so stale params from an earlier attempt must not be returned.
The guard now also rejects on_approval_timeout without require_approval.
* Clarify timeout defaults and guard replayed HITL fallback
"reject" only fails plain LLM operators, so the docs and example now
show "approve" as the value that keeps a pipeline moving. A replayed
HITLTrigger after its own timeout fallback has no responder, which
crashed the trigger on cores below 3.3. Timeout-driven outcomes are
now named in the exception message and in the approval log.
* Name the actor behind unattended HITL approvals
A replayed HITL trigger that resumed with a stored timeout default logged
that the fallback was skipped, so an on-call reader could not tell an
unattended approval from a human one. The common.ai reject and approve
messages also printed the raw HITLUser dict instead of the reviewer name,
and the standard provider floor let a common.ai upgrade pair the new
on_approval_timeout defaults with a trigger that still dereferenced a
missing responder.
* Reject non-positive approval timeouts for timeout defaults
* Fail zero approval timeouts on pre-3.3 cores instead of polling forever
Dropping timeout= from defer() left the trigger's timeout_datetime as the
only deadline on cores older than 3.3. The truthiness check skipped
timedelta(0), so the default on_approval_timeout configuration polled forever
where it used to fail at the deadline. The guard message also claimed the
setting had no effect, but the review form still pre-highlights the option;
only the deadline needs a positive approval_timeout.
---
providers/common/ai/docs/operators/llm.rst | 14 ++-
providers/common/ai/docs/operators/llm_branch.rst | 11 +-
.../common/ai/docs/operators/llm_file_analysis.rst | 4 +-
.../ai/docs/operators/llm_schema_compare.rst | 10 +-
providers/common/ai/docs/operators/llm_sql.rst | 4 +-
providers/common/ai/pyproject.toml | 2 +-
.../common/ai/example_dags/example_llm.py | 1 +
.../airflow/providers/common/ai/mixins/approval.py | 50 +++++++--
.../airflow/providers/common/ai/operators/llm.py | 25 ++++-
.../providers/common/ai/operators/llm_branch.py | 5 +-
.../common/ai/operators/llm_file_analysis.py | 3 +-
.../common/ai/operators/llm_schema_compare.py | 3 +-
.../providers/common/ai/operators/llm_sql.py | 3 +-
.../tests/unit/common/ai/mixins/test_approval.py | 124 +++++++++++++++++----
.../ai/tests/unit/common/ai/operators/test_llm.py | 47 +++++++-
.../unit/common/ai/operators/test_llm_branch.py | 37 ++++--
.../common/ai/operators/test_llm_file_analysis.py | 13 ++-
.../common/ai/operators/test_llm_schema_compare.py | 4 +-
.../tests/unit/common/ai/operators/test_llm_sql.py | 15 ++-
.../airflow/providers/standard/triggers/hitl.py | 71 +++++++-----
.../tests/unit/standard/triggers/test_hitl.py | 43 +++++++
21 files changed, 387 insertions(+), 102 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index b9537e2f301..50f8ae56216 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -238,7 +238,16 @@ Set ``require_approval=True`` to pause the task after the
LLM generates its
output and wait for a human reviewer to approve or reject it via the Airflow
HITL interface. Optionally allow the reviewer to edit the output before
approving with ``allow_modifications=True``, and set a deadline with
-``approval_timeout``:
+``approval_timeout``.
+
+When ``approval_timeout`` expires without a review, the task fails by default.
+Set ``on_approval_timeout="approve"`` to return the generated output instead,
so
+an unattended pipeline keeps moving. ``"reject"`` answers the review with a
+rejection, which still fails this operator; only
+:class:`~airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator`
+turns a rejection into a downstream skip. The chosen option is also
+pre-highlighted as the default in the review form, so ``"reject"`` makes
+Reject the primary button:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
:language: python
@@ -265,6 +274,9 @@ Parameters
for human review. Default ``False``.
- ``approval_timeout``: Maximum time to wait for a review (``timedelta``).
``None``
means wait indefinitely. Default ``None``.
+- ``on_approval_timeout``: Outcome when ``approval_timeout`` expires without a
+ review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires
+ ``require_approval=True`` and a positive ``approval_timeout``.
- ``allow_modifications``: If ``True``, the reviewer can edit the output before
approving. Default ``False``.
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index d1f11a7309c..eba21bd3e55 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -102,15 +102,17 @@ task on rejection instead (generally discouraged), or
``ignore_downstream_trigger_rules=True`` to skip every downstream task rather
than only the direct ones, so a task whose trigger rule would still run it is
skipped too. Letting ``approval_timeout`` expire fails the task
-(``HITLTimeoutError``).
+(``HITLTimeoutError``) unless ``on_approval_timeout`` answers the review for
+you; a timeout-driven rejection then skips downstream like any other rejection.
``require_approval=True`` requires a string prompt: a decorated callable
returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM
call.
Apart from ``fail_on_reject`` and ``ignore_downstream_trigger_rules``, which
-are specific to this operator, ``approval_timeout`` and the rest of the
-approval behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.
+are specific to this operator, ``approval_timeout``, ``on_approval_timeout``,
+and the rest of the approval behaviour are inherited from
+:ref:`LLMOperator <howto/operator:llm>`.
How It Works
------------
@@ -140,6 +142,9 @@ Parameters
branch(es) and waits for human review before branching. Default ``False``.
- ``approval_timeout``: Maximum time to wait for a review (``timedelta``).
``None``
means wait indefinitely. Default ``None``.
+- ``on_approval_timeout``: Outcome when ``approval_timeout`` expires without a
+ review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires
+ ``require_approval=True`` and a positive ``approval_timeout``.
- ``allow_modifications``: If ``True``, the reviewer can change the chosen
branch(es) before approving. Default ``False``.
- ``fail_on_reject``: If ``True``, a rejected review fails the task instead of
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index ac688ebfaf4..dc65e31bde0 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -159,8 +159,8 @@ Parameters
downstream consumer needs the dict shape.
This operator also inherits ``LLMOperator``'s HITL review parameters --
-``require_approval``, ``approval_timeout``, and ``allow_modifications`` -- see
-:doc:`llm` for details.
+``require_approval``, ``approval_timeout``, ``on_approval_timeout``, and
+``allow_modifications`` -- see :doc:`llm` for details.
Supported Formats
-----------------
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index 768d96280be..a8c3c29deec 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -121,7 +121,7 @@ Set ``require_approval=True`` to pause the task after the
comparison and wait
for a human reviewer to approve the result before it is returned. The review
body shows the compatibility verdict, a mismatch severity summary, and the
full result JSON. Rejecting the review, or letting ``approval_timeout``
-expire, fails the task:
+expire with the default ``on_approval_timeout="fail"``, fails the task:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
:language: python
@@ -132,8 +132,9 @@ expire, fails the task:
returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM
call.
-``approval_timeout``, ``allow_modifications``, and the rest of the approval
-behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.
+``approval_timeout``, ``on_approval_timeout``, ``allow_modifications``, and
+the rest of the approval behaviour are inherited from
+:ref:`LLMOperator <howto/operator:llm>`.
Conditional ETL Based on Schema Compatibility
----------------------------------------------
@@ -191,6 +192,9 @@ Parameters
waits for human review before returning the result. Default ``False``.
- ``approval_timeout``: Maximum time to wait for a review (``timedelta``).
``None``
means wait indefinitely. Default ``None``.
+- ``on_approval_timeout``: Outcome when ``approval_timeout`` expires without a
+ review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires
+ ``require_approval=True`` and a positive ``approval_timeout``.
- ``allow_modifications``: If ``True``, the reviewer can edit the result JSON
before approving. Default ``False``.
diff --git a/providers/common/ai/docs/operators/llm_sql.rst
b/providers/common/ai/docs/operators/llm_sql.rst
index 233f7262a95..bb328f1e67f 100644
--- a/providers/common/ai/docs/operators/llm_sql.rst
+++ b/providers/common/ai/docs/operators/llm_sql.rst
@@ -141,7 +141,9 @@ Human-in-the-Loop Approval
Set ``require_approval=True`` to pause the task after SQL generation and wait
for a human reviewer to approve the query before it is returned.
When ``allow_modifications=True``, the reviewer can also edit the SQL — the
-modified query is re-validated against the same safety rules automatically:
+modified query is re-validated against the same safety rules automatically.
+``approval_timeout`` and ``on_approval_timeout`` behave as on
+:ref:`LLMOperator <howto/operator:llm>`:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
:language: python
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index 78dbde5c959..a27011e7f69 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -69,7 +69,7 @@ requires-python = ">=3.10"
dependencies = [
"apache-airflow>=3.0.0",
"apache-airflow-providers-common-compat>=1.15.0",
- "apache-airflow-providers-standard>=1.12.1",
+ "apache-airflow-providers-standard>=1.12.1", # use next version
# Requires the pydantic-ai cost API (RunUsage.cost,
UsageLimits.cost_limit),
# landed in 2.23.0 via https://github.com/pydantic/pydantic-ai/pull/2684.
"pydantic-ai-slim>=2.23.0",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py
index 833d4b09645..ee1a678ce89 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py
@@ -190,6 +190,7 @@ def example_llm_operator_approval():
system_prompt="You are a financial analyst. Be concise and accurate.",
require_approval=True,
approval_timeout=timedelta(hours=24),
+ on_approval_timeout="approve",
allow_modifications=True,
)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
index cccff32bfd1..3153eb3a999 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
@@ -20,7 +20,7 @@ from __future__ import annotations
import json
import logging
from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol
from pydantic import BaseModel, TypeAdapter
@@ -42,6 +42,7 @@ class DeferForApprovalProtocol(Protocol):
approval_timeout: timedelta | None
allow_modifications: bool
+ on_approval_timeout: Literal["fail", "approve", "reject"]
prompt: str
task_id: str
defer: Any
@@ -62,16 +63,24 @@ class LLMApprovalMixin:
before approving. The (possibly modified) output is then returned as the
task result.
+ ``on_approval_timeout`` decides what happens when ``approval_timeout``
+ expires without a response: ``"fail"`` raises ``HITLTimeoutError``, while
+ ``"approve"`` and ``"reject"`` answer the review with that option so the
+ task resumes as if a reviewer had chosen it. The chosen option is also
+ pre-highlighted for the reviewer in the HITL form.
+
Operators that use this mixin must set the following attributes:
- ``require_approval`` (``bool``)
- ``allow_modifications`` (``bool``)
- ``approval_timeout`` (``timedelta | None``)
+ - ``on_approval_timeout`` (``Literal["fail", "approve", "reject"]``)
- ``prompt`` (``str``)
"""
APPROVE = "Approve"
REJECT = "Reject"
+ TIMEOUT_DEFAULTS: ClassVar[dict[str, list[str]]] = {"approve": [APPROVE],
"reject": [REJECT]}
def validate_approval_prompt(self: DeferForApprovalProtocol) -> None:
"""Fail fast when the prompt cannot be rendered as text in the
approval review body."""
@@ -99,7 +108,8 @@ class LLMApprovalMixin:
On Airflow 3.3+ the task parks in the ``awaiting_input`` state (no
trigger or triggerer
involved); on older versions it defers to :class:`HITLTrigger`. Either
way it resumes in
- ``execute_complete`` once a response (or timeout default) arrives.
+ ``execute_complete`` once a response (or timeout default) arrives.
``on_approval_timeout``
+ supplies that timeout default; ``"fail"`` supplies none, so the review
times out as an error.
:param context: Airflow task context.
:param output: The generated output to present for review.
@@ -130,6 +140,7 @@ class LLMApprovalMixin:
output = TypeAdapter(type(output)).dump_json(output).decode()
ti_id = context["task_instance"].id
+ timeout_defaults =
LLMApprovalMixin.TIMEOUT_DEFAULTS.get(self.on_approval_timeout)
if subject is None:
subject = f"Review output for task `{self.task_id}`"
@@ -160,7 +171,7 @@ class LLMApprovalMixin:
options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT],
subject=subject,
body=body,
- defaults=None,
+ defaults=timeout_defaults,
multiple=False,
params=hitl_params,
)
@@ -179,16 +190,24 @@ class LLMApprovalMixin:
trigger=HITLTrigger(
ti_id=ti_id,
options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT],
- defaults=None,
+ defaults=timeout_defaults,
params=hitl_params,
multiple=False,
- timeout_datetime=utcnow() + self.approval_timeout if
self.approval_timeout else None,
+ timeout_datetime=(
+ utcnow() + self.approval_timeout if self.approval_timeout
is not None else None
+ ),
),
method_name="execute_complete",
kwargs={"generated_output": output},
- timeout=self.approval_timeout,
)
+ @staticmethod
+ def _describe_responder(event: dict[str, Any]) -> str:
+ responded_by_user = event.get("responded_by_user")
+ if responded_by_user is None:
+ return "the approval timeout default"
+ return responded_by_user["name"]
+
def execute_complete(self, context: Context, generated_output: str, event:
dict[str, Any]) -> str:
"""
Resume after human review.
@@ -199,8 +218,9 @@ class LLMApprovalMixin:
:param context: Airflow task context.
:param generated_output: The output that was deferred for review.
:param event: Trigger event payload containing ``chosen_options``,
- ``params_input``, and ``responded_by_user``.
- :raises HITLRejectException: If the reviewer rejected the output.
+ ``params_input``, ``responded_by_user``, and ``timedout``.
+ :raises HITLRejectException: If the reviewer, or the
+ ``on_approval_timeout="reject"`` default, rejected the output.
:raises HITLTriggerEventError: If the trigger reported an error.
:raises HITLTimeoutError: If the approval timed out.
"""
@@ -216,13 +236,19 @@ class LLMApprovalMixin:
raise HITLTimeoutError(f"Approval timed out: {event['error']}")
raise HITLTriggerEventError(event)
- responded_by_user = event.get("responded_by_user")
+ responder = self._describe_responder(event)
chosen = event["chosen_options"]
if self.APPROVE not in chosen:
- raise HITLRejectException(f"Output was rejected by the reviewer
{responded_by_user}.")
+ if event.get("timedout"):
+ raise HITLRejectException(
+ "Output was rejected automatically: approval_timeout
expired with "
+ "on_approval_timeout='reject'."
+ )
+ raise HITLRejectException(f"Output was rejected by the reviewer
{responder}.")
+ log.info("Output approved by %s.", responder)
output = generated_output
- params_input: dict[str, Any] = event.get("params_input") or {}
+ params_input: dict[str, Any] = {} if event.get("timedout") else
event.get("params_input") or {}
# Only accept modified output when the operator explicitly allows
modifications.
# Without this guard a reviewer could craft a request with
params_input even
@@ -260,7 +286,7 @@ class LLMApprovalMixin:
}
)
if modified is not None and modified != generated_output:
- log.info("output=%s modified by the reviewer=%s ", modified,
responded_by_user)
+ log.info("output=%s modified by the reviewer=%s ", modified,
responder)
return modified
return output
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index 6b28dbf7232..2de3762de6a 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -21,7 +21,7 @@ from __future__ import annotations
from collections.abc import Sequence
from datetime import timedelta
from functools import cached_property
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
from pydantic import BaseModel
@@ -100,7 +100,14 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
output and waits for a human reviewer to approve or reject via the
HITL interface. Default ``False``.
:param approval_timeout: Maximum time to wait for a review. When
- exceeded, the task fails with ``TimeoutError``.
+ exceeded, ``on_approval_timeout`` decides the outcome.
+ :param on_approval_timeout: What to do when ``approval_timeout`` expires
+ without a review. ``"fail"`` (default) fails the task with
+ ``HITLTimeoutError``; ``"approve"`` and ``"reject"`` answer the review
+ with that option, so the task resumes as if a reviewer had chosen it.
+ The chosen option is also pre-highlighted for the reviewer in the HITL
+ form. Requires ``require_approval=True`` and a positive
+ ``approval_timeout``.
:param allow_modifications: If ``True``, the reviewer can edit the output
before approving. The modified value is returned as the task result.
Default ``False``.
@@ -135,6 +142,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
usage_limits: UsageLimits | dict[str, Any] | None = None,
require_approval: bool = False,
approval_timeout: timedelta | None = None,
+ on_approval_timeout: Literal["fail", "approve", "reject"] = "fail",
allow_modifications: bool = False,
serialize_output: bool = False,
**kwargs: Any,
@@ -153,8 +161,21 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
self.agent_params = agent_params or {}
# No validation here -- see coerce_usage_limits() docstring for why.
self.usage_limits = usage_limits
+ if on_approval_timeout not in ("fail",
*LLMApprovalMixin.TIMEOUT_DEFAULTS):
+ raise ValueError(
+ f"on_approval_timeout must be 'fail', 'approve', or 'reject',
got {on_approval_timeout!r}."
+ )
+ if on_approval_timeout != "fail" and not (
+ require_approval and approval_timeout is not None and
approval_timeout > timedelta(0)
+ ):
+ raise ValueError(
+ f"on_approval_timeout={on_approval_timeout!r} needs
require_approval=True and "
+ "a positive approval_timeout to fire. "
+ "Set both, or leave on_approval_timeout as 'fail'."
+ )
self.require_approval = require_approval
self.approval_timeout = approval_timeout
+ self.on_approval_timeout = on_approval_timeout
self.allow_modifications = allow_modifications
@cached_property
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
index 1cec2f83e4d..e646bc0d8c8 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
@@ -61,7 +61,8 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
- (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
+ ``allow_modifications``).
The task pauses after the LLM chooses the branch(es) and only skips the
unselected downstream tasks once a reviewer approves. Rejecting the
review skips the direct downstream tasks except teardowns, matching
@@ -160,7 +161,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
except HITLRejectException:
if self.fail_on_reject:
raise
- self.log.info("Rejected by %s. Skipping downstream tasks...",
event.get("responded_by_user"))
+ self.log.info("Rejected by %s. Skipping downstream tasks...",
self._describe_responder(event))
task = context["task"]
tasks = (
task.get_flat_relatives(upstream=False)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
index 7767e5a534d..c671587a3d0 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
@@ -72,7 +72,8 @@ class LLMFileAnalysisOperator(LLMOperator):
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
- (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
+ ``allow_modifications``).
The task pauses after the file analysis and only returns the result once a
reviewer approves.
"""
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
index 8fc82982d16..9e5127f8da9 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
@@ -111,7 +111,8 @@ class LLMSchemaCompareOperator(LLMOperator):
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
- (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
+ ``allow_modifications``).
The task pauses after the comparison and only returns the result once a
reviewer approves. The review body shows the compatibility verdict, a
mismatch severity summary, and the full result JSON.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
index 2ac0cbe40dc..87a81e0882c 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
@@ -84,7 +84,8 @@ class LLMSQLQueryOperator(LLMOperator):
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
- (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
+ ``allow_modifications``).
When ``allow_modifications=True`` and the reviewer edits the SQL, the
modified query is re-validated against the same safety rules before being
returned.
diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
index ace3ce7d3d9..6c9f55224a6 100644
--- a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
+++ b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
@@ -52,11 +52,13 @@ class FakeOperator(LLMApprovalMixin):
prompt: str = "Summarize this",
task_id: str = "test_task",
approval_timeout: timedelta | None = None,
+ on_approval_timeout: str = "fail",
allow_modifications: bool = False,
):
self.prompt = prompt
self.task_id = task_id
self.approval_timeout = approval_timeout
+ self.on_approval_timeout = on_approval_timeout
self.allow_modifications = allow_modifications
self.defer = MagicMock()
@@ -203,8 +205,39 @@ class TestDeferForApproval:
trigger_call_kwargs = mock_trigger_cls.call_args[1]
assert trigger_call_kwargs["timeout_datetime"] == fake_now + timeout
- defer_kwargs = op.defer.call_args[1]
- assert defer_kwargs["timeout"] == timeout
+ assert "timeout" not in op.defer.call_args[1]
+
+ @patch(UTCNOW_PATH)
+ @patch(HITL_TRIGGER_PATH, autospec=True)
+ @patch(UPSERT_HITL_PATH)
+ def test_zero_timeout_sets_timeout_datetime_to_now(
+ self, mock_upsert, mock_trigger_cls, mock_utcnow, context
+ ):
+ from datetime import datetime
+
+ fake_now = datetime(2025, 1, 1, 12, 0, 0)
+ mock_utcnow.return_value = fake_now
+ op = FakeOperator(approval_timeout=timedelta(0))
+
+ op.defer_for_approval(context, "output")
+
+ assert mock_trigger_cls.call_args[1]["timeout_datetime"] == fake_now
+
+ @pytest.mark.parametrize(
+ ("on_approval_timeout", "expected_defaults"),
+ [("fail", None), ("approve", ["Approve"]), ("reject", ["Reject"])],
+ )
+ @patch(HITL_TRIGGER_PATH, autospec=True)
+ @patch(UPSERT_HITL_PATH)
+ def test_on_approval_timeout_sets_hitl_defaults(
+ self, mock_upsert, mock_trigger_cls, context, on_approval_timeout,
expected_defaults
+ ):
+ op = FakeOperator(approval_timeout=timedelta(hours=1),
on_approval_timeout=on_approval_timeout)
+
+ op.defer_for_approval(context, "output")
+
+ assert mock_upsert.call_args[1]["defaults"] == expected_defaults
+ assert mock_trigger_cls.call_args[1]["defaults"] == expected_defaults
@patch(HITL_TRIGGER_PATH, autospec=True)
@patch(UPSERT_HITL_PATH)
@@ -213,9 +246,7 @@ class TestDeferForApproval:
trigger_call_kwargs = mock_trigger_cls.call_args[1]
assert trigger_call_kwargs["timeout_datetime"] is None
-
- defer_kwargs = approval_op.defer.call_args[1]
- assert defer_kwargs["timeout"] is None
+ assert "timeout" not in approval_op.defer.call_args[1]
@patch(HITL_TRIGGER_PATH, autospec=True)
@patch(UPSERT_HITL_PATH)
@@ -248,20 +279,20 @@ class TestDeferForApproval:
assert "Paris is the capital of France." in call_kwargs["body"]
def test_approved_returns_generated_output(self, approval_op):
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
result = approval_op.execute_complete({}, generated_output="hello
world", event=event)
assert result == "hello world"
def test_rejected_raises_rejection_exception(self, approval_op):
- event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "admin"}}
with pytest.raises(HITLRejectException, match="Output was rejected by
the reviewer admin."):
approval_op.execute_complete({}, generated_output="output",
event=event)
def test_empty_chosen_options_raises_rejection(self, approval_op):
- event = {"chosen_options": [], "responded_by_user": "admin"}
+ event = {"chosen_options": [], "responded_by_user": {"id": "u1",
"name": "admin"}}
with pytest.raises(HITLRejectException, match="Output was rejected by
the reviewer admin."):
approval_op.execute_complete({}, generated_output="output",
event=event)
@@ -283,7 +314,7 @@ class TestDeferForApproval:
def test_approved_with_modified_output(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": "modified output"},
}
@@ -299,7 +330,7 @@ class TestDeferForApproval:
# string contract instead of returning a dict as the task's output.
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": {"sneaky": "dict"}},
}
@@ -311,7 +342,7 @@ class TestDeferForApproval:
def test_approved_with_list_modified_output_is_serialized(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": ["task_b", "task_c"]},
}
@@ -324,7 +355,7 @@ class TestDeferForApproval:
def test_approved_with_unmodified_list_output_returns_original(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": ["task_a"]},
}
@@ -337,7 +368,7 @@ class TestDeferForApproval:
def test_approved_with_non_string_list_items_raises(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": ["task_a", 2]},
}
@@ -347,7 +378,7 @@ class TestDeferForApproval:
def test_approved_with_cleared_output_raises(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": None},
}
@@ -357,7 +388,7 @@ class TestDeferForApproval:
def test_approved_with_unmodified_output(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": "same output"},
}
@@ -370,7 +401,7 @@ class TestDeferForApproval:
def test_approved_modifications_allowed_but_no_params_input(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": None,
}
@@ -381,7 +412,7 @@ class TestDeferForApproval:
def test_approved_modifications_allowed_empty_output_key(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": "original"},
}
@@ -393,7 +424,7 @@ class TestDeferForApproval:
"""When allow_modifications=False, params will be empty so
params_input is empty too."""
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {},
}
@@ -405,7 +436,7 @@ class TestDeferForApproval:
"""When allow_modifications=False, tampered params_input with output
must be ignored."""
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "reviewer",
+ "responded_by_user": {"id": "u1", "name": "reviewer"},
"params_input": {"output": "tampered output"},
}
@@ -420,8 +451,46 @@ class TestDeferForApproval:
assert result == "output"
+ def test_timed_out_rejection_names_the_timeout_default(self, approval_op):
+ event = {"chosen_options": ["Reject"], "responded_by_user": None,
"timedout": True}
+
+ with pytest.raises(
+ HITLRejectException,
+ match="Output was rejected automatically: approval_timeout expired
with on_approval_timeout='reject'.",
+ ):
+ approval_op.execute_complete({}, generated_output="output",
event=event)
+
+ @pytest.mark.parametrize(
+ ("event", "expected_approver"),
+ [
+ ({"chosen_options": ["Approve"], "responded_by_user": {"id": "u1",
"name": "admin"}}, "admin"),
+ (
+ {"chosen_options": ["Approve"], "responded_by_user": None,
"timedout": True},
+ "the approval timeout default",
+ ),
+ ],
+ ids=["reviewer", "timeout_default"],
+ )
+ @patch("airflow.providers.common.ai.mixins.approval.log", autospec=True)
+ def test_approval_logs_who_approved(self, mock_log, approval_op, event,
expected_approver):
+ approval_op.execute_complete({}, generated_output="output",
event=event)
+
+ mock_log.info.assert_called_once_with("Output approved by %s.",
expected_approver)
+
+ def test_timed_out_approval_ignores_stale_params_input(self,
approval_op_with_modifications):
+ event = {
+ "chosen_options": ["Approve"],
+ "params_input": {"output": "stale output"},
+ "responded_by_user": None,
+ "timedout": True,
+ }
+
+ result = approval_op_with_modifications.execute_complete({},
generated_output="output", event=event)
+
+ assert result == "output"
+
def test_rejection_message_includes_username(self, approval_op):
- event = {"chosen_options": ["Reject"], "responded_by_user": "alice"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "alice"}}
with pytest.raises(HITLRejectException, match="alice"):
approval_op.execute_complete({}, generated_output="output",
event=event)
@@ -453,6 +522,21 @@ class TestAwaitInputForApproval:
assert exc_info.value.timeout == timeout
+ @pytest.mark.parametrize(
+ ("on_approval_timeout", "expected_defaults"),
+ [("fail", None), ("approve", ["Approve"]), ("reject", ["Reject"])],
+ )
+ @patch(UPSERT_HITL_PATH)
+ def test_on_approval_timeout_sets_hitl_defaults_on_await(
+ self, mock_upsert, context, on_approval_timeout, expected_defaults
+ ):
+ op = FakeOperator(approval_timeout=timedelta(hours=1),
on_approval_timeout=on_approval_timeout)
+
+ with pytest.raises(TaskAwaitingInput):
+ op.defer_for_approval(context, "output")
+
+ assert mock_upsert.call_args[1]["defaults"] == expected_defaults
+
@patch(UPSERT_HITL_PATH)
def test_pydantic_output_stringified_on_await(self, mock_upsert,
approval_op, context):
class Answer(BaseModel):
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index 95f0e569672..5ed69546493 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -326,6 +326,38 @@ class TestLLMOperatorApproval:
assert op.require_approval is False
assert op.allow_modifications is False
assert op.approval_timeout is None
+ assert op.on_approval_timeout == "fail"
+
+ def test_unknown_on_approval_timeout_raises(self):
+ with pytest.raises(ValueError, match="on_approval_timeout must be"):
+ LLMOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ approval_timeout=timedelta(hours=1),
+ on_approval_timeout="skip",
+ )
+
+ @pytest.mark.parametrize(
+ "kwargs",
+ [
+ {"require_approval": True},
+ {"require_approval": True, "approval_timeout": timedelta(0)},
+ {"require_approval": True, "approval_timeout":
timedelta(hours=-1)},
+ {"approval_timeout": timedelta(hours=1)},
+ ],
+ ids=[
+ "no_approval_timeout",
+ "zero_approval_timeout",
+ "negative_approval_timeout",
+ "no_require_approval",
+ ],
+ )
+ def test_on_approval_timeout_without_prerequisites_raises(self, kwargs):
+ with pytest.raises(
+ ValueError, match="needs require_approval=True and a positive
approval_timeout to fire"
+ ):
+ LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
on_approval_timeout="approve", **kwargs)
@patch("airflow.providers.standard.triggers.hitl.HITLTrigger",
autospec=True)
@patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")
@@ -430,7 +462,10 @@ class TestLLMOperatorApproval:
with pytest.raises(ApprovalPauseSignal) as exc_info:
op.execute(context=ctx)
- assert exc_info.value.timeout == timeout
+ if AIRFLOW_V_3_3_PLUS:
+ assert exc_info.value.timeout == timeout
+ else:
+ assert mock_trigger_cls.call_args[1]["timeout_datetime"] is not
None
@patch("airflow.providers.standard.triggers.hitl.HITLTrigger",
autospec=True)
@patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")
@@ -472,7 +507,7 @@ class TestLLMOperatorApproval:
def test_execute_complete_approved(self):
"""execute_complete returns output when approved."""
op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
result = op.execute_complete({}, generated_output="the output",
event=event)
@@ -483,7 +518,7 @@ class TestLLMOperatorApproval:
from airflow.providers.standard.exceptions import HITLRejectException
op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
- event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "admin"}}
with pytest.raises(HITLRejectException):
op.execute_complete({}, generated_output="output", event=event)
@@ -503,7 +538,7 @@ class TestLLMOperatorApproval:
op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_modifications=True)
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "editor",
+ "responded_by_user": {"id": "u1", "name": "editor"},
"params_input": {"output": "edited"},
}
@@ -515,7 +550,7 @@ class TestLLMOperatorApproval:
def test_execute_complete_rehydrates_pydantic_for_structured_output(self):
"""When output_type is a BaseModel, execute_complete returns the
model, not the JSON string."""
op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=Summary)
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
result = op.execute_complete({}, generated_output='{"text":"hello"}',
event=event)
@@ -535,7 +570,7 @@ class TestLLMOperatorApproval:
op = LLMOperator(
task_id="t", prompt="p", llm_conn_id="c", output_type=output_type,
require_approval=True
)
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
result = op.execute_complete({}, generated_output=generated_output,
event=event)
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index d6f05dd2b7d..a4cdaee9879 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
@@ -396,7 +396,7 @@ class TestLLMBranchOperatorApproval:
mock_do_branch.return_value = "task_a"
op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c")
op.downstream_task_ids = {"task_a", "task_b"}
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
ctx = _make_context()
result = op.execute_complete(ctx, generated_output="task_a",
event=event)
@@ -410,7 +410,7 @@ class TestLLMBranchOperatorApproval:
mock_do_branch.return_value = ["task_a", "task_c"]
op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_multiple_branches=True)
op.downstream_task_ids = {"task_a", "task_b", "task_c"}
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
ctx = _make_context()
result = op.execute_complete(ctx,
generated_output='["task_a","task_c"]', event=event)
@@ -447,7 +447,7 @@ class TestLLMBranchOperatorApproval:
if with_teardown:
op4.as_teardown()
op1 >> op2 >> op3 >> op4
- event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "admin"}}
ti = MagicMock()
ctx = MagicMock(**{"__getitem__": lambda self, key: {"task": op1,
"ti": ti}[key]})
@@ -459,11 +459,30 @@ class TestLLMBranchOperatorApproval:
assert {t.task_id for t in mock_skip.call_args.kwargs["tasks"]} ==
expected
mock_do_branch.assert_not_called()
+ @patch.object(LLMBranchOperator, "log")
+ @patch.object(LLMBranchOperator, "skip")
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_timed_out_reject_names_the_timeout_default(
+ self, mock_do_branch, mock_skip, mock_log
+ ):
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c")
+ op.downstream_task_ids = {"task_a"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": None,
"timedout": True}
+ task = MagicMock()
+ task.get_direct_relatives.return_value = []
+ ctx = MagicMock(**{"__getitem__": lambda self, key: {"task": task,
"ti": MagicMock()}[key]})
+
+ op.execute_complete(ctx, generated_output="task_a", event=event)
+
+ mock_log.info.assert_called_once_with(
+ "Rejected by %s. Skipping downstream tasks...", "the approval
timeout default"
+ )
+
@patch.object(LLMBranchOperator, "do_branch")
def test_execute_complete_reject_fails_with_fail_on_reject(self,
mock_do_branch):
op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c",
fail_on_reject=True)
op.downstream_task_ids = {"task_a", "task_b"}
- event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "admin"}}
with pytest.raises(HITLRejectException, match="rejected"):
op.execute_complete(_make_context(), generated_output="task_a",
event=event)
@@ -478,7 +497,7 @@ class TestLLMBranchOperatorApproval:
op.downstream_task_ids = {"task_a", "task_b"}
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "admin",
+ "responded_by_user": {"id": "u1", "name": "admin"},
"params_input": {"output": "task_b"},
}
ctx = _make_context()
@@ -502,7 +521,7 @@ class TestLLMBranchOperatorApproval:
op.downstream_task_ids = {"task_a", "task_b", "task_c"}
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "admin",
+ "responded_by_user": {"id": "u1", "name": "admin"},
"params_input": {"output": ["task_b", "task_c"]},
}
ctx = _make_context()
@@ -519,7 +538,7 @@ class TestLLMBranchOperatorApproval:
op.downstream_task_ids = {"task_a", "task_b"}
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "admin",
+ "responded_by_user": {"id": "u1", "name": "admin"},
"params_input": {"output": "task_x"},
}
@@ -541,7 +560,7 @@ class TestLLMBranchOperatorApproval:
op.downstream_task_ids = {"task_a", "task_b"}
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "admin",
+ "responded_by_user": {"id": "u1", "name": "admin"},
"params_input": {"output": "[]"},
}
@@ -568,7 +587,7 @@ class TestLLMBranchOperatorApproval:
op.downstream_task_ids = {"task_a", "task_b"}
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "admin",
+ "responded_by_user": {"id": "u1", "name": "admin"},
"params_input": {"output": modified},
}
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
index 74d9901ab18..ff511183b16 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
@@ -315,7 +315,11 @@ class TestLLMFileAnalysisOperatorApproval:
output_type=Summary,
require_approval=True,
)
- event = {"chosen_options": [op.APPROVE], "params_input": {},
"responded_by_user": "reviewer"}
+ event = {
+ "chosen_options": [op.APPROVE],
+ "params_input": {},
+ "responded_by_user": {"id": "u1", "name": "reviewer"},
+ }
result = op.execute_complete({}, generated_output='{"findings":["error
spike"]}', event=event)
@@ -336,7 +340,7 @@ class TestLLMFileAnalysisOperatorApproval:
event = {
"chosen_options": [op.APPROVE],
"params_input": {"output": '{"findings":["reviewed output"]}'},
- "responded_by_user": "reviewer",
+ "responded_by_user": {"id": "u1", "name": "reviewer"},
}
result = op.execute_complete({}, generated_output='{"findings":["error
spike"]}', event=event)
@@ -375,4 +379,7 @@ class TestLLMFileAnalysisOperatorApproval:
with pytest.raises(ApprovalPauseSignal) as exc_info:
op.execute(context=_make_context())
- assert exc_info.value.timeout == timeout
+ if AIRFLOW_V_3_3_PLUS:
+ assert exc_info.value.timeout == timeout
+ else:
+ assert mock_trigger_cls.call_args[1]["timeout_datetime"] is not
None
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
index 21ac1a41aef..1443ce54661 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
@@ -668,7 +668,7 @@ class TestLLMSchemaCompareOperatorApproval:
def test_execute_complete_approved_returns_dict(self):
result = SchemaCompareResult(compatible=True, mismatches=[],
summary="All good")
op = LLMSchemaCompareOperator(**_BASE_KWARGS, **self._APPROVAL_KWARGS)
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
resumed = op.execute_complete({},
generated_output=result.model_dump_json(), event=event)
@@ -684,7 +684,7 @@ class TestLLMSchemaCompareOperatorApproval:
op = LLMSchemaCompareOperator(**_BASE_KWARGS, **self._APPROVAL_KWARGS,
allow_modifications=True)
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "admin",
+ "responded_by_user": {"id": "u1", "name": "admin"},
"params_input": {"output": modified},
}
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
index 1742959c519..b96d836c490 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
@@ -718,7 +718,10 @@ class TestLLMSQLQueryOperatorApproval:
with pytest.raises(ApprovalPauseSignal) as exc_info:
op.execute(context=ctx)
- assert exc_info.value.timeout == timeout
+ if AIRFLOW_V_3_3_PLUS:
+ assert exc_info.value.timeout == timeout
+ else:
+ assert mock_trigger_cls.call_args[1]["timeout_datetime"] is not
None
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_without_approval_returns_sql(self, mock_hook_cls,
make_mock_run_result):
@@ -765,7 +768,7 @@ class TestLLMSQLQueryOperatorApproval:
def test_execute_complete_approved(self):
"""execute_complete returns SQL when approved."""
op = LLMSQLQueryOperator(task_id="t", prompt="p", llm_conn_id="c")
- event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
result = op.execute_complete({}, generated_output="SELECT * FROM
orders", event=event)
@@ -774,7 +777,7 @@ class TestLLMSQLQueryOperatorApproval:
def test_execute_complete_rejected(self):
"""execute_complete raises HITLRejectException when SQL is rejected."""
op = LLMSQLQueryOperator(task_id="t", prompt="p", llm_conn_id="c")
- event = {"chosen_options": ["Reject"], "responded_by_user": "dba"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "dba"}}
from airflow.providers.standard.exceptions import HITLRejectException
with pytest.raises(HITLRejectException, match="Output was rejected by
the reviewer"):
@@ -795,7 +798,7 @@ class TestLLMSQLQueryOperatorApproval:
op = LLMSQLQueryOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_modifications=True)
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "dba",
+ "responded_by_user": {"id": "u1", "name": "dba"},
"params_input": {"output": "SELECT id, name FROM users LIMIT 10"},
}
@@ -808,7 +811,7 @@ class TestLLMSQLQueryOperatorApproval:
op = LLMSQLQueryOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_modifications=True)
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "john",
+ "responded_by_user": {"id": "u1", "name": "john"},
"params_input": {"output": "DROP TABLE users"},
}
@@ -820,7 +823,7 @@ class TestLLMSQLQueryOperatorApproval:
op = LLMSQLQueryOperator(task_id="t", prompt="p", llm_conn_id="c")
event = {
"chosen_options": ["Approve"],
- "responded_by_user": "john",
+ "responded_by_user": {"id": "u1", "name": "john"},
"params_input": {},
}
diff --git a/providers/standard/src/airflow/providers/standard/triggers/hitl.py
b/providers/standard/src/airflow/providers/standard/triggers/hitl.py
index d7aff2e389c..a48e7c0226d 100644
--- a/providers/standard/src/airflow/providers/standard/triggers/hitl.py
+++ b/providers/standard/src/airflow/providers/standard/triggers/hitl.py
@@ -41,6 +41,9 @@ from airflow.sdk.execution_time.hitl import (
from airflow.sdk.timezone import utcnow
from airflow.triggers.base import BaseTrigger, TriggerEvent
+if TYPE_CHECKING:
+ from airflow.sdk.api.datamodels._generated import HITLDetailResponse
+
class HITLTriggerEventSuccessPayload(TypedDict, total=False):
"""Minimum required keys for a success Human-in-the-loop TriggerEvent."""
@@ -114,6 +117,12 @@ class HITLTrigger(BaseTrigger):
},
)
+ @staticmethod
+ def _get_responder(resp: HITLDetailResponse) -> HITLUser | None:
+ if resp.responded_by_user is None:
+ return None
+ return HITLUser(id=resp.responded_by_user.id,
name=resp.responded_by_user.name)
+
async def _handle_timeout(self) -> TriggerEvent:
"""Handle HITL timeout logic and yield appropriate event."""
resp = await
sync_to_async(get_hitl_detail_content_detail)(ti_id=self.ti_id)
@@ -121,27 +130,32 @@ class HITLTrigger(BaseTrigger):
# Case 1: Response arrived just before timeout
if resp.response_received and resp.chosen_options:
if TYPE_CHECKING:
- assert resp.responded_by_user is not None
assert resp.responded_at is not None
chosen_options_list = list(resp.chosen_options or [])
- self.log.info(
- "[HITL] responded_by=%s (id=%s) options=%s at %s (timeout
fallback skipped)",
- resp.responded_by_user.name,
- resp.responded_by_user.id,
- chosen_options_list,
- resp.responded_at,
- )
+ responded_by_user = self._get_responder(resp)
+ if responded_by_user is None:
+ self.log.info(
+ "[HITL] resuming with the timeout default %s recorded at
%s "
+ "(applied by an earlier run of this trigger)",
+ chosen_options_list,
+ resp.responded_at,
+ )
+ else:
+ self.log.info(
+ "[HITL] responded_by=%s (id=%s) options=%s at %s (timeout
fallback skipped)",
+ responded_by_user["name"],
+ responded_by_user["id"],
+ chosen_options_list,
+ resp.responded_at,
+ )
return TriggerEvent(
HITLTriggerEventSuccessPayload(
chosen_options=chosen_options_list,
params_input=resp.params_input or {},
responded_at=resp.responded_at,
- responded_by_user=HITLUser(
- id=resp.responded_by_user.id,
- name=resp.responded_by_user.name,
- ),
- timedout=False,
+ responded_by_user=responded_by_user,
+ timedout=responded_by_user is None,
)
)
@@ -181,7 +195,6 @@ class HITLTrigger(BaseTrigger):
"""Check if HITL response is ready and yield success if so."""
resp = await
sync_to_async(get_hitl_detail_content_detail)(ti_id=self.ti_id)
if TYPE_CHECKING:
- assert resp.responded_by_user is not None
assert resp.responded_at is not None
if not (resp.response_received and resp.chosen_options):
@@ -201,23 +214,29 @@ class HITLTrigger(BaseTrigger):
)
chosen_options_list = list(resp.chosen_options or [])
- self.log.info(
- "[HITL] responded_by=%s (id=%s) options=%s at %s",
- resp.responded_by_user.name,
- resp.responded_by_user.id,
- chosen_options_list,
- resp.responded_at,
- )
+ responded_by_user = self._get_responder(resp)
+ if responded_by_user is None:
+ self.log.info(
+ "[HITL] resuming with the timeout default %s recorded at %s "
+ "(applied by an earlier run of this trigger)",
+ chosen_options_list,
+ resp.responded_at,
+ )
+ else:
+ self.log.info(
+ "[HITL] responded_by=%s (id=%s) options=%s at %s",
+ responded_by_user["name"],
+ responded_by_user["id"],
+ chosen_options_list,
+ resp.responded_at,
+ )
return TriggerEvent(
HITLTriggerEventSuccessPayload(
chosen_options=chosen_options_list,
params_input=params_input or {},
responded_at=resp.responded_at,
- responded_by_user=HITLUser(
- id=resp.responded_by_user.id,
- name=resp.responded_by_user.name,
- ),
- timedout=False,
+ responded_by_user=responded_by_user,
+ timedout=responded_by_user is None,
)
)
diff --git a/providers/standard/tests/unit/standard/triggers/test_hitl.py
b/providers/standard/tests/unit/standard/triggers/test_hitl.py
index 7166d4428bb..bd4a50e73af 100644
--- a/providers/standard/tests/unit/standard/triggers/test_hitl.py
+++ b/providers/standard/tests/unit/standard/triggers/test_hitl.py
@@ -206,6 +206,49 @@ class TestHITLTrigger:
action_datetime,
)
+ @pytest.mark.db_test
+ @pytest.mark.asyncio
+ @mock.patch.object(HITLTrigger, "log")
+
@mock.patch("airflow.providers.standard.triggers.hitl.update_hitl_detail_response")
+ async def test_run_replayed_after_timeout_fallback_has_no_responder(
+ self, mock_update, mock_log, mock_supervisor_comms,
default_trigger_args
+ ):
+ fallback_datetime = utcnow()
+ trigger = HITLTrigger(
+ defaults=["1"],
+ timeout_datetime=utcnow() + timedelta(seconds=0.1),
+ poke_interval=5,
+ **default_trigger_args,
+ )
+ mock_supervisor_comms.send.return_value = HITLDetailResponse(
+ response_received=True,
+ responded_by_user=None,
+ responded_at=fallback_datetime,
+ chosen_options=["1"],
+ params_input={},
+ )
+
+ gen = trigger.run()
+ await asyncio.sleep(0.3)
+ event = await asyncio.create_task(gen.__anext__())
+
+ assert event == TriggerEvent(
+ HITLTriggerEventSuccessPayload(
+ chosen_options=["1"],
+ params_input={},
+ responded_at=fallback_datetime,
+ responded_by_user=None,
+ timedout=True,
+ )
+ )
+ mock_update.assert_not_called()
+ assert mock_log.info.call_args == mock.call(
+ "[HITL] resuming with the timeout default %s recorded at %s "
+ "(applied by an earlier run of this trigger)",
+ ["1"],
+ fallback_datetime,
+ )
+
@pytest.mark.db_test
@pytest.mark.asyncio
@mock.patch.object(HITLTrigger, "log")
