feat(client): send retry count header (#314)

Author: stainless-app[bot]

lithic-java-core/src/main/kotlin/com/lithic/api/core/http/RetryingHttpClient.kt

@@ -40,9 +40,16 @@ private constructor(
 
         maybeAddIdempotencyHeader(request)
 
+        // Don't send the current retry count in the headers if the caller set their own value.
+        val shouldSendRetryCount = !request.headers.containsKey("x-stainless-retry-count")
+
         var retries = 0
 
         while (true) {
+            if (shouldSendRetryCount) {
+                setRetryCountHeader(request, retries)
+            }
+
             val response =
                 try {
                     val response = httpClient.execute(request, requestOptions)
@@ -74,10 +81,21 @@ private constructor(
 
         maybeAddIdempotencyHeader(request)
 
+        // Don't send the current retry count in the headers if the caller set their own value.
+        val shouldSendRetryCount = !request.headers.containsKey("x-stainless-retry-count")
+
         var retries = 0
 
-        fun wrap(future: CompletableFuture<HttpResponse>): CompletableFuture<HttpResponse> {
-            return future
+        fun executeWithRetries(
+            request: HttpRequest,
+            requestOptions: RequestOptions,
+        ): CompletableFuture<HttpResponse> {
+            if (shouldSendRetryCount) {
+                setRetryCountHeader(request, retries)
+            }
+
+            return httpClient
+                .executeAsync(request, requestOptions)
                 .handleAsync(
                     fun(
                         response: HttpResponse?,
@@ -97,15 +115,15 @@ private constructor(
 
                         val backoffMillis = getRetryBackoffMillis(retries, response)
                         return sleepAsync(backoffMillis.toMillis()).thenCompose {
-                            wrap(httpClient.executeAsync(request, requestOptions))
+                            executeWithRetries(request, requestOptions)
                         }
                     },
                     MoreExecutors.directExecutor()
                 )
                 .thenCompose(Function.identity())
         }
 
-        return wrap(httpClient.executeAsync(request, requestOptions))
+        return executeWithRetries(request, requestOptions)
     }
 
     override fun close() {
@@ -118,6 +136,11 @@ private constructor(
         return request.body?.repeatable() ?: true
     }
 
+    private fun setRetryCountHeader(request: HttpRequest, retries: Int) {
+        request.headers.removeAll("x-stainless-retry-count")
+        request.headers.put("x-stainless-retry-count", retries.toString())
+    }
+
     private fun idempotencyKey(): String = "stainless-java-retry-${UUID.randomUUID()}"
 
     private fun maybeAddIdempotencyHeader(request: HttpRequest) {

lithic-java-core/src/test/kotlin/com/lithic/api/core/http/RetryingHttpClientTest.kt

@@ -8,6 +8,8 @@ import com.lithic.api.client.okhttp.OkHttpClient
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.ValueSource
 
 @WireMockTest
 internal class RetryingHttpClientTest {
@@ -50,8 +52,9 @@ internal class RetryingHttpClientTest {
         verify(1, postRequestedFor(urlPathEqualTo("/something")))
     }
 
-    @Test
-    fun retryAfterHeader() {
+    @ParameterizedTest
+    @ValueSource(booleans = [false, true])
+    fun retryAfterHeader(async: Boolean) {
         val request =
             HttpRequest.builder().method(HttpMethod.POST).addPathSegment("something").build()
         stubFor(
@@ -79,9 +82,67 @@ internal class RetryingHttpClientTest {
         )
         val retryingClient =
             RetryingHttpClient.builder().httpClient(httpClient).maxRetries(2).build()
-        val response = retryingClient.execute(request)
+
+        val response =
+            if (async) retryingClient.executeAsync(request).get()
+            else retryingClient.execute(request)
+
         assertThat(response.statusCode()).isEqualTo(200)
-        verify(3, postRequestedFor(urlPathEqualTo("/something")))
+        verify(
+            1,
+            postRequestedFor(urlPathEqualTo("/something"))
+                .withHeader("x-stainless-retry-count", equalTo("0"))
+        )
+        verify(
+            1,
+            postRequestedFor(urlPathEqualTo("/something"))
+                .withHeader("x-stainless-retry-count", equalTo("1"))
+        )
+        verify(
+            1,
+            postRequestedFor(urlPathEqualTo("/something"))
+                .withHeader("x-stainless-retry-count", equalTo("2"))
+        )
+    }
+
+    @ParameterizedTest
+    @ValueSource(booleans = [false, true])
+    fun overwriteRetryCountHeader(async: Boolean) {
+        val request =
+            HttpRequest.builder()
+                .method(HttpMethod.POST)
+                .addPathSegment("something")
+                .putHeader("x-stainless-retry-count", "42")
+                .build()
+        stubFor(
+            post(urlPathEqualTo("/something"))
+                .inScenario("foo") // first we fail with a retry after header given as a date
+                .whenScenarioStateIs(Scenario.STARTED)
+                .willReturn(
+                    serviceUnavailable().withHeader("Retry-After", "Wed, 21 Oct 2015 07:28:00 GMT")
+                )
+                .willSetStateTo("RETRY_AFTER_DATE")
+        )
+        stubFor(
+            post(urlPathEqualTo("/something"))
+                .inScenario("foo") // then we return a success
+                .whenScenarioStateIs("RETRY_AFTER_DATE")
+                .willReturn(ok())
+                .willSetStateTo("COMPLETED")
+        )
+        val retryingClient =
+            RetryingHttpClient.builder().httpClient(httpClient).maxRetries(2).build()
+
+        val response =
+            if (async) retryingClient.executeAsync(request).get()
+            else retryingClient.execute(request)
+
+        assertThat(response.statusCode()).isEqualTo(200)
+        verify(
+            2,
+            postRequestedFor(urlPathEqualTo("/something"))
+                .withHeader("x-stainless-retry-count", equalTo("42"))
+        )
     }
 
     @Test