From 7b9723d66f459a5ac945b605752561b2f1518dd8 Mon Sep 17 00:00:00 2001
From: Stainless Bot <dev@stainlessapi.com>
Date: Thu, 19 Sep 2024 16:11:24 +0000
Subject: [PATCH] feat(client): send retry count header

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

diff --git a/src/core.ts b/src/core.ts
index a4bb87a32..faedeb1da 100644
--- a/src/core.ts
+++ b/src/core.ts
@@ -288,7 +288,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 =
@@ -320,7 +323,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,
@@ -339,10 +342,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) {
@@ -358,6 +363,8 @@ export abstract class APIClient {
       delete reqHeaders['content-type'];
     }
 
+    reqHeaders['x-stainless-retry-count'] = String(retryCount);
+
     this.validateHeaders(reqHeaders, headers);
 
     return reqHeaders;
@@ -409,13 +416,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 cd5f2a0a9..a6fa97199 100644
--- a/tests/index.test.ts
+++ b/tests/index.test.ts
@@ -241,6 +241,31 @@ 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 OpenAI({ apiKey: 'My API Key', 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> => {