This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch anthropic-agent-config-docs in repository https://gitbox.apache.org/repos/asf/airflow.git
commit aeeae1fed77decf2a7bb5af82a04873ea0f79d88 Author: Kaxil Naik <[email protected]> AuthorDate: Tue Aug 11 19:28:59 2026 +0530 Record Anthropic agent session token usage and cost in XCom A session 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. Measured against the live API, a $0.01 ceiling admitted between $0.32 and $0.61 of usage across four runs. Setting a budget therefore tells you nothing about what a run actually cost, and the operator returned only a session ID. ``AnthropicAgentSessionOperator`` now pushes the session's usage to XCom under ``usage``, making cost per Dag run queryable: {"input_tokens": 827, "output_tokens": 14002, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, "server_tool_use": {"web_search_requests": 0, "web_fetch_requests": 0}, "active_seconds": 213.7, "list_cost": {"amount": "36", "currency": "USD"}, "try_number": 1} ``AnthropicHook.get_session_usage`` flattens the SDK models to plain scalars so the value survives XCom serialization. ``amount`` stays the API's minor-unit string rather than becoming a float, because a cost figure must not pick up binary rounding. ``list_cost`` is ``None`` when usage includes a model with no list price -- and that is exactly when a caller has to price the run from usage instead, so every billable dimension is reported: cache writes are billed above base input and server tool calls are billed per request. Usage is recorded on failure as well as success, since a budget-stopped session is precisely the one whose spend needs recording. Teardown runs first on the failure paths: it is the time-critical call, and ``sessions.archive`` returns the session carrying its final usage, so ``summarize_usage`` reports spend without a second request against an API that may be why the task is failing. If archiving fails, usage falls back to a fetch. The whole push is best effort -- it runs immediately before re-raising, so anything that throws there would otherwise replace the exception the task should fail with. Airflow clears XCom at the start of each attempt, so ``usage`` holds the final attempt only; ``try_number`` records which, and the docs say total spend across retries must be summed from the session records. Verified end to end against the live API: a real deferred run resumed by a real triggerer failed with AnthropicSessionBudgetExceeded and left the usage above in XCom, read back over the REST API. --- providers/anthropic/docs/operators/anthropic.rst | 35 ++++ .../airflow/providers/anthropic/hooks/anthropic.py | 43 ++++- .../airflow/providers/anthropic/operators/agent.py | 84 +++++++++- .../tests/unit/anthropic/hooks/test_anthropic.py | 102 +++++++++++- .../tests/unit/anthropic/operators/test_agent.py | 182 ++++++++++++++++++--- 5 files changed, 412 insertions(+), 34 deletions(-) diff --git a/providers/anthropic/docs/operators/anthropic.rst b/providers/anthropic/docs/operators/anthropic.rst index d4498e6864c..0537ed42408 100644 --- a/providers/anthropic/docs/operators/anthropic.rst +++ b/providers/anthropic/docs/operators/anthropic.rst @@ -248,6 +248,41 @@ rather than treated as a fault. budgeted sessions: the operator archives a budget-stopped session, so there is no running session left to raise the ceiling on. +Recording what a session actually spent +""""""""""""""""""""""""""""""""""""""" + +Because the ceiling is not a cap, it does not tell you the spend. The operator pushes the +session's usage to XCom under ``usage`` on **both** success and failure, so cost per Dag run +can be queried and a budget-stopped run still records what it consumed: + +.. code-block:: python + + { + "input_tokens": 827, + "output_tokens": 17065, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "server_tool_use": {"web_search_requests": 0, "web_fetch_requests": 0}, + "active_seconds": 91.2, + "list_cost": {"amount": "44", "currency": "USD"}, + "try_number": 1, + } + +``amount`` is the API's **minor-unit string** (``"44"`` is $0.44), kept as a string so no +rounding is applied to a cost figure. ``list_cost`` is ``None`` when usage includes a model +with no list price -- which is precisely when a caller has to price the run from the token +counts, so every billable dimension is reported: cache *writes* are billed above base input, +and server tool calls are billed per request. Reading usage is best effort: if it fails, the +task's real outcome is preserved and a warning is logged. + +.. warning:: + + Airflow clears a task's XCom at the start of every attempt, so ``usage`` holds the + **final attempt only** and ``try_number`` records which one that was. With retries + enabled, total spend across attempts is not recoverable from this key; sum it from the + session records instead. This is the same scenario as the retry warning above, so + ``retries=0`` keeps both problems away. + .. exampleinclude:: /../tests/system/anthropic/example_anthropic_agent.py :language: python :dedent: 4 diff --git a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py index bae659dfb2f..f3e250c572c 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py +++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py @@ -660,6 +660,40 @@ class AnthropicHook(BaseHook): self._require_first_party("Managed Agents") return self._first_party_conn.beta.sessions.retrieve(session_id) + def get_session_usage(self, session_id: str) -> dict[str, Any]: + """ + Return a JSON-serializable token/cost summary for a session. + + Plain scalars and a nested ``list_cost`` mapping rather than SDK models, so the + result survives XCom serialization and can be queried across runs. ``amount`` is + kept as the API's **minor-unit string** (``"44"`` is $0.44) rather than converted to + a float, so no rounding is applied to a cost figure. + + Every field is optional server-side -- ``list_cost`` is absent when usage includes a + model with no list price -- so missing values come back as ``None``. That absence is + why every billable dimension is reported and not just the token totals: it is + exactly when a caller has to reconstruct cost from usage that the breakdown must be + complete. Cache *writes* (``cache_creation``) are billed above base input, and + server tool calls are billed per request. + """ + return self.summarize_usage(self.get_session(session_id)) + + @staticmethod + def summarize_usage(session: BetaManagedAgentsSession) -> dict[str, Any]: + """ + Flatten an already-retrieved session's usage; see :meth:`get_session_usage`. + + Split out because ``sessions.archive`` also returns the session, so a caller that + is tearing a session down can report its usage without a second request. + + Dumps the model rather than copying a fixed list of fields. The usage model sets + ``extra="allow"``, so a billable dimension added by the API is kept on the object -- + an allowlist here would drop it silently, which is worst precisely when ``list_cost`` + is ``None`` and a caller has to price the run from the breakdown. ``mode="json"`` + keeps the result XCom-safe and leaves ``amount`` a minor-unit string. + """ + return session.usage.model_dump(mode="json") + def send_event(self, session_id: str, event: dict[str, Any]) -> Any: """Send a single event (e.g. a ``user.message`` or ``user.define_outcome``).""" self._require_first_party("Managed Agents") @@ -668,8 +702,13 @@ class AnthropicHook(BaseHook): session_id, events=cast("list[BetaManagedAgentsEventParams]", [event]) ) - def archive_session(self, session_id: str) -> Any: - """Archive a session (frees the server-side container). Best-effort teardown.""" + def archive_session(self, session_id: str) -> BetaManagedAgentsSession: + """ + Archive a session (frees the server-side container). Best-effort teardown. + + Returns the archived session, which carries its final ``usage`` -- so a caller + tearing a session down does not need a separate retrieve to report what it spent. + """ self._require_first_party("Managed Agents") return self._first_party_conn.beta.sessions.archive(session_id) diff --git a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py index f214404c71c..a57302e8bed 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py +++ b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py @@ -67,6 +67,11 @@ class AnthropicAgentSessionOperator(BaseOperator): Outputs the agent writes to ``/mnt/session/outputs/`` are retrieved afterwards via the Files API (``scope_id=<session_id>``); the operator returns the **session ID only**. + The session's token counts and list cost are pushed to XCom under ``usage`` on both + success and failure, so cost per Dag run can be queried. This matters because a budget is + a stop trigger rather than a cap: the ceiling is checked between model requests, so it + does not tell you what the session actually spent. + .. seealso:: For more information, take a look at the guide: :ref:`howto/operator:AnthropicAgentSessionOperator` @@ -178,7 +183,10 @@ class AnthropicAgentSessionOperator(BaseOperator): ) except Exception: # send_event failed after create_session allocated the container; tear it down. - self._archive_session(session.id) + # The event may still have been accepted server-side, so the agent can already be + # spending -- record usage here too rather than leaving this the one failure path + # with no cost trail. + self._tear_down(context, session.id) raise # Correlate completion against the kickoff event so a message run is not fooled by a # just-created idle session (the start race). @@ -216,8 +224,11 @@ class AnthropicAgentSessionOperator(BaseOperator): except Exception: # Any failure after the session starts (timeout, SDK 5xx, auth expiry) leaves # the server-side container running; archive it best-effort before failing. - self._archive_session(session.id) + # Teardown goes first: it is the time-critical call, and its response carries + # the usage, so reporting spend costs no extra request. + self._tear_down(context, session.id) raise + self._push_usage(context, session.id) return session.id def execute_complete(self, context: Context, event: Any = None) -> str: @@ -226,24 +237,83 @@ class AnthropicAgentSessionOperator(BaseOperator): self.session_id = event["session_id"] status = event["status"] if status == "timeout": - self._archive_session(self.session_id) + self._tear_down(context, self.session_id) raise AnthropicAgentSessionTimeout(event["message"]) if status == "error": # 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) + self._tear_down(context, self.session_id) raise _create_session_error(event["message"], event.get("stop_reason")) self.log.info("Session %s completed.", self.session_id) + self._push_usage(context, self.session_id) return self.session_id - def _archive_session(self, session_id: str | None) -> None: - """Best-effort teardown of the server-side session (frees its container).""" + def _tear_down(self, context: Context, session_id: str | None) -> None: + """ + Archive the session and record what it spent, in that order. + + Every failure path needs both. Teardown goes first because it is the + time-critical call, and ``sessions.archive`` returns the session carrying its final + usage, so recording spend costs no extra request. + """ + archived = self._archive_session(session_id) + self._push_usage(context, session_id, session=archived) + + def _push_usage(self, context: Context, session_id: str | None, session: Any = None) -> None: + """ + Push the session's token/cost usage to XCom under ``usage``, best effort. + + Runs on failure as well as success: a session that stopped against its budget is + exactly the one whose spend you want recorded. Never allowed to raise -- a usage + read that fails must not mask the task's real outcome. + + On the failure paths the session has just been archived, and ``sessions.archive`` + returns the session with its final usage; pass it as ``session`` so teardown is not + delayed by an extra retrieve against an API that may be why the task is failing. + """ if not session_id: return + # The whole body is guarded, not just the API call: this runs on the failure path + # immediately before re-raising, so anything that throws here -- an unavailable + # session, a serialization error, a context without a task instance -- would + # replace the exception the task should actually fail with. + try: + # XCom is cleared at the start of every attempt, so this key only ever holds + # the last one. Stamping the attempt makes that visible rather than silently + # under-reporting total spend across retries. Built as a new dict rather than + # mutating what the hook returned. + summary = ( + self.hook.summarize_usage(session) + if session is not None + else self.hook.get_session_usage(session_id) + ) + usage = {**summary, "try_number": getattr(context["ti"], "try_number", None)} + context["ti"].xcom_push(key="usage", value=usage) + cost = usage.get("list_cost") + self.log.info( + "Session %s used %s input / %s output tokens; list cost %s", + session_id, + usage.get("input_tokens"), + usage.get("output_tokens"), + f"{cost['currency']} {cost['amount']} (minor units)" if cost else "unavailable", + ) + except Exception as e: + self.log.warning("Could not record usage for session %s: %s", session_id, e) + + def _archive_session(self, session_id: str | None) -> Any: + """ + Best-effort teardown of the server-side session (frees its container). + + Returns the archived session, which carries its final usage, or ``None`` if the + archive call failed. + """ + if not session_id: + return None try: - self.hook.archive_session(session_id) + return self.hook.archive_session(session_id) except Exception as e: self.log.warning("Failed to archive session %s: %s", session_id, e) + return None def on_kill(self) -> None: """ diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py index a3867899fae..9253b4c1407 100644 --- a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py +++ b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py @@ -17,6 +17,7 @@ from __future__ import annotations from decimal import Decimal +from types import SimpleNamespace from unittest import mock import pytest @@ -44,6 +45,12 @@ from airflow.providers.anthropic.hooks.anthropic import ( pytest.importorskip("anthropic") +from anthropic.types import BetaMonetaryAmount +from anthropic.types.beta import BetaManagedAgentsServerToolUsage, BetaManagedAgentsSessionUsage +from anthropic.types.beta.beta_managed_agents_cache_creation_usage import ( + BetaManagedAgentsCacheCreationUsage, +) + 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. @@ -321,7 +328,100 @@ class TestCreateSessionBudget: assert "budget" not in client.beta.sessions.create.call_args.kwargs -class TestCreateSessionError: +class TestGetSessionUsage: + """ + Built on the real SDK usage models rather than mocks. + + Constructing the real model is not on its own a check on field names: it sets + ``extra="allow"``, so a misspelling is accepted as an extra field rather than rejected. + ``test_usage_fields_exist_on_the_sdk_model`` is what pins the names to the declared + schema. Only the session wrapper is a stand-in, since ``get_session_usage`` reads + nothing from it but ``.usage``. + """ + + def test_usage_fields_exist_on_the_sdk_model(self): + # The reported breakdown must be fields the SDK actually declares. active_seconds + # lived on session.stats rather than session.usage as recently as 0.117.0, and the + # floor is an open-ended >= on a beta API -- if a field moves again, summarize_usage + # would return it as None (or raise, swallowed into a warning) with no other signal. + assert set(BetaManagedAgentsSessionUsage.model_fields) >= { + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation", + "server_tool_use", + "active_seconds", + "list_cost", + } + + @staticmethod + def _create_session(list_cost, cache_creation=None, server_tool_use=None): + usage = BetaManagedAgentsSessionUsage( + input_tokens=827, + output_tokens=17065, + cache_read_input_tokens=0, + active_seconds=91.2, + list_cost=list_cost, + cache_creation=cache_creation, + server_tool_use=server_tool_use, + ) + return SimpleNamespace(usage=usage) + + def test_flattens_to_json_safe_scalars(self): + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = self._create_session( + BetaMonetaryAmount(amount="44", currency="USD") + ) + assert hook.get_session_usage("s") == { + "input_tokens": 827, + "output_tokens": 17065, + "cache_read_input_tokens": 0, + "cache_creation": None, + "server_tool_use": None, + "active_seconds": 91.2, + "list_cost": {"amount": "44", "currency": "USD"}, + } + + def test_reports_every_billable_dimension(self): + # list_cost is None exactly when a caller must price the run from usage, so the + # expensive side (cache writes, server tool calls) cannot be missing. + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = self._create_session( + None, + cache_creation=BetaManagedAgentsCacheCreationUsage( + ephemeral_5m_input_tokens=120, ephemeral_1h_input_tokens=340 + ), + server_tool_use=BetaManagedAgentsServerToolUsage(web_search_requests=3, web_fetch_requests=5), + ) + assert hook.get_session_usage("s") == { + "input_tokens": 827, + "output_tokens": 17065, + "cache_read_input_tokens": 0, + "cache_creation": { + "ephemeral_5m_input_tokens": 120, + "ephemeral_1h_input_tokens": 340, + }, + "server_tool_use": {"web_search_requests": 3, "web_fetch_requests": 5}, + "active_seconds": 91.2, + "list_cost": None, + } + + def test_amount_stays_a_minor_unit_string(self): + # Never floated: a cost figure must not pick up binary rounding. + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = self._create_session( + BetaMonetaryAmount(amount="44", currency="USD") + ) + assert hook.get_session_usage("s")["list_cost"]["amount"] == "44" + + def test_absent_list_cost_is_none(self): + # Happens when usage includes a model with no list price. + hook, client = _make_hook() + client.beta.sessions.retrieve.return_value = self._create_session(None) + assert hook.get_session_usage("s")["list_cost"] is None + + +class TestSessionError: def test_budget_reached_maps_to_budget_exception(self): err = _create_session_error("over budget", "budget_reached") assert isinstance(err, AnthropicSessionBudgetExceeded) diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py index d14472121a9..d7a05e8217c 100644 --- a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py +++ b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py @@ -34,8 +34,10 @@ from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrig pytest.importorskip("anthropic") -def _context(): - return {"ti": mock.MagicMock()} +def _create_context(try_number=1): + ti = mock.MagicMock() + ti.try_number = try_number + return {"ti": ti} def _create_op(**kwargs) -> AnthropicAgentSessionOperator: @@ -54,13 +56,13 @@ def _create_op(**kwargs) -> AnthropicAgentSessionOperator: def test_requires_exactly_one_of_message_or_outcome(): op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env") with pytest.raises(ValueError, match="exactly one"): - op.execute(_context()) + op.execute(_create_context()) op = AnthropicAgentSessionOperator( task_id="a", agent_id="ag", environment_id="env", message="hi", outcome={"description": "x"} ) with pytest.raises(ValueError, match="exactly one"): - op.execute(_context()) + op.execute(_create_context()) def test_outcome_requires_description_and_rubric(): @@ -69,7 +71,7 @@ def test_outcome_requires_description_and_rubric(): task_id="a", agent_id="ag", environment_id="env", outcome={"description": "x"} ) with pytest.raises(ValueError, match="description.*rubric"): - op.execute(_context()) + op.execute(_create_context()) # missing description op = AnthropicAgentSessionOperator( @@ -79,7 +81,7 @@ def test_outcome_requires_description_and_rubric(): outcome={"rubric": {"type": "text", "content": "c"}}, ) with pytest.raises(ValueError, match="description.*rubric"): - op.execute(_context()) + op.execute(_create_context()) def test_init_does_not_validate_message_or_outcome(): @@ -104,14 +106,20 @@ class TestExecute: op = AnthropicAgentSessionOperator( task_id="a", agent_id="ag", environment_id="env", message="summarize", deferrable=False ) - context = _context() + context = _create_context() assert op.execute(context) == "sess_1" hook.create_session.assert_called_once_with(agent="ag", environment_id="env") hook.send_event.assert_called_once_with( "sess_1", {"type": "user.message", "content": [{"type": "text", "text": "summarize"}]} ) hook.wait_for_session.assert_called_once() - context["ti"].xcom_push.assert_called_once_with(key="session_id", value="sess_1") + # Two pushes now: the session id up front, then usage once the run finishes. + # Asserting the keys in order keeps this stricter than assert_any_call would. + assert [c.kwargs["key"] for c in context["ti"].xcom_push.call_args_list] == [ + "session_id", + "usage", + ] + context["ti"].xcom_push.assert_any_call(key="session_id", value="sess_1") @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) def test_outcome_sends_define_outcome(self, mock_hook_prop): @@ -123,7 +131,7 @@ class TestExecute: op = AnthropicAgentSessionOperator( task_id="a", agent_id="ag", environment_id="env", outcome=outcome, deferrable=False ) - op.execute(_context()) + op.execute(_create_context()) hook.send_event.assert_called_once_with("sess_1", {"type": "user.define_outcome", **outcome}) @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -141,7 +149,7 @@ class TestExecute: vault_ids=["vlt_1"], session_resources=[{"type": "file", "file_id": "f1", "mount_path": "/workspace/f"}], ) - op.execute(_context()) + op.execute(_create_context()) hook.create_session.assert_called_once_with( agent="ag", environment_id="env", @@ -160,7 +168,7 @@ class TestExecute: task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=False ) with pytest.raises(AnthropicAgentSessionTimeout, match="too slow"): - op.execute(_context()) + op.execute(_create_context()) hook.archive_session.assert_called_once_with("sess_1") @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -176,7 +184,7 @@ class TestExecute: task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=False ) with pytest.raises(RuntimeError, match="api 5xx"): - op.execute(_context()) + op.execute(_create_context()) hook.archive_session.assert_called_once_with("sess_1") @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -191,7 +199,7 @@ class TestExecute: task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=False ) with pytest.raises(RuntimeError, match="send boom"): - op.execute(_context()) + op.execute(_create_context()) hook.archive_session.assert_called_once_with("sess_1") hook.wait_for_session.assert_not_called() @@ -205,7 +213,7 @@ class TestExecute: task_id="a", agent_id="ag", environment_id="env", message="hi", deferrable=True ) with pytest.raises(TaskDeferred) as exc: - op.execute(_context()) + op.execute(_create_context()) assert isinstance(exc.value.trigger, AnthropicAgentSessionTrigger) assert exc.value.trigger.session_id == "sess_1" assert exc.value.method_name == "execute_complete" @@ -217,7 +225,7 @@ class TestBudgetParam: def test_dollar_amount_is_converted_to_minor_units(self, mock_hook_prop): hook = mock.MagicMock(spec=AnthropicHook) mock_hook_prop.return_value = hook - _create_op(budget=25).execute(_context()) + _create_op(budget=25).execute(_create_context()) assert hook.create_session.call_args.kwargs["budget"] == { "type": "limit", "max_list_cost": {"amount": "2500", "currency": "USD"}, @@ -228,14 +236,14 @@ class TestBudgetParam: hook = mock.MagicMock(spec=AnthropicHook) mock_hook_prop.return_value = hook raw = {"type": "limit", "max_list_cost": {"amount": "750", "currency": "USD"}} - _create_op(budget=raw).execute(_context()) + _create_op(budget=raw).execute(_create_context()) assert hook.create_session.call_args.kwargs["budget"] == raw @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) def test_no_budget_key_when_unset(self, mock_hook_prop): hook = mock.MagicMock(spec=AnthropicHook) mock_hook_prop.return_value = hook - _create_op().execute(_context()) + _create_op().execute(_create_context()) assert "budget" not in hook.create_session.call_args.kwargs @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -244,7 +252,7 @@ class TestBudgetParam: mock_hook_prop.return_value = hook op = _create_op(budget=25, session_kwargs={"budget": {"type": "limit"}}) with pytest.raises(ValueError, match="not both"): - op.execute(_context()) + op.execute(_create_context()) hook.create_session.assert_not_called() @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -253,17 +261,141 @@ class TestBudgetParam: hook = mock.MagicMock(spec=AnthropicHook) mock_hook_prop.return_value = hook with pytest.raises(ValueError, match="positive"): - _create_op(budget=-1).execute(_context()) + _create_op(budget=-1).execute(_create_context()) hook.create_session.assert_not_called() def test_budget_is_templated(self): assert "budget" in AnthropicAgentSessionOperator.template_fields +class TestUsageXCom: + USAGE = { + "input_tokens": 827, + "output_tokens": 17065, + "cache_read_input_tokens": 0, + "active_seconds": 91.2, + "list_cost": {"amount": "44", "currency": "USD"}, + } + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_pushed_on_success(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.get_session_usage.return_value = self.USAGE + mock_hook_prop.return_value = hook + + context = _create_context() + _create_op().execute(context) + context["ti"].xcom_push.assert_any_call(key="usage", value={**self.USAGE, "try_number": 1}) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_read_from_the_archive_response_after_teardown(self, mock_hook_prop): + # Teardown is the time-critical call on a failure path, so it goes first; its + # response carries the usage, so recording spend costs no extra request. + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.wait_for_session.side_effect = AnthropicSessionBudgetExceeded("over budget") + archived = object() + calls = [] + hook.archive_session.side_effect = lambda *a, **k: calls.append("archive") or archived + hook.summarize_usage.side_effect = lambda *a, **k: calls.append("usage") or self.USAGE + mock_hook_prop.return_value = hook + + context = _create_context() + with pytest.raises(AnthropicSessionBudgetExceeded): + _create_op().execute(context) + context["ti"].xcom_push.assert_any_call(key="usage", value={**self.USAGE, "try_number": 1}) + assert calls == ["archive", "usage"] + hook.summarize_usage.assert_called_once_with(archived) + hook.get_session_usage.assert_not_called() + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_falls_back_to_a_fetch_when_archiving_fails(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.wait_for_session.side_effect = AnthropicSessionBudgetExceeded("over budget") + hook.archive_session.side_effect = RuntimeError("archive 500") + hook.get_session_usage.return_value = self.USAGE + mock_hook_prop.return_value = hook + + context = _create_context() + with pytest.raises(AnthropicSessionBudgetExceeded, match="over budget"): + _create_op().execute(context) + context["ti"].xcom_push.assert_any_call(key="usage", value={**self.USAGE, "try_number": 1}) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_read_failure_does_not_mask_the_real_error(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.wait_for_session.side_effect = AnthropicSessionBudgetExceeded("over budget") + # summarize_usage, not get_session_usage: the failure path archives first and reads + # usage off that response, so failing the fetch here would never be reached. + hook.summarize_usage.side_effect = RuntimeError("usage api 500") + mock_hook_prop.return_value = hook + + with pytest.raises(AnthropicSessionBudgetExceeded, match="over budget"): + _create_op().execute(_create_context()) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_read_failure_does_not_break_success(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.get_session_usage.side_effect = RuntimeError("usage api 500") + mock_hook_prop.return_value = hook + + assert _create_op().execute(_create_context()) == "sess_1" + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_missing_list_cost_is_logged_as_unavailable(self, mock_hook_prop): + # list_cost is absent when usage includes a model with no list price. + hook = mock.MagicMock(spec=AnthropicHook) + hook.create_session.return_value.id = "sess_1" + hook.get_session_usage.return_value = {**self.USAGE, "list_cost": None} + mock_hook_prop.return_value = hook + + context = _create_context() + _create_op().execute(context) + context["ti"].xcom_push.assert_any_call( + key="usage", value={**self.USAGE, "list_cost": None, "try_number": 1} + ) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_pushed_on_deferrable_success(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.get_session_usage.return_value = self.USAGE + mock_hook_prop.return_value = hook + + context = _create_context() + op = _create_op() + assert op.execute_complete(context, {"status": "success", "session_id": "sess_1"}) == "sess_1" + context["ti"].xcom_push.assert_any_call(key="usage", value={**self.USAGE, "try_number": 1}) + + @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) + def test_usage_pushed_on_deferrable_budget_error(self, mock_hook_prop): + hook = mock.MagicMock(spec=AnthropicHook) + hook.summarize_usage.return_value = self.USAGE + mock_hook_prop.return_value = hook + + context = _create_context() + with pytest.raises(AnthropicSessionBudgetExceeded): + _create_op().execute_complete( + context, + { + "status": "error", + "session_id": "sess_1", + "message": "over budget", + "stop_reason": "budget_reached", + }, + ) + context["ti"].xcom_push.assert_any_call(key="usage", value={**self.USAGE, "try_number": 1}) + + class TestExecuteComplete: def test_success_returns_session_id(self): op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") - assert op.execute_complete({}, {"status": "success", "session_id": "sess_1"}) == "sess_1" + assert ( + op.execute_complete(_create_context(), {"status": "success", "session_id": "sess_1"}) == "sess_1" + ) @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) def test_error_archives_and_raises(self, mock_hook_prop): @@ -273,7 +405,7 @@ class TestExecuteComplete: mock_hook_prop.return_value = hook op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") with pytest.raises(AnthropicAgentSessionError, match="boom"): - op.execute_complete({}, {"status": "error", "session_id": "s", "message": "boom"}) + op.execute_complete(_create_context(), {"status": "error", "session_id": "s", "message": "boom"}) hook.archive_session.assert_called_once_with("s") @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -302,7 +434,7 @@ class TestExecuteComplete: 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"}) + op.execute_complete(_create_context(), {"status": "error", "session_id": "s", "message": "boom"}) assert not isinstance(exc.value, AnthropicSessionBudgetExceeded) @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) @@ -311,7 +443,9 @@ class TestExecuteComplete: mock_hook_prop.return_value = hook op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") with pytest.raises(AnthropicAgentSessionTimeout): - op.execute_complete({}, {"status": "timeout", "session_id": "s", "message": "slow"}) + op.execute_complete( + _create_context(), {"status": "timeout", "session_id": "s", "message": "slow"} + ) hook.archive_session.assert_called_once_with("s") @pytest.mark.parametrize( @@ -328,7 +462,7 @@ class TestExecuteComplete: def test_invalid_event_raises(self, event, match): op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") with pytest.raises(AnthropicTriggerEventError, match=match): - op.execute_complete({}, event) + op.execute_complete(_create_context(), event) class TestOnKill:
