diff --git a/src/anthropic/lib/environments/_poller.py b/src/anthropic/lib/environments/_poller.py index e1f27d788..d40515885 100644 --- a/src/anthropic/lib/environments/_poller.py +++ b/src/anthropic/lib/environments/_poller.py @@ -137,25 +137,37 @@ def iter_work( time.sleep(_jitter(1.0, 3.0)) continue log.info("claimed work work_id=%s work_type=%s", item.id, getattr(item.data, "type", None)) - try: - work.ack( - item.id, - environment_id=environment_id, - extra_headers=extra_headers, - ) - except TRANSIENT_ERRORS as e: - if _is_fatal_4xx(e): - log.error("ack failed permanently; force-stopping work_id=%s error=%s", item.id, e) - _force_stop_quietly(work, item.id, environment_id=environment_id, extra_headers=extra_headers) + + acked = False + while not acked: + try: + work.ack( + item.id, + environment_id=environment_id, + extra_headers=extra_headers, + ) + except TRANSIENT_ERRORS as e: + if _is_fatal_4xx(e): + log.error("ack failed permanently; force-stopping work_id=%s error=%s", item.id, e) + _force_stop_quietly(work, item.id, environment_id=environment_id, extra_headers=extra_headers) + ack_attempt = 0 + break + ack_attempt += 1 + wait = _backoff(ack_attempt) + _jitter(0.0, 1.0) + log.warning( + "ack failed, backing off work_id=%s attempt=%d backoff=%.1fs error=%s", + item.id, + ack_attempt, + wait, + e, + ) + time.sleep(wait) continue - ack_attempt += 1 - wait = _backoff(ack_attempt) + _jitter(0.0, 1.0) - log.warning( - "ack failed, backing off work_id=%s attempt=%d backoff=%.1fs error=%s", item.id, ack_attempt, wait, e - ) - time.sleep(wait) + ack_attempt = 0 + acked = True + + if not acked: continue - ack_attempt = 0 if not auto_stop: yield item continue @@ -230,25 +242,42 @@ async def aiter_work( await anyio.sleep(_jitter(1.0, 3.0)) continue log.info("claimed work work_id=%s work_type=%s", item.id, getattr(item.data, "type", None)) - try: - await work.ack( - item.id, - environment_id=environment_id, - extra_headers=extra_headers, - ) - except TRANSIENT_ERRORS as e: - if _is_fatal_4xx(e): - log.error("ack failed permanently; force-stopping work_id=%s error=%s", item.id, e) - await _aforce_stop_quietly(work, item.id, environment_id=environment_id, extra_headers=extra_headers) + + acked = False + while not acked: + try: + await work.ack( + item.id, + environment_id=environment_id, + extra_headers=extra_headers, + ) + except TRANSIENT_ERRORS as e: + if _is_fatal_4xx(e): + log.error("ack failed permanently; force-stopping work_id=%s error=%s", item.id, e) + await _aforce_stop_quietly( + work, + item.id, + environment_id=environment_id, + extra_headers=extra_headers, + ) + ack_attempt = 0 + break + ack_attempt += 1 + wait = _backoff(ack_attempt) + _jitter(0.0, 1.0) + log.warning( + "ack failed, backing off work_id=%s attempt=%d backoff=%.1fs error=%s", + item.id, + ack_attempt, + wait, + e, + ) + await anyio.sleep(wait) continue - ack_attempt += 1 - wait = _backoff(ack_attempt) + _jitter(0.0, 1.0) - log.warning( - "ack failed, backing off work_id=%s attempt=%d backoff=%.1fs error=%s", item.id, ack_attempt, wait, e - ) - await anyio.sleep(wait) + ack_attempt = 0 + acked = True + + if not acked: continue - ack_attempt = 0 if not auto_stop: yield item continue diff --git a/tests/lib/environments/test_poller_ack_retry.py b/tests/lib/environments/test_poller_ack_retry.py new file mode 100644 index 000000000..02716149c --- /dev/null +++ b/tests/lib/environments/test_poller_ack_retry.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from typing import Any, cast + +import httpx +import pytest + +from anthropic import APIStatusError +from anthropic.lib.environments import _poller +from anthropic.lib.environments._poller import aiter_work, iter_work + + +class _WorkData: + type = "session" + + +class _WorkItem: + def __init__(self, work_id: str = "work_1") -> None: + self.id = work_id + self.data = _WorkData() + + +def _api_status_error(code: int) -> APIStatusError: + request = httpx.Request("POST", "https://api.example/work") + response = httpx.Response(code, request=request, content=b"{}") + return APIStatusError("boom", response=response, body=None) + + +class _SyncWork: + def __init__(self) -> None: + self.item = _WorkItem() + self.events: list[str] = [] + self.poll_calls = 0 + self.ack_calls = 0 + + def poll(self, _environment_id: str, **_kwargs: Any) -> _WorkItem: + self.poll_calls += 1 + self.events.append("poll") + if self.poll_calls > 1: + raise AssertionError("polled again before the claimed item was acknowledged") + return self.item + + def ack(self, work_id: str, **_kwargs: Any) -> None: + assert work_id == self.item.id + self.ack_calls += 1 + self.events.append("ack") + if self.ack_calls == 1: + raise _api_status_error(500) + + def stop(self, _work_id: str, **_kwargs: Any) -> None: + raise AssertionError("auto_stop=False should not stop the item") + + +class _AsyncWork: + def __init__(self) -> None: + self.item = _WorkItem() + self.events: list[str] = [] + self.poll_calls = 0 + self.ack_calls = 0 + + async def poll(self, _environment_id: str, **_kwargs: Any) -> _WorkItem: + self.poll_calls += 1 + self.events.append("poll") + if self.poll_calls > 1: + raise AssertionError("polled again before the claimed item was acknowledged") + return self.item + + async def ack(self, work_id: str, **_kwargs: Any) -> None: + assert work_id == self.item.id + self.ack_calls += 1 + self.events.append("ack") + if self.ack_calls == 1: + raise _api_status_error(500) + + async def stop(self, _work_id: str, **_kwargs: Any) -> None: + raise AssertionError("auto_stop=False should not stop the item") + + +def test_sync_transient_ack_failure_retries_same_claim_before_polling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(_poller.time, "sleep", lambda _seconds: None) + work = _SyncWork() + + item = next(iter_work(cast(Any, work), environment_id="env_1", auto_stop=False)) + + assert item is work.item + assert work.events == ["poll", "ack", "ack"] + assert work.poll_calls == 1 + + +async def test_async_transient_ack_failure_retries_same_claim_before_polling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def _no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(_poller.anyio, "sleep", _no_sleep) + work = _AsyncWork() + iterator = aiter_work(cast(Any, work), environment_id="env_1", auto_stop=False) + + item = await iterator.__anext__() + + assert item is work.item + assert work.events == ["poll", "ack", "ack"] + assert work.poll_calls == 1