diff --git a/src/anthropic/_base_client.py b/src/anthropic/_base_client.py index 16450584..4cba75e9 100644 --- a/src/anthropic/_base_client.py +++ b/src/anthropic/_base_client.py @@ -413,7 +413,10 @@ def _build_headers(self, options: FinalRequestOptions, *, retries_taken: int = 0 if idempotency_header and options.method.lower() != "get" and idempotency_header not in headers: headers[idempotency_header] = options.idempotency_key or self._idempotency_key() - headers.setdefault("x-stainless-retry-count", str(retries_taken)) + # Don't set the retry count header if it was already set or removed by the caller. We check + # `custom_headers`, which can contain `Omit()`, instead of `headers` to account for the removal case. + if "x-stainless-retry-count" not in (header.lower() for header in custom_headers): + headers["x-stainless-retry-count"] = str(retries_taken) return headers diff --git a/tests/test_client.py b/tests/test_client.py index 35a07afe..73d296a7 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -900,6 +900,72 @@ def retry_handler(_request: httpx.Request) -> httpx.Response: assert response.retries_taken == failures_before_success assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) + @mock.patch("anthropic._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + @pytest.mark.respx(base_url=base_url) + def test_omit_retry_count_header( + self, client: Anthropic, failures_before_success: int, respx_mock: MockRouter + ) -> None: + client = client.with_options(max_retries=4) + + nb_retries = 0 + + def retry_handler(_request: httpx.Request) -> httpx.Response: + nonlocal nb_retries + if nb_retries < failures_before_success: + nb_retries += 1 + return httpx.Response(500) + return httpx.Response(200) + + respx_mock.post("/v1/messages").mock(side_effect=retry_handler) + + response = client.messages.with_raw_response.create( + max_tokens=1024, + messages=[ + { + "content": "Hello, world", + "role": "user", + } + ], + model="claude-3-5-sonnet-20240620", + extra_headers={"x-stainless-retry-count": Omit()}, + ) + + assert len(response.http_request.headers.get_list("x-stainless-retry-count")) == 0 + + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) + @mock.patch("anthropic._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + @pytest.mark.respx(base_url=base_url) + def test_overwrite_retry_count_header( + self, client: Anthropic, failures_before_success: int, respx_mock: MockRouter + ) -> None: + client = client.with_options(max_retries=4) + + nb_retries = 0 + + def retry_handler(_request: httpx.Request) -> httpx.Response: + nonlocal nb_retries + if nb_retries < failures_before_success: + nb_retries += 1 + return httpx.Response(500) + return httpx.Response(200) + + respx_mock.post("/v1/messages").mock(side_effect=retry_handler) + + response = client.messages.with_raw_response.create( + max_tokens=1024, + messages=[ + { + "content": "Hello, world", + "role": "user", + } + ], + model="claude-3-5-sonnet-20240620", + extra_headers={"x-stainless-retry-count": "42"}, + ) + + assert response.http_request.headers.get("x-stainless-retry-count") == "42" + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) @mock.patch("anthropic._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url) @@ -1802,6 +1868,74 @@ def retry_handler(_request: httpx.Request) -> httpx.Response: assert response.retries_taken == failures_before_success assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) + @mock.patch("anthropic._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + @pytest.mark.respx(base_url=base_url) + @pytest.mark.asyncio + async def test_omit_retry_count_header( + self, async_client: AsyncAnthropic, failures_before_success: int, respx_mock: MockRouter + ) -> None: + client = async_client.with_options(max_retries=4) + + nb_retries = 0 + + def retry_handler(_request: httpx.Request) -> httpx.Response: + nonlocal nb_retries + if nb_retries < failures_before_success: + nb_retries += 1 + return httpx.Response(500) + return httpx.Response(200) + + respx_mock.post("/v1/messages").mock(side_effect=retry_handler) + + response = await client.messages.with_raw_response.create( + max_tokens=1024, + messages=[ + { + "content": "Hello, world", + "role": "user", + } + ], + model="claude-3-5-sonnet-20240620", + extra_headers={"x-stainless-retry-count": Omit()}, + ) + + assert len(response.http_request.headers.get_list("x-stainless-retry-count")) == 0 + + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) + @mock.patch("anthropic._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) + @pytest.mark.respx(base_url=base_url) + @pytest.mark.asyncio + async def test_overwrite_retry_count_header( + self, async_client: AsyncAnthropic, failures_before_success: int, respx_mock: MockRouter + ) -> None: + client = async_client.with_options(max_retries=4) + + nb_retries = 0 + + def retry_handler(_request: httpx.Request) -> httpx.Response: + nonlocal nb_retries + if nb_retries < failures_before_success: + nb_retries += 1 + return httpx.Response(500) + return httpx.Response(200) + + respx_mock.post("/v1/messages").mock(side_effect=retry_handler) + + response = await client.messages.with_raw_response.create( + max_tokens=1024, + messages=[ + { + "content": "Hello, world", + "role": "user", + } + ], + model="claude-3-5-sonnet-20240620", + extra_headers={"x-stainless-retry-count": "42"}, + ) + + assert response.http_request.headers.get("x-stainless-retry-count") == "42" + @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) @mock.patch("anthropic._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) @pytest.mark.respx(base_url=base_url)