From f202eef237847862097850bb4387452df762bce6 Mon Sep 17 00:00:00 2001
From: "stainless-app[bot]"
 <142633134+stainless-app[bot]@users.noreply.github.com>
Date: Thu, 19 Sep 2024 16:08:39 +0000
Subject: [PATCH] feat(client): send retry count header (#493)

---
 src/core.ts         | 16 ++++++++++++----
 tests/index.test.ts | 30 ++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/core.ts b/src/core.ts
index 70dc1b46..7d228834 100644
--- a/src/core.ts
+++ b/src/core.ts
@@ -262,7 +262,10 @@ export abstract class APIClient {
     return null;
   }
 
-  buildRequest<Req>(options: FinalRequestOptions<Req>): { req: RequestInit; url: string; timeout: number } {
+  buildRequest<Req>(
+    options: FinalRequestOptions<Req>,
+    { retryCount = 0 }: { retryCount?: number } = {},
+  ): { req: RequestInit; url: string; timeout: number } {
     const { method, path, query, headers: headers = {} } = options;
 
     const body =
@@ -292,7 +295,7 @@ export abstract class APIClient {
       headers[this.idempotencyHeader] = options.idempotencyKey;
     }
 
-    const reqHeaders = this.buildHeaders({ options, headers, contentLength });
+    const reqHeaders = this.buildHeaders({ options, headers, contentLength, retryCount });
 
     const req: RequestInit = {
       method,
@@ -311,10 +314,12 @@ export abstract class APIClient {
     options,
     headers,
     contentLength,
+    retryCount,
   }: {
     options: FinalRequestOptions;
     headers: Record<string, string | null | undefined>;
     contentLength: string | null | undefined;
+    retryCount: number;
   }): Record<string, string> {
     const reqHeaders: Record<string, string> = {};
     if (contentLength) {
@@ -330,6 +335,8 @@ export abstract class APIClient {
       delete reqHeaders['content-type'];
     }
 
+    reqHeaders['x-stainless-retry-count'] = String(retryCount);
+
     this.validateHeaders(reqHeaders, headers);
 
     return reqHeaders;
@@ -381,13 +388,14 @@ export abstract class APIClient {
     retriesRemaining: number | null,
   ): Promise<APIResponseProps> {
     const options = await optionsInput;
+    const maxRetries = options.maxRetries ?? this.maxRetries;
     if (retriesRemaining == null) {
-      retriesRemaining = options.maxRetries ?? this.maxRetries;
+      retriesRemaining = maxRetries;
     }
 
     await this.prepareOptions(options);
 
-    const { req, url, timeout } = this.buildRequest(options);
+    const { req, url, timeout } = this.buildRequest(options, { retryCount: maxRetries - retriesRemaining });
 
     await this.prepareRequest(req, { url, options });
 
diff --git a/tests/index.test.ts b/tests/index.test.ts
index e802c5cf..85d6e777 100644
--- a/tests/index.test.ts
+++ b/tests/index.test.ts
@@ -268,6 +268,36 @@ describe('retries', () => {
     expect(count).toEqual(3);
   });
 
+  test('retry count header', async () => {
+    let count = 0;
+    let capturedRequest: RequestInit | undefined;
+    const testFetch = async (url: RequestInfo, init: RequestInit = {}): Promise<Response> => {
+      count++;
+      if (count <= 2) {
+        return new Response(undefined, {
+          status: 429,
+          headers: {
+            'Retry-After': '0.1',
+          },
+        });
+      }
+      capturedRequest = init;
+      return new Response(JSON.stringify({ a: 1 }), { headers: { 'Content-Type': 'application/json' } });
+    };
+
+    const client = new Mux({
+      tokenId: 'my token id',
+      tokenSecret: 'my secret',
+      fetch: testFetch,
+      maxRetries: 4,
+    });
+
+    expect(await client.request({ path: '/foo', method: 'get' })).toEqual({ a: 1 });
+
+    expect((capturedRequest!.headers as Headers)['x-stainless-retry-count']).toEqual('2');
+    expect(count).toEqual(3);
+  });
+
   test('retry on 429 with retry-after', async () => {
     let count = 0;
     const testFetch = async (url: RequestInfo, { signal }: RequestInit = {}): Promise<Response> => {