This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch anthropic-interrupt-before-archive in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 56057b65198d2de94cd002d12ae504e95a65a42a Author: Kaxil Naik <[email protected]> AuthorDate: Wed Aug 12 01:13:35 2026 +0530 Interrupt a running Anthropic session before archiving it The API refuses to archive **or** delete a session while its status is ``running``, rejecting both with a 400. ``AnthropicAgentSessionOperator`` treats archiving as best-effort teardown on every failure path, so a session that will not stop on its own is left behind with no way to release it -- and it keeps accruing ``active_seconds``, which the SDK describes as the duration the session's runtime cost is priced on. Observed directly: two sessions halted by a $0.01 budget sat at status ``running`` for over 40 minutes with model spend frozen at the ceiling while ``active_seconds`` kept climbing. ``archive`` and ``delete`` both returned 400. Sending ``user.interrupt`` moved them to ``idle`` immediately, after which archiving succeeded. ``AnthropicHook.archive_session`` now interrupts and retries when the first archive attempt fails, and ``interrupt_session`` exposes the event on its own. The retry is bounded, so an unarchivable session still surfaces its error rather than looping. This is reachable well beyond budgets -- any session still working when a task fails, times out, or is killed hits the same 400 -- but a budget halt makes it routine, because that is a session which has stopped spending without stopping. The interrupt-then-archive sequence was verified by hand against the live API on two genuinely stuck sessions. Its wiring through ``archive_session`` is covered by unit tests rather than a live run. --- .../airflow/providers/anthropic/hooks/anthropic.py | 48 +++++++++++++++++++++- .../airflow/providers/anthropic/operators/agent.py | 17 ++++++-- .../tests/unit/anthropic/hooks/test_anthropic.py | 38 +++++++++++++++++ .../tests/unit/anthropic/operators/test_agent.py | 2 +- 4 files changed, 98 insertions(+), 7 deletions(-) diff --git a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py index 0096b95ff12..864a6bbb813 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py +++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py @@ -31,6 +31,7 @@ from anthropic import ( AnthropicBedrock, AnthropicFoundry, AnthropicVertex, + BadRequestError, IdentityTokenFile, WorkloadIdentityCredentials, ) @@ -708,15 +709,58 @@ class AnthropicHook(BaseHook): session_id, events=cast("list[BetaManagedAgentsEventParams]", [event]) ) - def archive_session(self, session_id: str) -> BetaManagedAgentsSession: + 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"}) + + 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") - return self._first_party_conn.beta.sessions.archive(session_id) + 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: """ diff --git a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py index afdd916d7ce..cf367d2197c 100644 --- a/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py +++ b/providers/anthropic/src/airflow/providers/anthropic/operators/agent.py @@ -302,17 +302,20 @@ class AnthropicAgentSessionOperator(BaseOperator): except Exception: self.log.exception("Could not record usage for session %s", session_id) - def _archive_session(self, session_id: str | None) -> BetaManagedAgentsSession | None: + def _archive_session( + self, session_id: str | None, **archive_kwargs: Any + ) -> BetaManagedAgentsSession | None: """ 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. + archive call failed. ``archive_kwargs`` tightens the hook's interrupt-and-retry + budget for callers that do not have its full ~25s. """ if not session_id: return None try: - return self.hook.archive_session(session_id) + return self.hook.archive_session(session_id, **archive_kwargs) except Exception as e: self.log.warning("Failed to archive session %s: %s", session_id, e) return None @@ -325,5 +328,11 @@ class AnthropicAgentSessionOperator(BaseOperator): (``deferrable=False``). On Airflow 3.3+ a killed deferred task is archived by the trigger's ``on_kill``. On older Airflow the session of a killed deferred task is not archived automatically; archive it manually via the hook. + + The supervisor escalates to SIGKILL a few seconds after asking the task to stop, so + the hook's default retry budget would never finish here -- the process dies inside + the first sleep and the session is never released. This budget gives the interrupt a + beat to land (the status change is not instant) while still fitting inside that + window: two attempts, one second apart. """ - self._archive_session(self.session_id) + self._archive_session(self.session_id, attempts=2, wait_seconds=1) diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py index dd0caaf96d9..816b238d401 100644 --- a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py +++ b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py @@ -45,6 +45,8 @@ from airflow.providers.anthropic.hooks.anthropic import ( pytest.importorskip("anthropic") +import httpx +from anthropic import BadRequestError from anthropic.types import BetaMonetaryAmount from anthropic.types.beta import BetaManagedAgentsServerToolUsage, BetaManagedAgentsSessionUsage from anthropic.types.beta.beta_managed_agents_cache_creation_usage import ( @@ -279,6 +281,42 @@ class TestBuildBudget: build_budget(amount) +def _make_bad_request(message="cannot be archived while its status is running") -> BadRequestError: + """A real SDK BadRequestError -- archive_session only interrupts on this, not on any error.""" + request = httpx.Request("POST", "https://api.anthropic.com/v1/sessions/sess_1/archive";) + return BadRequestError(message, response=httpx.Response(400, request=request), body=None) + + +class TestArchiveSession: + def test_archives_directly_when_the_session_is_stoppable(self): + hook, client = _make_hook() + hook.archive_session("sess_1") + client.beta.sessions.archive.assert_called_once_with("sess_1") + client.beta.sessions.events.send.assert_not_called() + + @mock.patch(f"{HOOK_PATH}.time.sleep", autospec=True) + def test_interrupts_and_retries_when_the_session_is_running(self, mock_sleep): + # A running session is refused by both archive and delete, so without the interrupt + # it keeps accruing billable runtime with no way to release it. + hook, client = _make_hook() + archived = object() + client.beta.sessions.archive.side_effect = [_make_bad_request(), archived] + assert hook.archive_session("sess_1") is archived + client.beta.sessions.events.send.assert_called_once_with( + "sess_1", events=[{"type": "user.interrupt"}] + ) + assert client.beta.sessions.archive.call_count == 2 + + @mock.patch(f"{HOOK_PATH}.time.sleep", autospec=True) + def test_gives_up_after_the_retry_budget(self, mock_sleep): + hook, client = _make_hook() + client.beta.sessions.archive.side_effect = _make_bad_request() + with pytest.raises(BadRequestError): + hook.archive_session("sess_1") + # one initial attempt plus the bounded retry loop + assert client.beta.sessions.archive.call_count == 7 + + class TestUpdateSession: def test_only_passes_supplied_keys(self): # The API distinguishes omitted (preserve) from None (clear), so an unmentioned diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py index d7a05e8217c..010e20d393e 100644 --- a/providers/anthropic/tests/unit/anthropic/operators/test_agent.py +++ b/providers/anthropic/tests/unit/anthropic/operators/test_agent.py @@ -473,7 +473,7 @@ class TestOnKill: op = AnthropicAgentSessionOperator(task_id="a", agent_id="ag", environment_id="env", message="hi") op.session_id = "sess_1" op.on_kill() - hook.archive_session.assert_called_once_with("sess_1") + hook.archive_session.assert_called_once_with("sess_1", attempts=2, wait_seconds=1) @mock.patch.object(AnthropicAgentSessionOperator, "hook", new_callable=mock.PropertyMock) def test_on_kill_noop_without_session(self, mock_hook_prop):
