kaxil commented on code in PR #71461:
URL: https://github.com/apache/airflow/pull/71461#discussion_r3774212894
##########
providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py:
##########
@@ -636,12 +696,13 @@ def wait_for_session(
idle event on a ``message`` run (defeats the start race).
:param poll_interval: Seconds to sleep between polls.
:param timeout: Maximum seconds to wait before raising
:class:`AnthropicAgentSessionTimeout`.
+ :raises AnthropicSessionBudgetExceeded: If the session stopped against
its budget.
"""
start = time.monotonic()
consecutive_failures = 0
while True:
try:
- done, error_message = self.poll_session_completion(
+ poll = self.poll_session_completion(
Review Comment:
Renamed -- `poll_result` here, in the trigger, and at the test call sites.
##########
providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py:
##########
@@ -94,7 +94,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
while True:
try:
# poll_session_completion does blocking SDK HTTP calls; run
off the event loop.
- done, error_message = await asyncio.to_thread(
+ poll = await asyncio.to_thread(
Review Comment:
Done here too.
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -156,39 +159,77 @@ class TestPollSessionCompletion:
def test_terminated_is_error(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("terminated")
- done, err = hook.poll_session_completion("s")
- assert done is True
- assert err is not None
+ poll = hook.poll_session_completion("s")
Review Comment:
Renamed to `poll_result`, and took your second option in spirit -- there was
a real mismatch behind the confusion, not just a short name.
`_session()` builds the mocked session with `id = "sess_1"`, but this class
was passing `"s"`. The two ids reach the assertions by different routes:
`evaluate_session_state` interpolates `session.id`, while
`poll_session_completion` interpolates the argument. So the messages under test
carried `"sess_1"` on one path and `"s"` on the other, and an assertion could
have matched the wrong one.
There is now a module-level `SESSION_ID = "sess_1"` used by both
`_session()` and every call in this class, so they cannot drift apart.
I left the pre-existing `"sess_1"` literals elsewhere in the file as they
are, to keep the diff off tests this PR does not otherwise touch -- happy to
move them onto the constant too. I also skipped the explicit `session_id=`
keyword, since the value now names itself; say the word if you would still
rather have it spelled out.
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -156,39 +159,77 @@ class TestPollSessionCompletion:
def test_terminated_is_error(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("terminated")
- done, err = hook.poll_session_completion("s")
- assert done is True
- assert err is not None
+ poll = hook.poll_session_completion("s")
+ assert poll.done is True
+ assert poll.error_message is not None
+ assert poll.stop_reason is None
def test_message_end_turn_success(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle")
client.beta.sessions.events.list.return_value =
[_idle_event("end_turn")]
- assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== (True, None)
+ assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== SessionPollResult(
+ done=True, error_message=None, stop_reason="end_turn"
+ )
@pytest.mark.parametrize("reason", ["requires_action",
"retries_exhausted"])
def test_message_blocked_is_error(self, reason):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle")
client.beta.sessions.events.list.return_value = [_idle_event(reason)]
- done, err = hook.poll_session_completion("s",
kickoff_event_id="evt_kick")
- assert done is True
- assert err is not None
- assert reason in err
+ poll = hook.poll_session_completion("s", kickoff_event_id="evt_kick")
+ assert poll.done is True
+ assert poll.error_message is not None
+ assert reason in poll.error_message
+ assert poll.stop_reason == reason
def test_message_no_response_yet_not_done(self):
# newest event is our kickoff (agent hasn't responded) -> keep waiting
(start race)
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle")
client.beta.sessions.events.list.return_value =
[mock.MagicMock(type="user.message", id="evt_kick")]
- assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== (False, None)
+ assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== SessionPollResult(
+ done=False, error_message=None, stop_reason=None
+ )
def test_outcome_satisfied_skips_event_check(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle",
["satisfied"])
- assert hook.poll_session_completion("s", expect_outcome=True) ==
(True, None)
+ assert hook.poll_session_completion("s", expect_outcome=True) ==
SessionPollResult(
+ done=True, error_message=None, stop_reason=None
+ )
client.beta.sessions.events.list.assert_not_called()
+ def test_budget_reached_names_both_causes(self):
+ # A budget stop must not be reported as "configure an autonomous
agent": that advice
+ # is wrong, and the no-list-price cause is invisible without being
named.
+ hook, client = _make_hook()
+ client.beta.sessions.retrieve.return_value = _session("idle")
+ client.beta.sessions.events.list.return_value =
[_idle_event("budget_reached")]
+ poll = hook.poll_session_completion("s", kickoff_event_id="evt_kick")
+ assert poll.done is True
+ assert poll.stop_reason == "budget_reached"
+ assert "budget" in poll.error_message
+ assert "no list price" in poll.error_message
+ assert "autonomous agent" not in poll.error_message
Review Comment:
Took the whole-object assertion with `error_message=mock.ANY`, and the loop
for the causes that must be named.
One correction on the loop, though: it inverts the last assertion. The
original line was
```python
assert "autonomous agent" not in poll.error_message
```
`not in`, and that negative is the point of the test -- a budget stop must
*not* carry the generic "configure an autonomous agent or use an outcome run"
advice, because that advice is wrong for this case and was the bug the PR
fixes. Folding it into a loop of `in` checks asserts the opposite of what the
test is for.
I ran your version verbatim to be sure rather than argue from reading:
```
FAILED test_budget_reached_names_both_causes - AssertionError: assert
'autonomous agent' in
'Session s stopped against its budget: the tracked list cost reached the
configured ceiling,
or its usage included a model with no list price (which a budget cannot
measure). ...'
```
So it is the loop over the two causes that must be present, with the
negative kept as its own line and a comment saying why it is there.
I applied the same whole-object shape to `test_message_blocked_is_error`
just above, since it had the same stack of single-field asserts.
--
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]
