kaxil commented on code in PR #76882:
URL: https://github.com/apache/airflow/pull/76882#discussion_r4555909692
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -237,11 +260,134 @@ def _answered_request(event: dict[str, Any]) -> int |
None:
return -1
-def _is_awaiting_input(error: BaseException) -> bool:
+def _is_awaiting_input(error: BaseException) -> TypeGuard[TaskAwaitingInput]:
"""Whether ``error`` is the pause for tool approval, which only Airflow
3.3+ raises."""
return AIRFLOW_V_3_3_PLUS and isinstance(error, TaskAwaitingInput)
+def _refuse_external_calls(requests: DeferredToolRequests) -> None:
+ """Fail on tool calls that hand their work to an external system, which
AgentOperator cannot resolve."""
+ if requests.calls:
+ raise UnsupportedToolDeferralError(
+ "The agent called tools that need external execution "
+ f"({', '.join(call.tool_name for call in requests.calls)});
AgentOperator only "
+ "supports tools that need approval."
+ )
+
+
+def _call_content(call: ToolCallPart) -> str | None:
Review Comment:
Done, it's `_fingerprint_call_content` now (8eae500733f).
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -237,11 +260,134 @@ def _answered_request(event: dict[str, Any]) -> int |
None:
return -1
-def _is_awaiting_input(error: BaseException) -> bool:
+def _is_awaiting_input(error: BaseException) -> TypeGuard[TaskAwaitingInput]:
"""Whether ``error`` is the pause for tool approval, which only Airflow
3.3+ raises."""
return AIRFLOW_V_3_3_PLUS and isinstance(error, TaskAwaitingInput)
+def _refuse_external_calls(requests: DeferredToolRequests) -> None:
+ """Fail on tool calls that hand their work to an external system, which
AgentOperator cannot resolve."""
+ if requests.calls:
+ raise UnsupportedToolDeferralError(
+ "The agent called tools that need external execution "
+ f"({', '.join(call.tool_name for call in requests.calls)});
AgentOperator only "
+ "supports tools that need approval."
+ )
+
+
+def _call_content(call: ToolCallPart) -> str | None:
+ """Fingerprint what a reviewer sees of a call, its tool and arguments, but
not the id the model gave it."""
+ return fingerprint_tool_call(call.tool_name, call.args_as_dict(), None)
+
+
+@dataclass
+class _ToolApprovalAnswer:
+ """
+ A reviewer's answer, and what it was given for.
+
+ :param step: Which approval step of the run asked, counting replayed ones,
so the answer
+ never reaches another step, such as a later call with the same
arguments.
+ :param calls: The calls the reviewer saw, as ``(tool name, content
fingerprint)``.
+ :param toolset_ids: The agent's rendered toolset ids when the reviewer
answered.
+ """
+
+ step: int
+ calls: list[tuple[str, str | None]]
+ toolset_ids: list[str]
+ approval: bool | ToolDenied
+
+ def take(self, pending: Sequence[ToolCallPart]) -> dict[str, bool |
ToolDenied]:
+ """
+ Return the answer for each pending call it was given for, by the
call's current id.
+
+ A call matches on its tool and arguments, so a model request asked
again on resume, whose
+ calls get new ids, still gets the answer, and a different call with a
reused id does not.
+ Each answered call is used once: a step asked again can hold the same
call twice where
+ the reviewer saw it once.
+ """
+ settled: dict[str, bool | ToolDenied] = {}
+ for call in pending:
+ answered = (call.tool_name, _call_content(call))
+ if answered[1] is not None and answered in self.calls:
+ self.calls.remove(answered)
+ settled[call.tool_call_id] = self.approval
+ return settled
+
+ def dump(self) -> dict[str, Any]:
+ denied = self.approval.message if isinstance(self.approval,
ToolDenied) else None
+ return {
+ "step": self.step,
+ "calls": [list(call) for call in self.calls],
+ "toolset_ids": self.toolset_ids,
+ "denied": denied,
+ }
+
+ @classmethod
+ def load(cls, stored: dict[str, Any]) -> _ToolApprovalAnswer:
+ denied = stored["denied"]
+ return cls(
+ step=stored["step"],
+ calls=[(name, fingerprint) for name, fingerprint in
stored["calls"]],
+ toolset_ids=stored["toolset_ids"],
+ approval=True if denied is None else ToolDenied(denied),
+ )
+
+
+@dataclass
+class _DurableToolApproval(AbstractCapability[Any]):
+ """
+ Settles the tool calls a durable run is waiting to have approved, as a
step of its journal.
+
+ It runs first (outermost) and settles every call or parks the task, so the
whole decision is
+ one step, ``decide``: the reviewers' answers, then the other capabilities'
+ ``handle_deferred_tool_calls`` handlers, then a request for whatever is
left. A parked step
+ is not recorded, so the resumed run, which replays every step before it,
reaches it again
+ with the answer. Once recorded, a retry replays the decision instead of
asking anyone again,
+ and replays the approved calls instead of running them twice.
+
+ :param next_step: Number the approval step about to run for these calls,
replayed or not,
+ or return ``None`` when ``decide`` itself is asking the other handlers.
+ :param settle: The body of ``decide``.
+ """
+
+ next_step: Callable[[DeferredToolRequests], int | None] = field(repr=False)
+ settle: Callable[[RunContext[Any], DeferredToolRequests, int],
Awaitable[DeferredToolResults]] = field(
+ repr=False
+ )
+
+ @classmethod
+ def get_serialization_name(cls) -> str | None:
+ return None
+
+ def get_ordering(self) -> CapabilityOrdering:
+ # Ahead of any handler the Dag author ordered first, so none decides a
call before the reviewer.
+ return CapabilityOrdering(position="outermost")
+
+ async def handle_deferred_tool_calls(
+ self, ctx: RunContext[Any], *, requests: DeferredToolRequests
+ ) -> DeferredToolResults | None:
+ if (step := self.next_step(requests)) is None:
+ return None
+ return await self.decide(ctx, requests=requests, step=step)
+
+ @durable_operation(name="tool_approval")
+ async def decide(
+ self, ctx: RunContext[Any], *, requests: DeferredToolRequests, step:
int
+ ) -> DeferredToolResults | None:
+ return await self.settle(ctx, requests, step)
+
+
+def _only_for(requests: DeferredToolRequests, results: DeferredToolResults) ->
DeferredToolResults:
Review Comment:
Done, renamed to `_filter_results_to_requests` in 8eae500733f.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1435,6 +1613,184 @@ def _resume_after_tool_approval(
context, result, attempt_usage=subtract_run_usage(run_usage,
self._run_usage_base)
)
+ def _check_review_answer(self, event: dict[str, Any], toolset_ids:
list[str]) -> None:
+ """Refuse to resume when the review ended without an answer, or the
toolsets changed while it waited."""
+ if "error" in event:
+ if event.get("error_type") == "timeout":
+ raise HITLTimeoutError(f"Tool approval timed out:
{event['error']}")
+ raise HITLTriggerEventError(event)
+ if (current := self._toolset_ids()) != toolset_ids:
+ raise ToolApprovalError(
+ f"The agent's toolsets changed while it waited for approval
(paused with {toolset_ids}, "
+ f"resumed with {current}), so the approved calls would reach a
connection the reviewer "
+ "did not see."
+ )
+
+ def _answers_waiting_request(self, event: dict[str, Any]) -> bool:
Review Comment:
Done, `_is_answer_for_waiting_request` in 8eae500733f.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1435,6 +1613,184 @@ def _resume_after_tool_approval(
context, result, attempt_usage=subtract_run_usage(run_usage,
self._run_usage_base)
)
+ def _check_review_answer(self, event: dict[str, Any], toolset_ids:
list[str]) -> None:
+ """Refuse to resume when the review ended without an answer, or the
toolsets changed while it waited."""
+ if "error" in event:
+ if event.get("error_type") == "timeout":
+ raise HITLTimeoutError(f"Tool approval timed out:
{event['error']}")
+ raise HITLTriggerEventError(event)
+ if (current := self._toolset_ids()) != toolset_ids:
+ raise ToolApprovalError(
+ f"The agent's toolsets changed while it waited for approval
(paused with {toolset_ids}, "
+ f"resumed with {current}), so the approved calls would reach a
connection the reviewer "
+ "did not see."
+ )
+
+ def _answers_waiting_request(self, event: dict[str, Any]) -> bool:
+ """Whether the answer is for the approval request waiting now, not one
a form left open showed."""
+ if (answered := _answered_request(event)) is None or answered ==
self._approval_pause_number:
+ return True
+ # A review form left open on an earlier pause: its answer was about
other calls.
+ self.log.warning(
+ "An answer for approval request %s arrived while request %s was
waiting; "
+ "asking again instead of applying it.",
+ answered,
+ self._approval_pause_number,
+ )
+ return False
+
+ def _approval_from_answer(self, event: dict[str, Any]) -> bool |
ToolDenied:
Review Comment:
Done, `_build_approval_from_answer` in 8eae500733f.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1435,6 +1613,184 @@ def _resume_after_tool_approval(
context, result, attempt_usage=subtract_run_usage(run_usage,
self._run_usage_base)
)
+ def _check_review_answer(self, event: dict[str, Any], toolset_ids:
list[str]) -> None:
+ """Refuse to resume when the review ended without an answer, or the
toolsets changed while it waited."""
+ if "error" in event:
+ if event.get("error_type") == "timeout":
+ raise HITLTimeoutError(f"Tool approval timed out:
{event['error']}")
+ raise HITLTriggerEventError(event)
+ if (current := self._toolset_ids()) != toolset_ids:
+ raise ToolApprovalError(
+ f"The agent's toolsets changed while it waited for approval
(paused with {toolset_ids}, "
+ f"resumed with {current}), so the approved calls would reach a
connection the reviewer "
+ "did not see."
+ )
+
+ def _answers_waiting_request(self, event: dict[str, Any]) -> bool:
+ """Whether the answer is for the approval request waiting now, not one
a form left open showed."""
+ if (answered := _answered_request(event)) is None or answered ==
self._approval_pause_number:
+ return True
+ # A review form left open on an earlier pause: its answer was about
other calls.
+ self.log.warning(
+ "An answer for approval request %s arrived while request %s was
waiting; "
+ "asking again instead of applying it.",
+ answered,
+ self._approval_pause_number,
+ )
+ return False
+
+ def _approval_from_answer(self, event: dict[str, Any]) -> bool |
ToolDenied:
+ """Turn the reviewer's answer into the result each paused call gets:
run it, or tell the agent why not."""
+ if event.get("timedout"):
+ # on_tool_approval_timeout="deny": nobody refused, so do not tell
the agent a person did.
+ self.log.info(
+ "Tool calls denied: nobody answered within
tool_approval_timeout=%s",
+ self.tool_approval_timeout,
+ )
+ return ToolDenied("No reviewer answered within the approval
timeout, so this call was not run.")
+ if LLMApprovalMixin.APPROVE in event["chosen_options"]:
+ self.log.info("Tool calls approved by %s",
LLMApprovalMixin._describe_responder(event))
+ return True
+ reason = (event.get("params_input") or {}).get("reason")
+ # Only a typed reason reaches the agent; an untouched field can come
back as None, "",
+ # or, from some UI versions, the whole parameter spec.
+ if not isinstance(reason, str) or not reason.strip():
+ reason = "A reviewer denied this tool call."
+ self.log.info("Tool calls rejected by %s",
LLMApprovalMixin._describe_responder(event))
+ return ToolDenied(reason)
+
+ def _load_tool_approval_answers(self, store: TaskStateStoreAccessor) ->
None:
+ """Read the answers reviewers gave this run, leaving out ones given
for other toolsets."""
+ stored =
_STORED_ANSWERS_ADAPTER.validate_python(store.get(_TOOL_APPROVAL_ANSWERS_KEY)
or [])
+ self._tool_approval_answers_stored = bool(stored)
+ toolset_ids = self._toolset_ids()
+ answers = [_ToolApprovalAnswer.load(raw) for raw in stored]
+ if self._new_tool_approval_answer is not None and answers:
+ # The resume stored it last; keep it as the object the warnings
look at.
+ answers[-1] = self._new_tool_approval_answer
+ # Like a resume, an answer never reaches a connection the reviewer did
not see.
+ self._tool_approval_answers = [answer for answer in answers if
answer.toolset_ids == toolset_ids]
+ if skipped := len(answers) - len(self._tool_approval_answers):
+ self.log.warning(
+ "Not applying %s approval answer(s) given while the agent had
other toolsets.",
+ skipped,
+ )
+
+ def _next_approval_step(self, requests: DeferredToolRequests) -> int |
None:
Review Comment:
Good catch, I could reproduce it: with `run_in_executor` + `asyncio.run` the
inner approval parked the task instead of failing it. Fixed in 8eae500733f. The
journal now tracks its running steps on the journal object itself, and
`_claim_next_approval_step` (renamed as well) asks it whether a tool call is
still running, so the refusal no longer depends on the context var. It only
looks at tool calls, since other journaled work like event handling can still
be running when the top-level approval is decided. The test is now parametrized
over a plain `await`, `asyncio.to_thread` and `loop.run_in_executor(None,
lambda: asyncio.run(...))`, and all three expect `UnsupportedToolDeferralError`.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/capability.py:
##########
@@ -803,6 +819,20 @@ def _is_journaled_tool_payload(payload: object) -> bool:
return False
+def _capability_fingerprint(cache_key: tuple[object, ...]) -> str | None:
Review Comment:
Done, `_fingerprint_capability_operation` in 8eae500733f.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
