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 a8e95debbc6 Fix flaky AgentOperator durable usage budget tests (#73996)
a8e95debbc6 is described below
commit a8e95debbc6cce2f2437d88daa529e3cd99d99bb
Author: Aaron Chen <[email protected]>
AuthorDate: Wed Sep 30 15:47:01 2026 -0700
Fix flaky AgentOperator durable usage budget tests (#73996)
* Fix flaky AgentOperator durable usage budget tests
* Refactor imports in test_agent.py to include Tool from pydantic_ai
---
providers/common/ai/tests/unit/common/ai/operators/test_agent.py | 6 ++++--
1 file changed, 4 insertions(+), 2 deletions(-)
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 274fb3ce3d8..1f8575be2ce 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -28,7 +28,7 @@ from unittest.mock import ANY, MagicMock, PropertyMock, call,
patch
import pytest
from pydantic import BaseModel
-from pydantic_ai import Agent, DeferredToolRequests
+from pydantic_ai import Agent, DeferredToolRequests, Tool
from pydantic_ai.capabilities import Toolset
from pydantic_ai.exceptions import UsageLimitExceeded
from pydantic_ai.messages import (
@@ -2884,7 +2884,9 @@ class _DurableBudgetScenario:
self.live_tool_calls += 1
return "B"
- return FunctionToolset(tools=[tool_a, tool_b])
+ # Cache tool_a before tool_b fails; sibling cancellation in a parallel
batch
+ # can otherwise prevent the successful result from being persisted.
+ return FunctionToolset(tools=[Tool(tool_a, sequential=True), tool_b])
def run_attempt(
self, *, max_tries=0, prompt="run", limits=None, fail_after_run=False,
interrupt_at_step=None
