This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch durable-tool-approval in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 035c5416d70fe472ab0b4132ebff47ee8867ce35 Author: Kaxil Naik <[email protected]> AuthorDate: Thu Nov 12 10:23:32 2026 +0000 Fail durable agent tasks once on invalid tool opt-out metadata A tool's airflow metadata that is not True, False or None (such as the string "false") raised pydantic-ai's UserError only when the model first called the tool, so every retry asked the model again up to that call and failed the same way. The metadata of every tool offered to the model is now checked before the model request, and an invalid value raises DurableToolMetadataError, an AirflowFailException, so the task fails once without retrying. Outside an Airflow task, where durable execution is off, the value is not checked. --- providers/common/ai/docs/durable_execution.rst | 6 +++--- .../providers/common/ai/durable/capability.py | 18 ++++++++++++---- .../src/airflow/providers/common/ai/exceptions.py | 9 ++++++++ .../unit/common/ai/durable/test_capability.py | 25 ++++++++++++++++++---- 4 files changed, 47 insertions(+), 11 deletions(-) diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index 679be441518..57cda143168 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/providers/common/ai/docs/durable_execution.rst @@ -468,12 +468,12 @@ The error from pydantic-ai's ``MCPToolset``: pydantic_ai.exceptions.UserError: Airflow durable config for MCP tool 'get_forecast' has been explicitly set to `False` (durable execution disabled), but MCP tools perform I/O and cannot run outside a durable step. Remove the metadata so the call stays durable. -Any value other than ``True``, ``False`` or ``None`` is refused when the tool is called. The -attempt fails, as do its retries: +Any value other than ``True``, ``False`` or ``None`` fails the task before the model is first +offered the tool, and the task is not retried: .. code-block:: text - pydantic_ai.exceptions.UserError: Tool 'today' has invalid 'airflow' metadata: expected True or False, got str. + airflow.providers.common.ai.exceptions.DurableToolMetadataError: Tool 'today' has invalid 'airflow' metadata: expected True or False, got str. Limitations ----------- diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py index e3c242e0102..1a8d73b8732 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py @@ -46,7 +46,6 @@ from pydantic_ai.exceptions import ( ModelRetry, ToolFailed, UnexpectedModelBehavior, - UserError, ) from pydantic_ai.messages import ModelResponse, ToolReturn from pydantic_ai.tools import AgentDepsT @@ -58,7 +57,7 @@ from airflow.providers.common.ai.durable.base import TOOL_CONFIG_KEY from airflow.providers.common.ai.durable.fingerprint import fingerprint_model_request, fingerprint_tool_call from airflow.providers.common.ai.durable.journal import Unrecordable, current_journal from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger, is_successful_tool_payload -from airflow.providers.common.ai.exceptions import DurableJournalError +from airflow.providers.common.ai.exceptions import DurableJournalError, DurableToolMetadataError from airflow.providers.common.ai.utils.masking import mask_secrets from airflow.providers.common.ai.utils.tool_metrics import record_tool_call from airflow.providers.common.ai.utils.toolset_base import AirflowToolset @@ -71,7 +70,7 @@ if TYPE_CHECKING: from pydantic_ai.durable_exec._operation import DurableOperation, ParamsT, ResultT, WireT from pydantic_ai.durable_exec._operation_backend import BoundDurableOperation from pydantic_ai.messages import ModelMessage - from pydantic_ai.models import Model, ModelRequestParameters + from pydantic_ai.models import Model, ModelRequestContext, ModelRequestParameters from pydantic_ai.run import AgentRunResult from pydantic_ai.settings import ModelSettings from pydantic_ai.tools import RunContext, ToolDefinition @@ -112,7 +111,7 @@ def _is_tool_journaled(tool_def: ToolDefinition) -> bool: return True if isinstance(setting, bool): return setting - raise UserError( + raise DurableToolMetadataError( f"Tool {tool_def.name!r} has invalid {TOOL_CONFIG_KEY!r} metadata: expected True or False, " f"got {type(setting).__name__}." ) @@ -329,6 +328,17 @@ class AirflowDurability(BaseDurabilityCapability[AgentDepsT]): def get_durable_operation_backend(self) -> _AirflowOperationBackend: return _AirflowOperationBackend(self) + async def before_model_request( + self, ctx: RunContext[AgentDepsT], request_context: ModelRequestContext + ) -> ModelRequestContext: + # A bad ``airflow`` metadata value fails the task here, before the model is first asked, + # rather than when the model calls the tool, which every retry would reach again. The + # request's tool definitions carry the metadata other capabilities set on the toolsets. + if self.in_durable_context: + for tool_def in request_context.model_request_parameters.function_tools: + _is_tool_journaled(tool_def) + return await super().before_model_request(ctx, request_context) + def get_wrapper_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT] | None: """Journal the leaf toolsets pydantic-ai's backend does not, then let the base wrap the rest.""" journaled = toolset.visit_and_replace(self._journal_other_leaf) diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py index 67facfa980e..201445725db 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py +++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py @@ -76,6 +76,15 @@ class UnsupportedToolDeferralError(AirflowFailException): """ +class DurableToolMetadataError(AirflowFailException): + """ + Raised when a tool's ``airflow`` metadata is not ``True``, ``False`` or ``None``. + + Durable execution checks it when the agent lists its tools, before the first model + request. Every attempt would fail the same way, so the task fails without retrying. + """ + + class DurableJournalError(AirflowFailException): """ Raised when durable execution cannot record or replay a step. diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py index 5d4284ee0e9..c9f9b6022c8 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py @@ -34,7 +34,7 @@ import pytest from pydantic import BaseModel from pydantic_ai import Agent, CancellationToken, RunContext from pydantic_ai.capabilities import AbstractCapability, SetToolMetadata, durable_operation -from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError +from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior from pydantic_ai.messages import ( ModelMessage, ModelResponse, @@ -53,6 +53,7 @@ from airflow.providers.common.ai.durable.capability import _JournalCodec from airflow.providers.common.ai.durable.journal import DurableJournal, Unrecordable, journal_scope from airflow.providers.common.ai.exceptions import DurableJournalError from airflow.providers.common.ai.utils.toolset_base import AirflowToolset, ensure_masked +from airflow.providers.common.compat.sdk import AirflowFailException from unit.common.ai.durable.memory_storage import MemoryStorage from unit.common.ai.tool_parts import is_tool_result, is_tool_retry @@ -1053,9 +1054,10 @@ class TestToolOptOut: @pytest.mark.asyncio @pytest.mark.parametrize("value", ["false", 0]) @pytest.mark.parametrize("build_toolset", [_function_toolset, _WarehouseToolset]) - async def test_an_airflow_metadata_value_that_is_not_a_bool_is_refused( + async def test_an_airflow_metadata_value_that_is_not_a_bool_fails_the_task_before_the_model_is_asked( self, memory_storage, value, build_toolset ): + """Every attempt would fail the same way, so the task fails once, without retrying.""" calls = Calls() agent = Agent( FunctionModel(tool_then_answer(calls, "query", fail_final=[False])), @@ -1064,9 +1066,24 @@ class TestToolOptOut: capabilities=[AirflowDurability(), SetToolMetadata(tools=["query"], airflow=value)], ) - with pytest.raises(UserError, match="Tool 'query' has invalid 'airflow' metadata"): + with pytest.raises(AirflowFailException, match="Tool 'query' has invalid 'airflow' metadata"): await attempt(memory_storage, agent) - assert calls["query"] == 0 + assert (calls["model"], calls["query"]) == (0, 0) + + @pytest.mark.asyncio + async def test_an_airflow_metadata_value_that_is_not_a_bool_is_ignored_outside_a_task(self): + calls = Calls() + agent = Agent( + FunctionModel(tool_then_answer(calls, "query", fail_final=[False])), + name="analyst", + toolsets=[_function_toolset(calls, "3 rows")], + capabilities=[AirflowDurability(), SetToolMetadata(tools=["query"], airflow="false")], + ) + + result = await agent.run("go") + + assert result.output == "answer: 3 rows" + assert calls["query"] == 1 class Answer(BaseModel):
