You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2026/08/13 12:06:29 UTC
(airflow) 01/02: Record Anthropic agent session token usage and cost in XCom
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 <ka...@apache.org>
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: