From 3e0ba409e57ce276fb1f95cd11c801e4ccaad572 Mon Sep 17 00:00:00 2001
From: Robert Craigie <robert@craigie.dev>
Date: Wed, 22 Jan 2025 20:34:34 +0000
Subject: [PATCH] fix(azure): include retry count header

---
 src/index.ts            |  7 +++++--
 tests/lib/azure.test.ts | 12 ++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index cf6aa89e3..944def00f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -577,7 +577,10 @@ export class AzureOpenAI extends OpenAI {
     this._deployment = deployment;
   }
 
-  override buildRequest(options: Core.FinalRequestOptions<unknown>): {
+  override buildRequest(
+    options: Core.FinalRequestOptions<unknown>,
+    props: { retryCount?: number } = {},
+  ): {
     req: RequestInit;
     url: string;
     timeout: number;
@@ -591,7 +594,7 @@ export class AzureOpenAI extends OpenAI {
         options.path = `/deployments/${model}${options.path}`;
       }
     }
-    return super.buildRequest(options);
+    return super.buildRequest(options, props);
   }
 
   private async _getAzureADToken(): Promise<string | undefined> {
diff --git a/tests/lib/azure.test.ts b/tests/lib/azure.test.ts
index 064a0098c..0e3c2c5a3 100644
--- a/tests/lib/azure.test.ts
+++ b/tests/lib/azure.test.ts
@@ -51,6 +51,18 @@ describe('instantiate azure client', () => {
       });
       expect(req.headers as Headers).not.toHaveProperty('x-my-default-header');
     });
+
+    test('includes retry count', () => {
+      const { req } = client.buildRequest(
+        {
+          path: '/foo',
+          method: 'post',
+          headers: { 'X-My-Default-Header': null },
+        },
+        { retryCount: 1 },
+      );
+      expect((req.headers as Headers)['x-stainless-retry-count']).toEqual('1');
+    });
   });
 
   describe('defaultQuery', () => {