This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch anthropic-session-budget-reached in repository https://gitbox.apache.org/repos/asf/airflow.git
commit a1c4cb0b288cbc6ee60dbb707c3894a1a62f1bf1 Author: Kaxil Naik <[email protected]> AuthorDate: Tue Aug 11 13:26:01 2026 +0530 Fix misleading error when an Anthropic agent session stops on its budget Anthropic SDK 0.121.0 adds Managed Agents session budgets: a spend ceiling passed to ``sessions.create``, after which the session stops issuing new model requests and goes idle with a new ``budget_reached`` stop reason. The provider treated every non-``end_turn`` idle reason the same way, so a budget stop surfaced as "configure an autonomous agent or use an outcome run" -- advice that has nothing to do with what happened. It also raised the generic ``AnthropicAgentSessionError``, giving Dag authors no way to tell a spend decision apart from a fault. ``budget_reached`` now gets its own branch and a new ``AnthropicSessionBudgetExceeded`` (a subclass of the existing session error, so current handlers keep working). The message names both causes, because a session also stops this way when its usage includes a model with no list price that a budget cannot measure -- and raising the ceiling does not unblock that case. It also says the operator archives the session on this path, so the ceiling has to be raised for the next run rather than on a session that no longer exists. ``poll_session_completion`` returns the SDK's ``stop_reason`` alongside the message, and the trigger carries it in its error event, so the deferrable path raises the same class as the synchronous one without matching on message text. A trigger serialized before this field existed omits it and falls back to the generic error. Verified against the live API, not only against mocks: ``sessions.create`` accepts the documented budget payload, the server really emits ``budget_reached``, and the hook classifies a genuinely budget-stopped session and raises ``AnthropicSessionBudgetExceeded``. That testing also showed a budget is a stop trigger, not a spend cap. The ceiling is checked between model requests, so an in-flight request runs to completion: a $0.01 ceiling admitted between $0.32 and $0.61 of usage across four runs. The docs now say so, since "hard spend ceiling" invites the wrong expectation. The ``anthropic`` floor moves to 0.121.0, the first release whose ``sessions.create`` accepts ``budget`` -- which is how the guide tells users to set a ceiling. Reading a ``budget_reached`` stop reason happens to work further back, because older clients mis-build the discriminated union into the wrong variant class while preserving ``.type``, but only with response validation left non-strict, so it is not something to depend on. --- providers/anthropic/README.rst | 2 +- providers/anthropic/docs/index.rst | 8 +- providers/anthropic/docs/operators/anthropic.rst | 67 ++++++++++++++- providers/anthropic/pyproject.toml | 15 ++-- .../src/airflow/providers/anthropic/exceptions.py | 11 +++ .../airflow/providers/anthropic/hooks/anthropic.py | 89 ++++++++++++++++---- .../airflow/providers/anthropic/operators/agent.py | 25 ++++-- .../airflow/providers/anthropic/triggers/agent.py | 18 +++- .../tests/unit/anthropic/hooks/test_anthropic.py | 95 ++++++++++++++++++---- .../tests/unit/anthropic/operators/test_agent.py | 30 +++++++ .../tests/unit/anthropic/test_exceptions.py | 2 + .../tests/unit/anthropic/triggers/test_agent.py | 37 +++++++-- uv.lock | 14 ++-- 13 files changed, 345 insertions(+), 68 deletions(-) diff --git a/providers/anthropic/README.rst b/providers/anthropic/README.rst index c895946d7c9..265ea60f94c 100644 --- a/providers/anthropic/README.rst +++ b/providers/anthropic/README.rst @@ -57,7 +57,7 @@ PIP package Version required ========================================== ================== ``apache-airflow`` ``>=3.0.0`` ``apache-airflow-providers-common-compat`` ``>=1.12.0`` -``anthropic`` ``>=0.101.0`` +``anthropic`` ``>=0.121.0`` ========================================== ================== Optional dependencies diff --git a/providers/anthropic/docs/index.rst b/providers/anthropic/docs/index.rst index cab3cc67c30..c6b97c9716c 100644 --- a/providers/anthropic/docs/index.rst +++ b/providers/anthropic/docs/index.rst @@ -138,7 +138,7 @@ PIP package Version required ========================================== ================== ``apache-airflow`` ``>=3.0.0`` ``apache-airflow-providers-common-compat`` ``>=1.12.0`` -``anthropic`` ``>=0.101.0`` +``anthropic`` ``>=0.121.0`` ========================================== ================== Optional dependencies @@ -155,9 +155,9 @@ Install them when installing from PyPI. For example: =========== =============================== Extra Dependencies =========== =============================== -``bedrock`` ``anthropic[bedrock]>=0.101.0`` -``vertex`` ``anthropic[vertex]>=0.101.0`` -``aws`` ``anthropic[aws]>=0.101.0`` +``bedrock`` ``anthropic[bedrock]>=0.121.0`` +``vertex`` ``anthropic[vertex]>=0.121.0`` +``aws`` ``anthropic[aws]>=0.121.0`` =========== =============================== Downloading official packages diff --git a/providers/anthropic/docs/operators/anthropic.rst b/providers/anthropic/docs/operators/anthropic.rst index 90b288a0c82..a904e687549 100644 --- a/providers/anthropic/docs/operators/anthropic.rst +++ b/providers/anthropic/docs/operators/anthropic.rst @@ -155,10 +155,69 @@ Parameters Completion is detected accurately for both modes. A ``message`` run inspects the terminal ``session.status_idle`` event's ``stop_reason`` (correlated against the - kickoff event): ``end_turn`` succeeds; ``requires_action`` and ``retries_exhausted`` - raise an error. An ``outcome`` run is judged from the ``outcome_evaluations`` verdict. - The agent must still be configured for autonomous operation (no client-side custom - tools / ``always_ask``). + kickoff event): ``end_turn`` succeeds; ``requires_action``, ``retries_exhausted`` and + ``budget_reached`` raise an error. An ``outcome`` run is judged from the + ``outcome_evaluations`` verdict. The agent must still be configured for autonomous + operation (no client-side custom tools / ``always_ask``). + +Session budgets +""""""""""""""" + +Pass a `session budget +<https://platform.claude.com/docs/en/managed-agents/overview>`__ through +``session_kwargs`` to cap what a single session may spend. The session stops issuing new +model requests once its tracked list cost reaches the ceiling: + +.. code-block:: python + + AnthropicAgentSessionOperator( + task_id="research", + agent_id="agt_...", + environment_id="env_...", + message="Summarise yesterday's incidents.", + session_kwargs={ + "budget": { + "type": "limit", + # Minor units as an integer string: "2500" is $25.00. + "max_list_cost": {"amount": "2500", "currency": "USD"}, + } + }, + ) + +On a ``message`` run, a session that stops this way raises +:class:`~airflow.providers.anthropic.exceptions.AnthropicSessionBudgetExceeded`, a subclass +of ``AnthropicAgentSessionError``, so it can be caught on its own and routed to review +rather than treated as a fault. + +.. warning:: + + On an ``outcome`` run, completion is judged from ``outcome_evaluations`` before the idle + event is read, so a budget stop raises nothing. The session stays non-terminal, polling + continues until ``timeout`` (24 hours by default), and the task then fails with + ``AnthropicAgentSessionTimeout`` -- a misleading error for a session that stopped + deliberately. Set a shorter ``timeout`` when combining ``outcome`` with a budget. + +.. warning:: + + **A budget is a stop trigger, not a spend cap.** The ceiling is checked *between* + model requests, so a request already in flight runs to completion and the session can + finish well above the limit -- in testing, by a large multiple of a very small + ceiling, because a single long generation overshoots before the next request can be + blocked. Size it as a circuit breaker rather than a guarantee, and read the session's + ``usage.list_cost`` for what was actually spent. + +.. warning:: + + A session also stops with ``budget_reached`` when its usage includes a model with **no + list price**, because a budget cannot measure that spend. Raising the ceiling does not + unblock that case; remove the budget instead. + +.. warning:: + + Airflow ``retries`` multiply spend. Each retry starts a **new** session with a **fresh** + budget, so ``retries=2`` with a $25 ceiling can spend $75. Prefer ``retries=0`` on + budgeted sessions: the operator archives a budget-stopped session, so there is no + running session left to raise the ceiling on. .. exampleinclude:: /../tests/system/anthropic/example_anthropic_agent.py :language: python diff --git a/providers/anthropic/pyproject.toml b/providers/anthropic/pyproject.toml index 6affe34490d..da81825cce7 100644 --- a/providers/anthropic/pyproject.toml +++ b/providers/anthropic/pyproject.toml @@ -61,17 +61,20 @@ requires-python = ">=3.10" dependencies = [ "apache-airflow>=3.0.0", "apache-airflow-providers-common-compat>=1.12.0", - # 0.101.0 is the first release that ships AnthropicAWS, the newest of the platform - # client classes the hook imports at module top. - "anthropic>=0.101.0", + # 0.121.0 is the first release whose ``sessions.create`` accepts ``budget``, which is + # how the guide tells users to set a session ceiling. Reading a ``budget_reached`` stop + # reason happens to work further back -- older clients mis-build the discriminated + # union into the wrong variant class while preserving ``.type`` -- but only with + # response validation left non-strict, so it is not something to depend on. + "anthropic>=0.121.0", ] # The optional dependencies should be modified in place in the generated file # Any change in the dependencies is preserved when the file is regenerated [project.optional-dependencies] -"bedrock" = ["anthropic[bedrock]>=0.101.0"] -"vertex" = ["anthropic[vertex]>=0.101.0"] -"aws" = ["anthropic[aws]>=0.101.0"] +"bedrock" = ["anthropic[bedrock]>=0.121.0"] +"vertex" = ["anthropic[vertex]>=0.121.0"] +"aws" = ["anthropic[aws]>=0.121.0"] [dependency-groups] dev = [ diff --git a/providers/anthropic/src/airflow/providers/anthropic/exceptions.py b/providers/anthropic/src/airflow/providers/anthropic/exceptions.py index 7250d1e848c..2cf58b8a709 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/exceptions.py +++ b/providers/anthropic/src/airflow/providers/anthropic/exceptions.py @@ -37,5 +37,16 @@ class AnthropicAgentSessionError(AnthropicError): """Raised when a Managed Agents session terminates or fails.""" +class AnthropicSessionBudgetExceeded(AnthropicAgentSessionError): + """ + Raised when a Managed Agents session stops against its configured budget. + + A subclass of :class:`AnthropicAgentSessionError` so existing handlers keep working, + but it can be caught on its own: a budget stop is a spend decision, not a fault, and + usually wants different handling (alert, raise the ceiling, route to review) than a + failed run. + """ + + class AnthropicAgentSessionTimeout(AnthropicError): """Raised when a Managed Agents session does not reach a terminal status in time.""" diff --git a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py index 5f0cbd40886..932b9215b35 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py +++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py @@ -20,7 +20,7 @@ import logging import time from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, NamedTuple, cast from anthropic import ( Anthropic, @@ -38,6 +38,7 @@ from airflow.providers.anthropic.exceptions import ( AnthropicBatchJobError, AnthropicBatchTimeout, AnthropicError, + AnthropicSessionBudgetExceeded, AnthropicTriggerEventError, ) from airflow.providers.common.compat.sdk import AirflowSkipException, BaseHook @@ -119,6 +120,35 @@ class SessionStatus(str, Enum): #: ``outcome_evaluations[].result`` values that mean the outcome did NOT succeed. OUTCOME_FAILURE_RESULTS = frozenset({"failed", "max_iterations_reached", "interrupted"}) +# ``session.status_idle`` stop reason emitted when a session stops against its budget. +BUDGET_REACHED = "budget_reached" + + +def _create_session_error(message: str, stop_reason: str | None) -> AnthropicAgentSessionError: + """ + Return the session error class matching an idle ``stop_reason``. + + Keyed on the SDK's own ``stop_reason`` value rather than on the message text, so the + synchronous path and the deferrable path (which only carries the reason as a string + through the trigger event) raise the same type for the same cause. + """ + if stop_reason == BUDGET_REACHED: + return AnthropicSessionBudgetExceeded(message) + return AnthropicAgentSessionError(message) + + +class SessionPollResult(NamedTuple): + """ + Verdict from one poll of a session; see :meth:`AnthropicHook.poll_session_completion`. + + Named rather than a bare tuple because ``error_message`` and ``stop_reason`` are both + ``str | None``, so transposing them at a call site would still type-check. + """ + + done: bool + error_message: str | None + stop_reason: str | None + def evaluate_session_state( session: BetaManagedAgentsSession, *, expect_outcome: bool @@ -587,13 +617,24 @@ class AnthropicHook(BaseHook): def poll_session_completion( self, session_id: str, *, expect_outcome: bool = False, kickoff_event_id: str | None = None - ) -> tuple[bool, str | None]: + ) -> SessionPollResult: """ - Return ``(done, error_message)`` for one poll of a session. + Return the :class:`SessionPollResult` for one poll of a session. Combines the session object (status / outcome verdict) with the event log (``stop_reason`` of the latest idle) so a ``message`` run distinguishes genuine - ``end_turn`` completion from ``requires_action`` / ``retries_exhausted``. + ``end_turn`` completion from ``requires_action`` / ``retries_exhausted`` / + ``budget_reached``. + + ``stop_reason`` is the SDK's own idle stop reason, or ``None`` when the verdict did + not come from an idle event (a ``terminated`` session, or an outcome verdict). It + exists so callers can pick an error class without matching on the message text; pass + it to :func:`_create_session_error`. + + .. note:: + A budget stop is classified on ``message`` runs only. An ``outcome`` run is + judged from ``outcome_evaluations`` before the event log is consulted, so a + budget stop there surfaces as whatever verdict the outcome recorded. """ session = self.get_session(session_id) done, error_message, needs_event_check = evaluate_session_state( @@ -607,15 +648,34 @@ class AnthropicHook(BaseHook): needs_event_check, ) if not needs_event_check: - return done, error_message + return SessionPollResult(done=done, error_message=error_message, stop_reason=None) reason = self._latest_idle_reason(session_id, kickoff_event_id) if reason is None: - return False, None + return SessionPollResult(done=False, error_message=None, stop_reason=None) if reason == "end_turn": - return True, None - return True, ( - f"Session {session_id} is idle but did not complete ({reason}); " - "configure an autonomous agent or use an outcome run." + return SessionPollResult(done=True, error_message=None, stop_reason=reason) + if reason == BUDGET_REACHED: + # Both causes are worth naming: a session also stops with ``budget_reached`` + # when its usage includes a model with no list price, because the budget cannot + # measure that spend -- and then raising the ceiling does not unblock it. + return SessionPollResult( + done=True, + error_message=( + f"Session {session_id} 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). The operator archives the " + "session on this path, so it cannot be resumed -- raise the ceiling for " + "the next run, or drop the budget if a model has no list price." + ), + stop_reason=reason, + ) + return SessionPollResult( + done=True, + error_message=( + f"Session {session_id} is idle but did not complete ({reason}); " + "configure an autonomous agent or use an outcome run." + ), + stop_reason=reason, ) def wait_for_session( @@ -636,12 +696,13 @@ class AnthropicHook(BaseHook): 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_result = self.poll_session_completion( session_id, expect_outcome=expect_outcome, kickoff_event_id=kickoff_event_id ) except Exception as e: @@ -657,9 +718,9 @@ class AnthropicHook(BaseHook): time.sleep(poll_interval) continue consecutive_failures = 0 - if done: - if error_message: - raise AnthropicAgentSessionError(error_message) + if poll_result.done: + if poll_result.error_message: + raise _create_session_error(poll_result.error_message, poll_result.stop_reason) return if time.monotonic() - start > timeout: raise AnthropicAgentSessionTimeout( diff --git a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py index e829ab0de6a..a8599a618e4 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py +++ b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py @@ -22,8 +22,12 @@ from datetime import timedelta from functools import cached_property from typing import TYPE_CHECKING, Any -from airflow.providers.anthropic.exceptions import AnthropicAgentSessionError, AnthropicAgentSessionTimeout -from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, validate_execute_complete_event +from airflow.providers.anthropic.exceptions import AnthropicAgentSessionTimeout +from airflow.providers.anthropic.hooks.anthropic import ( + AnthropicHook, + _create_session_error, + validate_execute_complete_event, +) from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrigger from airflow.providers.common.compat.sdk import BaseOperator, conf @@ -49,9 +53,9 @@ class AnthropicAgentSessionOperator(BaseOperator): Completion is detected accurately for both modes. A ``message`` run reads the terminal ``session.status_idle`` event's ``stop_reason`` (correlated against the kickoff event to avoid a start-race false positive): ``end_turn`` succeeds, while - ``requires_action`` (the agent is blocked on input) and ``retries_exhausted`` - raise an error rather than silently passing. An ``outcome`` run is judged from the - session's ``outcome_evaluations`` verdict (``satisfied`` vs. + ``requires_action`` (the agent is blocked on input), ``retries_exhausted`` and + ``budget_reached`` raise an error rather than silently passing. An ``outcome`` run + is judged from the session's ``outcome_evaluations`` verdict (``satisfied`` vs. ``failed``/``max_iterations_reached``/``interrupted``). Agents and environments are created once (see @@ -81,7 +85,14 @@ class AnthropicAgentSessionOperator(BaseOperator): :param session_resources: Session resources (files, GitHub repos, memory stores). Named ``session_resources`` to avoid colliding with the reserved ``BaseOperator.resources``; forwarded to ``sessions.create`` as ``resources``. - :param session_kwargs: Extra keyword arguments forwarded to ``sessions.create``. + :param session_kwargs: Extra keyword arguments forwarded to ``sessions.create``, such as + ``budget`` (a spend ceiling for the session). A session that stops against its + budget raises + :class:`~airflow.providers.anthropic.exceptions.AnthropicSessionBudgetExceeded`. + The ceiling is a stop trigger rather than a cap -- it is checked between model + requests, so an in-flight request can carry the session past it. Airflow + ``retries`` also each start a new session with a fresh budget, so prefer + ``retries=0`` when a budget is set. """ template_fields: Sequence[str] = ("agent_id", "environment_id", "message", "outcome") @@ -202,7 +213,7 @@ class AnthropicAgentSessionOperator(BaseOperator): # The trigger yields "error" when polling gives up while the session may still # be running; archive it best-effort so its container does not linger. self._archive_session(self.session_id) - raise AnthropicAgentSessionError(event["message"]) + raise _create_session_error(event["message"], event.get("stop_reason")) self.log.info("Session %s completed.", self.session_id) return self.session_id diff --git a/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py b/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py index 2045f4cc2d0..875f0c26538 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py +++ b/providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py @@ -94,7 +94,7 @@ class AnthropicAgentSessionTrigger(BaseTrigger): while True: try: # poll_session_completion does blocking SDK HTTP calls; run off the event loop. - done, error_message = await asyncio.to_thread( + poll_result = await asyncio.to_thread( hook.poll_session_completion, self.session_id, expect_outcome=self.expect_outcome, @@ -111,10 +111,20 @@ class AnthropicAgentSessionTrigger(BaseTrigger): continue consecutive_failures = 0 - if done: - if error_message: + if poll_result.done: + if poll_result.error_message: + # ``stop_reason`` is the SDK's idle stop reason; it lets the resuming + # worker raise the same exception class the synchronous path would, + # without matching on the message text. A trigger serialized before this + # field existed simply omits it, and the operator falls back to the + # generic session error. yield TriggerEvent( - {"status": "error", "session_id": self.session_id, "message": error_message} + { + "status": "error", + "session_id": self.session_id, + "message": poll_result.error_message, + "stop_reason": poll_result.stop_reason, + } ) else: yield TriggerEvent( diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py index 2fe3f6fe1f7..30ad61170b7 100644 --- a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py +++ b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py @@ -25,6 +25,7 @@ from airflow.providers.anthropic.exceptions import ( AnthropicAgentSessionTimeout, AnthropicBatchTimeout, AnthropicError, + AnthropicSessionBudgetExceeded, AnthropicTriggerEventError, ) from airflow.providers.anthropic.hooks.anthropic import ( @@ -32,7 +33,9 @@ from airflow.providers.anthropic.hooks.anthropic import ( MAX_CONSECUTIVE_POLL_FAILURES, AnthropicHook, BatchStatus, + SessionPollResult, SessionStatus, + _create_session_error, evaluate_session_state, validate_execute_complete_event, ) @@ -40,6 +43,9 @@ from airflow.providers.anthropic.hooks.anthropic import ( pytest.importorskip("anthropic") HOOK_PATH = "airflow.providers.anthropic.hooks.anthropic" +# One id for both the mocked session object and the id passed to the hook, so an assertion +# on a message that interpolates the session id cannot pass against the wrong one. +SESSION_ID = "sess_1" def _conn(password="sk-ant-test", host=None, extra=None): @@ -108,7 +114,7 @@ class TestValidateTriggerEvent: def _session(status, outcome_results=None): s = mock.MagicMock() s.status = status - s.id = "sess_1" + s.id = SESSION_ID s.outcome_evaluations = [mock.MagicMock(result=r) for r in (outcome_results or [])] return s @@ -156,39 +162,80 @@ 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_result = hook.poll_session_completion(SESSION_ID) + assert poll_result.done is True + assert poll_result.error_message is not None + assert poll_result.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(SESSION_ID, 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_result = hook.poll_session_completion(SESSION_ID, kickoff_event_id="evt_kick") + assert poll_result == SessionPollResult(done=True, error_message=mock.ANY, stop_reason=reason) + assert reason in poll_result.error_message 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) + # Stub event: the SDK's event models are a discriminated union, so there is no single + # class to spec against -- only ``type`` and ``id`` are read by the code under test. + kickoff = mock.MagicMock(type="user.message", id="evt_kick") # noqa: spec + client.beta.sessions.events.list.return_value = [kickoff] + assert hook.poll_session_completion(SESSION_ID, 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(SESSION_ID, 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_result = hook.poll_session_completion(SESSION_ID, kickoff_event_id="evt_kick") + assert poll_result == SessionPollResult( + done=True, error_message=mock.ANY, stop_reason="budget_reached" + ) + for named_cause in ("budget", "no list price"): + assert named_cause in poll_result.error_message + # The inverse matters just as much: the generic advice must not appear here. + assert "autonomous agent" not in poll_result.error_message + + +class TestCreateSessionError: + def test_budget_reached_maps_to_budget_exception(self): + err = _create_session_error("over budget", "budget_reached") + assert isinstance(err, AnthropicSessionBudgetExceeded) + assert str(err) == "over budget" + + @pytest.mark.parametrize("stop_reason", [None, "requires_action", "retries_exhausted", "end_turn"]) + def test_other_reasons_map_to_generic_session_error(self, stop_reason): + err = _create_session_error("boom", stop_reason) + assert isinstance(err, AnthropicAgentSessionError) + assert not isinstance(err, AnthropicSessionBudgetExceeded) + + def test_budget_exception_is_catchable_as_session_error(self): + # Existing `except AnthropicAgentSessionError` handlers must keep catching it. + assert isinstance(_create_session_error("over budget", "budget_reached"), AnthropicAgentSessionError) + class TestDefaultModel: def test_defaults_to_constant(self): @@ -304,7 +351,10 @@ class TestWaitForSession: @mock.patch.object(AnthropicHook, "get_connection") def test_returns_when_done(self, mock_get_connection, mock_poll, mock_sleep): mock_get_connection.return_value = _conn() - mock_poll.side_effect = [(False, None), (True, None)] + mock_poll.side_effect = [ + SessionPollResult(done=False, error_message=None, stop_reason=None), + SessionPollResult(done=True, error_message=None, stop_reason="end_turn"), + ] AnthropicHook().wait_for_session("sess_1", poll_interval=0.01) assert mock_poll.call_count == 2 @@ -314,7 +364,7 @@ class TestWaitForSession: @mock.patch.object(AnthropicHook, "get_connection") def test_raises_on_timeout(self, mock_get_connection, mock_poll, mock_sleep, mock_monotonic): mock_get_connection.return_value = _conn() - mock_poll.return_value = (False, None) + mock_poll.return_value = SessionPollResult(done=False, error_message=None, stop_reason=None) mock_monotonic.side_effect = [0, 100] with pytest.raises(AnthropicAgentSessionTimeout, match="did not reach a terminal status"): AnthropicHook().wait_for_session("sess_1", poll_interval=0.01, timeout=10) @@ -324,10 +374,25 @@ class TestWaitForSession: @mock.patch.object(AnthropicHook, "get_connection") def test_failure_raises(self, mock_get_connection, mock_poll, mock_sleep): mock_get_connection.return_value = _conn() - mock_poll.return_value = (True, "Outcome not satisfied for session sess_1: failed.") + mock_poll.return_value = SessionPollResult( + done=True, error_message="Outcome not satisfied for session sess_1: failed.", stop_reason=None + ) with pytest.raises(AnthropicAgentSessionError, match="not satisfied"): AnthropicHook().wait_for_session("sess_1", expect_outcome=True, poll_interval=0.01) + @mock.patch(f"{HOOK_PATH}.time.sleep", autospec=True) + @mock.patch.object(AnthropicHook, "poll_session_completion", autospec=True) + @mock.patch.object(AnthropicHook, "get_connection", autospec=True) + def test_budget_stop_raises_budget_exception(self, mock_get_connection, mock_poll, mock_sleep): + mock_get_connection.return_value = _conn() + mock_poll.return_value = SessionPollResult( + done=True, + error_message="Session sess_1 stopped against its budget.", + stop_reason="budget_reached", + ) + with pytest.raises(AnthropicSessionBudgetExceeded, match="budget"): + AnthropicHook().wait_for_session("sess_1", poll_interval=0.01) + class TestAnthropicHookGetConn: @mock.patch(f"{HOOK_PATH}.Anthropic") diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py index 3c4f595ad51..d269b1bfc2d 100644 --- a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py +++ b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py @@ -24,6 +24,7 @@ from airflow.exceptions import TaskDeferred from airflow.providers.anthropic.exceptions import ( AnthropicAgentSessionError, AnthropicAgentSessionTimeout, + AnthropicSessionBudgetExceeded, AnthropicTriggerEventError, ) from airflow.providers.anthropic.hooks.anthropic import AnthropicHook @@ -214,6 +215,35 @@ class TestExecuteComplete: op.execute_complete({}, {"status": "error", "session_id": "s", "message": "boom"}) hook.archive_session.assert_called_once_with("s") + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_budget_stop_raises_budget_exception(self, mock_hook_prop): + # The deferrable path must raise the same class the synchronous path does; the + # trigger event's stop_reason is the only classification channel across the boundary. + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + with pytest.raises(AnthropicSessionBudgetExceeded, match="over budget"): + op.execute_complete( + {}, + { + "status": "error", + "session_id": "s", + "message": "over budget", + "stop_reason": "budget_reached", + }, + ) + hook.archive_session.assert_called_once_with("s") + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_error_event_without_stop_reason_raises_generic(self, mock_hook_prop): + # Version skew: a trigger serialized before stop_reason existed omits the key. + hook = mock.MagicMock(spec=AnthropicHook) + mock_hook_prop.return_value = hook + op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") + with pytest.raises(AnthropicAgentSessionError) as exc: + op.execute_complete({}, {"status": "error", "session_id": "s", "message": "boom"}) + assert not isinstance(exc.value, AnthropicSessionBudgetExceeded) + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) def test_timeout_archives_and_raises(self, mock_hook_prop): hook = mock.MagicMock(spec=AnthropicHook) diff --git a/providers/anthropic/tests/unit/anthropic/test_exceptions.py b/providers/anthropic/tests/unit/anthropic/test_exceptions.py index 365ba875eb4..bce1deea58b 100644 --- a/providers/anthropic/tests/unit/anthropic/test_exceptions.py +++ b/providers/anthropic/tests/unit/anthropic/test_exceptions.py @@ -24,6 +24,7 @@ from airflow.providers.anthropic.exceptions import ( AnthropicBatchJobError, AnthropicBatchTimeout, AnthropicError, + AnthropicSessionBudgetExceeded, ) @@ -35,6 +36,7 @@ from airflow.providers.anthropic.exceptions import ( AnthropicBatchTimeout, AnthropicAgentSessionError, AnthropicAgentSessionTimeout, + AnthropicSessionBudgetExceeded, ], ) def test_provider_errors_share_base_and_are_not_airflow_exceptions(exc): diff --git a/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py b/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py index 6e01b4259b8..aa5df4b105e 100644 --- a/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py +++ b/providers/anthropic/tests/unit/anthropic/triggers/test_agent.py @@ -21,6 +21,7 @@ from unittest import mock import pytest +from airflow.providers.anthropic.hooks.anthropic import SessionPollResult from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrigger from airflow.triggers.base import TriggerEvent @@ -29,6 +30,9 @@ pytest.importorskip("anthropic") TRIGGER_PATH = "airflow.providers.anthropic.triggers.agent" POLL = "airflow.providers.anthropic.hooks.anthropic.AnthropicHook.poll_session_completion" +STILL_RUNNING = SessionPollResult(done=False, error_message=None, stop_reason=None) +COMPLETED = SessionPollResult(done=True, error_message=None, stop_reason="end_turn") + def _trigger(end_time=None, expect_outcome=False): return AnthropicAgentSessionTrigger( @@ -66,7 +70,7 @@ async def test_on_kill_archives_session(mock_hook_cls): @pytest.mark.asyncio @mock.patch(POLL) async def test_done_success_yields_success(mock_poll): - mock_poll.return_value = (True, None) + mock_poll.return_value = COMPLETED event = await _trigger().run().__anext__() assert event.payload["status"] == "success" assert event.payload["session_id"] == "sess_1" @@ -75,7 +79,9 @@ async def test_done_success_yields_success(mock_poll): @pytest.mark.asyncio @mock.patch(POLL) async def test_done_error_yields_error(mock_poll): - mock_poll.return_value = (True, "Session sess_1 terminated.") + mock_poll.return_value = SessionPollResult( + done=True, error_message="Session sess_1 terminated.", stop_reason=None + ) event = await _trigger().run().__anext__() assert event.payload["status"] == "error" assert "terminated" in event.payload["message"] @@ -84,7 +90,7 @@ async def test_done_error_yields_error(mock_poll): @pytest.mark.asyncio @mock.patch(POLL) async def test_timeout_yields_timeout(mock_poll): - mock_poll.return_value = (False, None) + mock_poll.return_value = STILL_RUNNING event = await _trigger(end_time=time.time() - 1).run().__anext__() assert event.payload["status"] == "timeout" @@ -93,7 +99,7 @@ async def test_timeout_yields_timeout(mock_poll): @mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") @mock.patch(POLL) async def test_polls_until_done(mock_poll, mock_sleep): - mock_poll.side_effect = [(False, None), (False, None), (True, None)] + mock_poll.side_effect = [STILL_RUNNING, STILL_RUNNING, COMPLETED] event = await _trigger().run().__anext__() assert event.payload["status"] == "success" assert mock_poll.call_count == 3 @@ -114,15 +120,34 @@ async def test_persistent_error_yields_error_after_retries(mock_poll, mock_sleep @mock.patch(f"{TRIGGER_PATH}.asyncio.sleep") @mock.patch(POLL) async def test_transient_error_then_success(mock_poll, mock_sleep): - mock_poll.side_effect = [RuntimeError("blip"), (True, None)] + mock_poll.side_effect = [RuntimeError("blip"), COMPLETED] event = await _trigger().run().__anext__() assert event.payload["status"] == "success" [email protected] [email protected](POLL, autospec=True) +async def test_budget_stop_carries_stop_reason(mock_poll): + # The resuming worker needs the reason to raise AnthropicSessionBudgetExceeded rather + # than the generic session error; the message text is not a classification channel. + mock_poll.return_value = SessionPollResult( + done=True, + error_message="Session sess_1 stopped against its budget.", + stop_reason="budget_reached", + ) + event = await _trigger().run().__anext__() + assert event.payload["status"] == "error" + assert event.payload["stop_reason"] == "budget_reached" + + @pytest.mark.asyncio @mock.patch(POLL) async def test_outcome_failure_yields_error(mock_poll): - mock_poll.return_value = (True, "Outcome not satisfied for session sess_1: max_iterations_reached.") + mock_poll.return_value = SessionPollResult( + done=True, + error_message="Outcome not satisfied for session sess_1: max_iterations_reached.", + stop_reason=None, + ) event = await _trigger(expect_outcome=True).run().__anext__() assert event.payload["status"] == "error" assert "max_iterations_reached" in event.payload["message"] diff --git a/uv.lock b/uv.lock index 940a928dcab..a5ea95aa159 100644 --- a/uv.lock +++ b/uv.lock @@ -932,7 +932,7 @@ wheels = [ [[package]] name = "anthropic" -version = "0.117.0" +version = "0.121.0" source = { registry = "https://pypi.org/simple"; } dependencies = [ { name = "anyio" }, @@ -944,9 +944,9 @@ dependencies = [ { name = "sniffio" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/41/0d/8f71d535edb0d438f023bd825fb65f67c14fa88a2bd6b75f292a58a63de4/anthropic-0.117.0.tar.gz";, hash = "sha256:98107f2b76439641e0ae2a1754087534b8f178dbab99d6eb1bc4b7bc8c744496", size = 989933, upload-time = "2026-07-16T19:36:13.07Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/ca/3cb2c20ee729736fbd4546d5d8b67e818288529fe70cb7a80dbf80aef70b/anthropic-0.121.0.tar.gz";, hash = "sha256:e79d6e08ab3376602fc9a70d4d5ea3540817c76cf7e16658bed790834e1833d6", size = 1013292, upload-time = "2026-08-07T17:11:07.241Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/4c/917d21d6619a4475cdafc6d13a69fdb3b901ddac57e76caca5a25c117b6d/anthropic-0.117.0-py3-none-any.whl";, hash = "sha256:451a0a6905f11dff7663d13e4ee5dbf909eb8942b1d049803c7b937a13ac47ec", size = 998327, upload-time = "2026-07-16T19:36:11.225Z" }, + { url = "https://files.pythonhosted.org/packages/fa/91/b3d41643f1f639927e8c5fb02c3bd8bffe6f1f29e219b3bd4c61e267b15c/anthropic-0.121.0-py3-none-any.whl";, hash = "sha256:6048713fa441e59e1cba8363171cd2a86273b25bd213e9c7ac70a523af88b011", size = 1035493, upload-time = "2026-08-07T17:11:08.508Z" }, ] [package.optional-dependencies] @@ -3255,10 +3255,10 @@ docs = [ [package.metadata] requires-dist = [ - { name = "anthropic", specifier = ">=0.101.0" }, - { name = "anthropic", extras = ["aws"], marker = "extra == 'aws'", specifier = ">=0.101.0" }, - { name = "anthropic", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=0.101.0" }, - { name = "anthropic", extras = ["vertex"], marker = "extra == 'vertex'", specifier = ">=0.101.0" }, + { name = "anthropic", specifier = ">=0.121.0" }, + { name = "anthropic", extras = ["aws"], marker = "extra == 'aws'", specifier = ">=0.121.0" }, + { name = "anthropic", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=0.121.0" }, + { name = "anthropic", extras = ["vertex"], marker = "extra == 'vertex'", specifier = ">=0.121.0" }, { name = "apache-airflow", editable = "." }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, ]
