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(options: FinalRequestOptions): { req: RequestInit; url: string; timeout: number } { + buildRequest( + options: FinalRequestOptions, + { 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; contentLength: string | null | undefined; + retryCount: number; }): Record { const reqHeaders: Record = {}; 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 { 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 => { + 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 => {