# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import logging
import time
from collections.abc import Mapping
from copy import deepcopy
from dataclasses import dataclass
from decimal import Decimal, InvalidOperation
from enum import Enum
from functools import cached_property
from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar, cast
from anthropic import (
Anthropic,
AnthropicAWS,
AnthropicBedrock,
AnthropicFoundry,
AnthropicVertex,
AsyncAnthropic,
AsyncAnthropicAWS,
AsyncAnthropicBedrock,
AsyncAnthropicFoundry,
AsyncAnthropicVertex,
BadRequestError,
IdentityTokenFile,
WorkloadIdentityCredentials,
)
from airflow.providers.anthropic.exceptions import (
AnthropicAgentSessionError,
AnthropicAgentSessionTimeout,
AnthropicBatchJobError,
AnthropicBatchTimeout,
AnthropicError,
AnthropicSessionBudgetExceeded,
AnthropicTriggerEventError,
)
from airflow.providers.common.compat.connection import get_async_connection
from airflow.providers.common.compat.sdk import AirflowSkipException, BaseHook
[docs]
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from collections.abc import Callable, Iterable, Iterator
from anthropic.types import Message
from anthropic.types.beta import (
BetaEnvironment,
BetaManagedAgentsAgent,
BetaManagedAgentsBudgetLimitParam,
BetaManagedAgentsSession,
agent_create_params,
environment_create_params,
)
from anthropic.types.beta.sessions import BetaManagedAgentsEventParams
from anthropic.types.messages import MessageBatch, MessageBatchIndividualResponse
from anthropic.types.messages.batch_create_params import Request
#: Default model used when an operator or hook caller does not specify one.
#: Prefer configuring the model on the connection so it can be updated without
#: a provider release when this model ID is retired.
[docs]
DEFAULT_MODEL = "claude-opus-4-8"
#: Platforms that serve the first-party-only endpoints (Message Batches, token
#: counting, the Models API). Amazon Bedrock, Google Vertex AI and Microsoft
#: Foundry do not serve these, so the hook fails fast rather than surfacing a
#: raw ``404`` from the SDK.
[docs]
AnthropicClient = Anthropic | AnthropicBedrock | AnthropicVertex | AnthropicAWS | AnthropicFoundry
[docs]
AsyncAnthropicClient = (
AsyncAnthropic | AsyncAnthropicBedrock | AsyncAnthropicVertex | AsyncAnthropicAWS | AsyncAnthropicFoundry
)
_ClientT = TypeVar("_ClientT")
@dataclass(frozen=True)
class _ClientFactories(Generic[_ClientT]):
"""
The client class to build for each platform.
:meth:`AnthropicHook.get_conn` and :meth:`AnthropicHook.get_async_conn` pass the sync and
async SDK classes through the same builder, so the two clients read the connection the
same way and cannot drift.
"""
anthropic: Callable[..., _ClientT]
bedrock: Callable[..., _ClientT]
vertex: Callable[..., _ClientT]
aws: Callable[..., _ClientT]
foundry: Callable[..., _ClientT]
#: Consecutive failed polls tolerated in the synchronous wait helpers before giving up
#: (transient errors). Mirrors the deferrable triggers' tolerance so a single blip does
#: not fail (and cancel/archive) a still-healthy batch or session.
[docs]
MAX_CONSECUTIVE_POLL_FAILURES = 5
[docs]
class BatchStatus(str, Enum):
"""Top-level ``processing_status`` of an Anthropic Message Batch."""
[docs]
IN_PROGRESS = "in_progress"
[docs]
CANCELING = "canceling"
@classmethod
[docs]
def is_in_progress(cls, status: str) -> bool:
"""
Return ``True`` while the batch has not reached the terminal ``ended`` status.
This is broader than the ``in_progress`` value: a ``canceling`` batch is also
non-terminal (cancellation is in flight but the batch has not ended yet), so it
returns ``True`` too. Read the name as "not yet terminal", not "equals the
``in_progress`` status".
"""
return status != cls.ENDED
[docs]
class SessionStatus(str, Enum):
"""Status of a Managed Agents session."""
[docs]
RESCHEDULING = "rescheduling"
[docs]
TERMINATED = "terminated"
@classmethod
[docs]
def is_terminal(cls, status: str) -> bool:
"""
Return ``True`` once the session has stopped working.
``idle`` means the agent finished its turn (done, for an autonomous run);
``terminated`` is an unrecoverable failure. Both stop the wait.
"""
return status in (cls.IDLE, cls.TERMINATED)
#: ``outcome_evaluations[].result`` values that mean the outcome did NOT succeed.
[docs]
OUTCOME_FAILURE_RESULTS = frozenset({"failed", "max_iterations_reached", "interrupted"})
# ``session.status_idle`` stop reason emitted when a session stops against its budget.
[docs]
BUDGET_REACHED = "budget_reached"
# What a caller may pass as a session budget: an amount in USD, or the raw API payload.
[docs]
BudgetSpec = str | int | float | Decimal | Mapping[str, Any]
[docs]
def build_budget(budget: BudgetSpec) -> BetaManagedAgentsBudgetLimitParam:
"""
Normalize a session budget into the API's ``max_list_cost`` payload.
A scalar is read as **US dollars** (``25``, ``25.0`` and ``"25.00"`` all mean $25.00).
The API wants minor units as an integer decimal string, so the conversion runs through
:class:`~decimal.Decimal` (never binary float) and rejects an amount finer than a cent
rather than silently rounding money. A mapping is deep-copied and otherwise returned
unchanged, so a raw payload the provider has not caught up with stays usable without a
provider release.
.. warning::
The ceiling is a stop trigger, not a cap: it is checked between model requests, so
a request already in flight can carry the session well past it.
"""
if isinstance(budget, Mapping):
# Deep, not ``dict()``: a shallow copy leaves the nested ``max_list_cost`` aliased
# to the caller's object, so a later edit of the returned payload would reach back
# into a templated operator field.
# Cast, not validate: a raw payload is accepted precisely so a caller can reach a
# field this provider does not model yet, so its shape cannot be checked here.
return cast("BetaManagedAgentsBudgetLimitParam", deepcopy(dict(budget)))
try:
dollars = Decimal(str(budget))
except (InvalidOperation, ValueError) as e:
raise ValueError(f"Invalid budget {budget!r}: not a decimal amount in USD.") from e
if not dollars.is_finite() or dollars <= 0:
raise ValueError(f"Invalid budget {budget!r}: must be a positive USD amount.")
minor_units = dollars * 100
if minor_units != minor_units.to_integral_value():
raise ValueError(
f"Invalid budget {budget!r}: amounts finer than a cent are not representable. "
"Pass a mapping if you need the raw payload."
)
return {
"type": "limit",
"max_list_cost": {"amount": str(int(minor_units)), "currency": "USD"},
}
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)
[docs]
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.
"""
[docs]
error_message: str | None
[docs]
stop_reason: str | None
[docs]
def evaluate_session_state(
session: BetaManagedAgentsSession, *, expect_outcome: bool
) -> tuple[bool, str | None, bool]:
"""
Judge a polled session from its object fields alone.
Returns ``(done, error_message, needs_event_check)``. ``done=False`` means keep
polling. ``needs_event_check=True`` means the session is ``idle`` on a ``message``
run and the object can't say *why* — the caller must inspect the event log (see
:meth:`AnthropicHook.poll_session_completion`).
The ``status`` field can't distinguish a genuine ``end_turn`` from ``requires_action``
or ``retries_exhausted``, nor a just-created ``idle``. For an outcome run the true
verdict is in ``outcome_evaluations`` (judged here, which also defeats the start race).
"""
if session.status == SessionStatus.TERMINATED:
return True, f"Session {session.id} terminated.", False
if session.status != SessionStatus.IDLE:
return False, None, False
if not expect_outcome:
return False, None, True
for evaluation in session.outcome_evaluations:
if evaluation.result == "satisfied":
return True, None, False
if evaluation.result in OUTCOME_FAILURE_RESULTS:
return True, f"Outcome not satisfied for session {session.id}: {evaluation.result}.", False
# idle but no terminal outcome verdict yet (e.g. the run has not started)
return False, None, False
#: Statuses the provider's triggers emit in their terminal event.
[docs]
TRIGGER_EVENT_STATUSES = frozenset({"success", "error", "timeout"})
[docs]
def validate_execute_complete_event(event: dict[str, Any] | None = None) -> dict[str, Any]:
"""
Validate the event a deferred task resumes with, returning it if well-formed.
The event crosses the triggerer/worker boundary through the metadata DB, so a
resuming task can receive ``None`` or a status its handlers do not recognize
(version skew, a custom trigger). Both must fail loudly: the ``execute_complete``
handlers raise on ``timeout``/``error`` and treat everything else as success, so
an unrecognized status would otherwise silently succeed.
"""
if event is None:
raise AnthropicTriggerEventError("Trigger error: event is None")
if event.get("status") not in TRIGGER_EVENT_STATUSES:
raise AnthropicTriggerEventError(
f"Unexpected trigger event status {event.get('status')!r}: {event!r}"
)
return event
[docs]
def evaluate_batch_counts(
*,
batch_id: str | None,
canceled: int,
errored: int,
expired: int,
succeeded: int,
fail_on_partial_error: bool,
) -> None:
"""
Apply the success/skip/fail policy for a terminal batch's request counts.
Lives in the hook module so both :class:`AnthropicBatchOperator` and
:class:`~airflow.providers.anthropic.sensors.batch.AnthropicBatchSensor` share it
without an operator/sensor cross-import. Raises ``AirflowSkipException`` for a
fully-cancelled batch, ``AnthropicBatchJobError`` when ``fail_on_partial_error`` and any
request failed, otherwise returns (logging a warning for partial failures).
"""
total = canceled + errored + expired + succeeded
if total and canceled == total:
raise AirflowSkipException(f"Batch {batch_id} was fully cancelled.")
failed = errored + expired
if failed:
message = (
f"Batch {batch_id} ended with {failed} failed request(s) "
f"(errored={errored}, expired={expired}, succeeded={succeeded})."
)
if fail_on_partial_error:
raise AnthropicBatchJobError(message)
logger.warning("%s Successful results are still available.", message)
[docs]
class AnthropicHook(BaseHook):
"""
Use the Anthropic SDK to interact with the Claude API.
The connection's ``password`` is used as the API key and ``host`` as an optional
base URL (for gateways/proxies). The ``extra`` field selects the platform client
and passes platform-specific configuration:
- ``platform``: one of ``anthropic`` (default), ``bedrock``, ``vertex``, ``aws``, ``foundry``.
- ``model``: default model id used when an operator/hook call omits ``model`` (lets you
change the model without editing Dags); falls back to :data:`DEFAULT_MODEL`.
- ``aws_region``: region for the ``bedrock`` and ``aws`` platforms. Required unless the
worker resolves one from ``AWS_REGION`` / ``AWS_DEFAULT_REGION`` or its AWS profile.
- ``project_id`` / ``region``: project and region for the ``vertex`` platform.
- ``resource``: Azure resource name for the ``foundry`` platform.
- ``anthropic_client_kwargs``: extra keyword arguments forwarded to the client
constructor (e.g. ``timeout``, ``max_retries``, ``default_headers``).
- ``workload_identity``: configure `Workload Identity Federation
<https://platform.claude.com/docs/en/manage-claude/workload-identity-federation>`__
(keyless OIDC auth) with ``identity_token_file``, ``federation_rule_id``,
``organization_id`` and optional ``service_account_id``, ``workspace_id`` / ``scope``.
When the ``anthropic`` platform has no API Key and no ``workload_identity`` block, the
client is built with no static credential so the SDK resolves them from the environment
— supporting env-driven Workload Identity Federation and ``ant`` profiles.
:meth:`get_conn` returns the synchronous client. Async callers, such as an agent loop or a
trigger, ``await`` :meth:`get_async_conn` for its async twin, built from the same connection.
.. seealso:: https://docs.claude.com/en/api/client-sdks
:param conn_id: :ref:`Anthropic connection id <howto/connection:anthropic>`.
"""
[docs]
conn_name_attr = "conn_id"
[docs]
default_conn_name = "anthropic_default"
[docs]
conn_type = "anthropic"
[docs]
hook_name = "Anthropic"
def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
@cached_property
def _connection(self):
return self.get_connection(self.conn_id)
@cached_property
@cached_property
[docs]
def default_model(self) -> str:
"""Default model id — connection ``extra['model']`` if set, else :data:`DEFAULT_MODEL`."""
return self._connection.extra_dejson.get("model") or DEFAULT_MODEL
@cached_property
[docs]
def conn(self) -> AnthropicClient:
"""Return the Anthropic client for the configured platform."""
return self.get_conn()
def _build_aws_client(
self,
factory: Callable[..., _ClientT],
aws_region: str | None,
client_kwargs: dict[str, Any],
) -> _ClientT:
try:
return factory(aws_region=aws_region, **client_kwargs)
except ValueError as exc:
# anthropic 1.x dropped the implicit us-east-1 fallback, so an unset region now
# fails at client construction. Point at the connection rather than the SDK.
if aws_region:
raise
raise AnthropicError(
f"No AWS region configured for the {self.platform!r} platform. Set 'aws_region' in the "
f"extra of connection {self.conn_id!r}, or set AWS_REGION / AWS_DEFAULT_REGION (or a "
"region on the AWS profile) on the worker or triggerer."
) from exc
def _build_client(self, factories: _ClientFactories[_ClientT]) -> _ClientT:
"""Build the client for the connection's platform from the given sync or async classes."""
conn = self._connection
extras = conn.extra_dejson
client_kwargs = dict(extras.get("anthropic_client_kwargs", {}))
platform = self.platform
self.log.debug("Building Anthropic client for platform %r (conn_id=%s)", platform, self.conn_id)
if platform == "bedrock":
return self._build_aws_client(factories.bedrock, extras.get("aws_region"), client_kwargs)
if platform == "vertex":
return factories.vertex(
project_id=extras.get("project_id"), region=extras.get("region"), **client_kwargs
)
if platform == "aws":
return self._build_aws_client(factories.aws, extras.get("aws_region"), client_kwargs)
if platform == "foundry":
api_key = client_kwargs.pop("api_key", None) or conn.password
return factories.foundry(api_key=api_key, resource=extras.get("resource"), **client_kwargs)
if platform != "anthropic":
raise AnthropicError(
f"Unknown Anthropic platform {platform!r}. "
"Expected one of: anthropic, bedrock, vertex, aws, foundry."
)
base_url = client_kwargs.pop("base_url", None) or conn.host or None
wif = extras.get("workload_identity")
if wif:
# The async client takes the same credential: the SDK runs a synchronous
# provider's token exchange in a worker thread, off the event loop.
return factories.anthropic(
credentials=self._workload_identity_credentials(wif), base_url=base_url, **client_kwargs
)
api_key = client_kwargs.pop("api_key", None) or conn.password
if api_key:
return factories.anthropic(api_key=api_key, base_url=base_url, **client_kwargs)
# No static key and no explicit federation config: let the SDK resolve credentials
# from the environment, which supports env-driven Workload Identity Federation
# (ANTHROPIC_FEDERATION_RULE_ID etc.) and ``ant`` profiles.
return factories.anthropic(base_url=base_url, **client_kwargs)
[docs]
def get_conn(self) -> AnthropicClient:
"""Build and return the Anthropic client for the configured platform."""
factories: _ClientFactories[AnthropicClient] = _ClientFactories(
anthropic=Anthropic,
bedrock=AnthropicBedrock,
vertex=AnthropicVertex,
aws=AnthropicAWS,
foundry=AnthropicFoundry,
)
return self._build_client(factories)
[docs]
async def get_async_conn(self) -> AsyncAnthropicClient:
"""
Build and return the async Anthropic client for the configured platform.
Reads the same connection fields as :meth:`get_conn` through the same builder, and
returns the async twin of the client it would build: ``AsyncAnthropic``,
``AsyncAnthropicBedrock``, ``AsyncAnthropicVertex``, ``AsyncAnthropicAWS`` or
``AsyncAnthropicFoundry``. The connection is looked up without blocking the event
loop, so a trigger can call this directly.
Each call builds a new client. Close it when done, with ``async with`` or
``await client.close()``, to release its HTTP connections.
"""
if "_connection" not in self.__dict__:
# Fills the cached property, so later reads of ``platform`` and ``default_model``
# use this connection instead of a second, blocking lookup.
self._connection = await get_async_connection(self.conn_id, hook=self)
factories: _ClientFactories[AsyncAnthropicClient] = _ClientFactories(
anthropic=AsyncAnthropic,
bedrock=AsyncAnthropicBedrock,
vertex=AsyncAnthropicVertex,
aws=AsyncAnthropicAWS,
foundry=AsyncAnthropicFoundry,
)
return self._build_client(factories)
@staticmethod
def _workload_identity_credentials(wif: dict[str, Any]) -> WorkloadIdentityCredentials:
"""
Build a WIF credential from the connection ``extra['workload_identity']`` mapping.
Exchanges a short-lived OIDC token (read from ``identity_token_file``) for an
Anthropic access token. See
https://platform.claude.com/docs/en/manage-claude/workload-identity-federation.
"""
required_fields = (
"identity_token_file",
"federation_rule_id",
"organization_id",
)
missing_fields = [field for field in required_fields if not wif.get(field)]
if missing_fields:
raise AnthropicError(
"The workload_identity configuration is missing or has empty required fields: "
f"{', '.join(missing_fields)}."
)
kwargs: dict[str, Any] = {
"identity_token_provider": IdentityTokenFile(wif["identity_token_file"]),
"federation_rule_id": wif["federation_rule_id"],
"organization_id": wif["organization_id"],
}
if wif.get("service_account_id"):
kwargs["service_account_id"] = wif["service_account_id"]
if wif.get("workspace_id"):
kwargs["workspace_id"] = wif["workspace_id"]
if wif.get("scope"):
kwargs["scope"] = wif["scope"]
return WorkloadIdentityCredentials(**kwargs)
def _resolve_model(self, model: str | None) -> str:
"""Resolve the effective model id; Bedrock rejects a bare id, so require its prefix."""
resolved = model or self.default_model
# Valid Bedrock ids either start with the ``anthropic.`` provider prefix or carry a
# region/profile prefix as a dotted component (e.g. ``us.anthropic.``, ``global.anthropic.``).
is_bedrock_model_id = resolved.startswith("anthropic.") or ".anthropic." in resolved
if self.platform == "bedrock" and not is_bedrock_model_id:
raise AnthropicError(
f"Model {resolved!r} is not a valid Amazon Bedrock model id. Bedrock ids carry a "
"provider/region prefix (e.g. 'global.anthropic.claude-opus-4-6-v1'); set one via "
"the 'model' argument or the connection's extra['model']."
)
return resolved
def _require_first_party(self, feature: str) -> None:
if self.platform not in FIRST_PARTY_PLATFORMS:
raise AnthropicError(
f"{feature} is not available on the {self.platform!r} platform. "
"Use the first-party Anthropic API (platform='anthropic') or "
"Claude Platform on AWS (platform='aws')."
)
@property
def _first_party_conn(self) -> Anthropic:
"""
Client cast to the first-party type for endpoints only it exposes.
Callers must guard with :meth:`_require_first_party` first; the Bedrock/Vertex/
Foundry clients don't expose ``beta.agents``/``beta.sessions``/``models``.
"""
return cast("Anthropic", self.conn)
[docs]
def test_connection(self) -> tuple[bool, str]:
"""Test the Anthropic connection."""
try:
if self.platform in FIRST_PARTY_PLATFORMS:
# Narrowed by the platform guard: only the first-party / AWS clients,
# which expose the Models API, reach this branch.
self._first_party_conn.models.list()
return True, "Connection established!"
# models.list() is not served on bedrock/vertex/foundry; building the
# client validates the configuration without a paid request.
self.get_conn()
return True, f"Connection configured for platform {self.platform!r} (no live check available)."
except Exception as e:
return False, str(e)
[docs]
def create_message(
self,
messages: list[dict[str, Any]],
model: str | None = None,
max_tokens: int = 1024,
system: str | None = None,
**kwargs: Any,
) -> Message:
"""
Create a single message response (one-shot ``messages.create``).
:param messages: The conversation so far, as a list of message dicts.
:param model: Model ID to use. Defaults to :attr:`default_model` (the connection's
``extra['model']`` or :data:`DEFAULT_MODEL`).
:param max_tokens: Maximum number of tokens to generate.
:param system: Optional system prompt.
"""
params: dict[str, Any] = {
"model": self._resolve_model(model),
"max_tokens": max_tokens,
"messages": messages,
**kwargs,
}
if system is not None:
params["system"] = system
return self.conn.messages.create(**params)
[docs]
def count_tokens(
self,
messages: list[dict[str, Any]],
model: str | None = None,
system: str | None = None,
**kwargs: Any,
) -> int:
"""Return the number of input tokens the given request would consume."""
self._require_first_party("Token counting")
params: dict[str, Any] = {"model": model or self.default_model, "messages": messages, **kwargs}
if system is not None:
params["system"] = system
return self.conn.messages.count_tokens(**params).input_tokens
@staticmethod
def _apply_default_model(request: dict[str, Any], default_model: str) -> dict[str, Any]:
"""
Fill ``params['model']`` from ``default_model`` when the request omits it.
The input dict is never mutated, and a request that sets its own ``model`` is
returned unchanged, so a single batch can still mix models across requests.
"""
params = request.get("params")
if not isinstance(params, dict) or params.get("model"):
return request
return {**request, "params": {**params, "model": default_model}}
[docs]
def create_batch(self, requests: list[dict[str, Any]], model: str | None = None) -> MessageBatch:
"""
Submit a Message Batch.
:param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where
``params`` is a ``messages.create`` payload (``model``, ``max_tokens``,
``messages``, ...). A request that omits ``model`` inherits ``model`` below,
or the connection's ``default_model`` (``extra['model']``) when that is unset too.
:param model: Default model id for requests that do not set their own. Falls back
to the connection's :attr:`default_model`.
"""
self._require_first_party("The Message Batches API")
default_model = model or self.default_model
prepared = [self._apply_default_model(request, default_model) for request in requests]
# ``Request`` is a TypedDict, so the plain dicts callers build match structurally.
return self.conn.messages.batches.create(requests=cast("Iterable[Request]", prepared))
[docs]
def get_batch(self, batch_id: str) -> MessageBatch:
"""Retrieve a Message Batch by ID."""
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.retrieve(batch_id)
[docs]
def cancel_batch(self, batch_id: str) -> MessageBatch:
"""Request cancellation of a Message Batch."""
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.cancel(batch_id)
[docs]
def list_batches(self, **kwargs: Any) -> Any:
"""Return a (paginated) list of Message Batches."""
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.list(**kwargs)
[docs]
def stream_batch_results(self, batch_id: str) -> Iterator[MessageBatchIndividualResponse]:
"""
Return a streaming iterator of per-request results, keyed by ``custom_id``.
Results stream from the API and arrive in **arbitrary order** — key them by
``result.custom_id``, never by position. Results are available for 29 days
after the batch is created. The result set can be very large: iterate and
persist to object storage; do not materialize it into XCom.
"""
# Return (don't ``yield``) so the platform guard fails fast at call time
# rather than only when the caller starts iterating.
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.results(batch_id)
[docs]
def wait_for_batch(
self, batch_id: str, wait_seconds: float = 3, timeout: float = 24 * 60 * 60
) -> MessageBatch:
"""
Poll a batch synchronously until it reaches the terminal ``ended`` status.
:param batch_id: The batch to wait for.
:param wait_seconds: Seconds to sleep between polls.
:param timeout: Maximum seconds to wait before raising :class:`AnthropicBatchTimeout`.
:return: The terminal :class:`~anthropic.types.messages.MessageBatch`.
"""
start = time.monotonic()
consecutive_failures = 0
while True:
try:
batch = self.get_batch(batch_id)
except Exception as e:
# Tolerate transient poll errors (as the deferrable trigger does) so a
# single blip does not fail — and cancel — a still-running batch whose
# results remain recoverable for 29 days.
consecutive_failures += 1
if (
consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES
or time.monotonic() - start > timeout
):
raise
self.log.warning("Polling batch %s failed (%s); retrying.", batch_id, e)
time.sleep(wait_seconds)
continue
consecutive_failures = 0
self.log.debug("Batch %s status=%s", batch_id, batch.processing_status)
if not BatchStatus.is_in_progress(batch.processing_status):
return batch
if time.monotonic() - start > timeout:
raise AnthropicBatchTimeout(
f"Batch {batch_id} did not reach a terminal status within {timeout} seconds."
)
time.sleep(wait_seconds)
# --- Managed Agents -------------------------------------------------------
# Agents and environments are persisted, reusable resources: create them once
# (these helpers, the ``ant`` CLI, or a setup script) and store the IDs. The
# operator references those IDs; it never creates an agent per run.
[docs]
def create_agent(
self, name: str, model: str | dict[str, Any] | None = None, **kwargs: Any
) -> BetaManagedAgentsAgent:
"""
Create a (reusable, versioned) Managed Agents agent. One-time setup.
``model`` defaults to :attr:`default_model` (the connection's ``extra['model']``
or :data:`DEFAULT_MODEL`). Pass a mapping instead of a bare id to set the model
config, e.g. ``{"id": "claude-opus-5", "inference_geo": "us"}``.
"""
self._require_first_party("Managed Agents")
agent = self._first_party_conn.beta.agents.create(
name=name,
model=cast("agent_create_params.Model", model or self.default_model),
**kwargs,
)
self.log.debug("Created agent %s (name=%r, model=%s)", agent.id, name, model or self.default_model)
return agent
[docs]
def create_environment(
self, name: str, config: dict[str, Any] | None = None, **kwargs: Any
) -> BetaEnvironment:
"""Create a (reusable) environment for agent sessions. One-time setup."""
self._require_first_party("Managed Agents")
if config is None:
config = {"type": "cloud", "networking": {"type": "unrestricted"}}
environment = self._first_party_conn.beta.environments.create(
name=name, config=cast("environment_create_params.Config", config), **kwargs
)
self.log.debug("Created environment %s (name=%r)", environment.id, name)
return environment
[docs]
def create_session(self, agent: str, environment_id: str, **kwargs: Any) -> BetaManagedAgentsSession:
"""
Start a session against a pre-created agent + environment.
A ``budget`` keyword accepts an amount in USD as well as the raw API payload; see
:func:`build_budget`.
"""
self._require_first_party("Managed Agents")
if "budget" in kwargs:
if kwargs["budget"] is None:
# ``SessionCreateParams.budget`` is not Optional -- unlike the update
# params, create has no "clear the ceiling" semantics, and sending
# ``"budget": null`` is rejected. Treat None as "no budget".
kwargs.pop("budget")
else:
kwargs["budget"] = build_budget(kwargs["budget"])
return self._first_party_conn.beta.sessions.create(
agent=agent, environment_id=environment_id, **kwargs
)
[docs]
def update_session(self, session_id: str, **kwargs: Any) -> BetaManagedAgentsSession:
"""
Update a live session -- raise or clear its ``budget``, swap ``agent`` tools, retitle.
Only the keywords you pass are sent, which matters because the API distinguishes
*omitted* (preserve) from ``None`` (clear). So ``update_session(sid)`` changes
nothing, while ``update_session(sid, budget=None)`` removes the ceiling -- the
escape hatch for a session stopped by a model with no list price, which raising the
budget cannot unblock.
``budget`` accepts an amount in USD or the raw payload (see :func:`build_budget`).
``agent={"tools": [...]}`` is a **full replacement** of the tool list, not a merge,
and needs the ``mid-conversation-tool-changes-2026-07-01`` beta.
"""
self._require_first_party("Managed Agents")
if kwargs.get("budget") is not None:
kwargs["budget"] = build_budget(kwargs["budget"])
return self._first_party_conn.beta.sessions.update(session_id, **kwargs)
[docs]
def get_session(self, session_id: str) -> BetaManagedAgentsSession:
"""Retrieve a session (carries its current ``status``)."""
self._require_first_party("Managed Agents")
return self._first_party_conn.beta.sessions.retrieve(session_id)
[docs]
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
[docs]
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")
[docs]
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")
# Event dicts callers build match the SDK's TypedDict union structurally.
return self._first_party_conn.beta.sessions.events.send(
session_id, events=cast("list[BetaManagedAgentsEventParams]", [event])
)
[docs]
def interrupt_session(self, session_id: str) -> Any:
"""
Send ``user.interrupt`` to pause a running session.
The API refuses to archive or delete a session while it is ``running``, so this is
the only way to release one that is not going to stop on its own -- see
:meth:`archive_session`.
"""
self._require_first_party("Managed Agents")
return self.send_event(session_id, {"type": "user.interrupt"})
[docs]
def archive_session(
self, session_id: str, *, attempts: int = 6, wait_seconds: float = 5
) -> 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.
A ``running`` session cannot be archived (nor deleted): the API rejects both with a
400. Only then does this interrupt the session and retry, because a session that
will not stop on its own otherwise accrues billable runtime with no way to release
it. Any other failure is re-raised untouched, so a transient 5xx does not send
``user.interrupt`` to a session that was working fine.
Retrying costs up to ``attempts`` further calls with ``wait_seconds`` between them
(about 25s at the defaults), which is longer than some callers have: a killed task's
``on_kill`` is SIGKILLed a few seconds in, so it passes a much tighter budget.
"""
self._require_first_party("Managed Agents")
try:
return self._first_party_conn.beta.sessions.archive(session_id)
except BadRequestError as e:
# Catching the SDK's published error type, not matching on message text: a 400
# here is the documented "cannot archive while running" rejection.
self.log.info("Archiving session %s failed (%s); interrupting and retrying.", session_id, e)
self.interrupt_session(session_id)
return self._wait_for_archive(session_id, attempts=attempts, wait_seconds=wait_seconds)
def _wait_for_archive(
self, session_id: str, attempts: int = 6, wait_seconds: float = 5
) -> BetaManagedAgentsSession:
"""Retry archiving while the interrupt takes effect; the status change is not instant."""
for attempt in range(attempts):
try:
return self._first_party_conn.beta.sessions.archive(session_id)
except Exception:
if attempt == attempts - 1:
raise
time.sleep(wait_seconds)
raise AnthropicError(f"Could not archive session {session_id}.") # pragma: no cover
def _latest_idle_reason(self, session_id: str, kickoff_event_id: str | None) -> str | None:
"""
Return the ``stop_reason`` of the newest ``session.status_idle`` event, or ``None``.
Walks the event log newest-first. Returns ``None`` if the kickoff event is the most
recent event (the agent has not responded yet — defeats the start race) or no idle
event is found in the scan window.
"""
# The SDK cursor auto-paginates (page size 20); cap the walk at 100 events so a
# long event log can't make one poll iterate unboundedly.
examined = 0
for event in self._first_party_conn.beta.sessions.events.list(session_id, order="desc", limit=20):
if kickoff_event_id is not None and event.id == kickoff_event_id:
return None
if event.type == "session.status_idle":
return event.stop_reason.type
examined += 1
if examined >= 100:
break
return None
[docs]
def poll_session_completion(
self, session_id: str, *, expect_outcome: bool = False, kickoff_event_id: str | None = None
) -> SessionPollResult:
"""
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`` /
``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(
session, expect_outcome=expect_outcome
)
self.log.debug(
"Session %s status=%s done=%s needs_event_check=%s",
session_id,
session.status,
done,
needs_event_check,
)
if not needs_event_check:
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 SessionPollResult(done=False, error_message=None, stop_reason=None)
if reason == "end_turn":
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,
)
[docs]
def wait_for_session(
self,
session_id: str,
expect_outcome: bool = False,
kickoff_event_id: str | None = None,
poll_interval: float = 30,
timeout: float = 24 * 60 * 60,
) -> None:
"""
Poll a session synchronously until it completes.
:param session_id: The session to wait for.
:param expect_outcome: Whether the session is running a ``user.define_outcome`` loop
(completion judged from ``outcome_evaluations``).
:param kickoff_event_id: ID of the kickoff event, used to correlate the terminal
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:
poll_result = self.poll_session_completion(
session_id, expect_outcome=expect_outcome, kickoff_event_id=kickoff_event_id
)
except Exception as e:
# Tolerate transient poll errors (as the deferrable trigger does) so a
# single blip does not fail — and archive — a still-running session.
consecutive_failures += 1
if (
consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES
or time.monotonic() - start > timeout
):
raise
self.log.warning("Polling session %s failed (%s); retrying.", session_id, e)
time.sleep(poll_interval)
continue
consecutive_failures = 0
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(
f"Session {session_id} did not reach a terminal status within {timeout} seconds."
)
time.sleep(poll_interval)