Honor Retry-After for provider rate limits

valhalla/jawn/src/lib/proxy/ProviderClient.ts

@@ -1,4 +1,3 @@
-import retry from "async-retry";
 import { HeliconeProxyRequest, RetryOptions } from "./HeliconeProxyRequest";
 import fetch from "node-fetch";
 import { Headers, Response } from "node-fetch";
@@ -12,7 +11,7 @@ export interface CallProps {
 }
 
 export function callPropsFromProxyRequest(
-  proxyRequest: HeliconeProxyRequest
+  proxyRequest: HeliconeProxyRequest,
 ): CallProps {
   return {
     apiBase: proxyRequest.api_base,
@@ -60,50 +59,92 @@ export function buildTargetUrl(originalUrl: URL, apiBase: string): URL {
   const apiBaseUrl = new URL(apiBase.replace(/\/$/, ""));
   const pathname = originalUrl.pathname.replace(
     /^\/v1\/gateway(\/[^\/]+)?/,
-    ""
+    "",
   );
 
   return new URL(`${apiBaseUrl.origin}${pathname}${originalUrl.search}`);
 }
 
+export function retryAfterMs(headers: Headers): number | null {
+  const retryAfter = headers.get("retry-after");
+  if (retryAfter === null) {
+    return null;
+  }
+
+  const retryAfterSeconds = Number(retryAfter);
+  if (!Number.isNaN(retryAfterSeconds)) {
+    return Math.max(0, retryAfterSeconds * 1000);
+  }
+
+  const retryAfterDate = Date.parse(retryAfter);
+  if (!Number.isNaN(retryAfterDate)) {
+    return Math.max(0, retryAfterDate - Date.now());
+  }
+
+  return null;
+}
+
+export function exponentialDelayMs(
+  attempt: number,
+  retryOptions: RetryOptions,
+): number {
+  const baseDelay =
+    retryOptions.minTimeout * Math.pow(retryOptions.factor, attempt - 1);
+
+  return Math.min(baseDelay, retryOptions.maxTimeout);
+}
+
+export function retryDelayMs(
+  attempt: number,
+  response: Response,
+  retryOptions: RetryOptions,
+): number {
+  const fallbackDelay = exponentialDelayMs(attempt, retryOptions);
+  const providerDelay =
+    response.status === 429 ? retryAfterMs(response.headers) : null;
+
+  return Math.max(fallbackDelay, providerDelay ?? 0);
+}
+
+export function isRetryableProviderResponse(response: Response): boolean {
+  return (
+    response.status === 429 ||
+    response.status === 500 ||
+    response.status === 522
+  );
+}
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 export async function callProviderWithRetry(
   callProps: CallProps,
-  retryOptions: RetryOptions
+  retryOptions: RetryOptions,
 ): Promise<Response> {
   let lastResponse;
 
   try {
-    // Use async-retry to call the forwardRequestToOpenAi function with exponential backoff
-    await retry(
-      async (bail, attempt) => {
-        try {
-          const res = await callProvider(callProps);
-
-          lastResponse = res;
-          // Throw an error if the status code is 429
-          if (res.status === 429 || res.status === 500 || res.status === 522) {
-            throw new Error(`Status code ${res.status}`);
-          }
-          return res;
-        } catch (e) {
-          // If we reach the maximum number of retries, bail with the error
-          if (attempt >= retryOptions.retries) {
-            bail(e as Error);
-          }
-          // Otherwise, retry with exponential backoff
-          throw e;
-        }
-      },
-      {
-        ...retryOptions,
-        onRetry: (error, attempt) => {
-          console.log(`Retry attempt ${attempt}. Error: ${error}`);
-        },
+    for (let attempt = 1; attempt <= retryOptions.retries; attempt++) {
+      const res = await callProvider(callProps);
+
+      lastResponse = res;
+      if (
+        !isRetryableProviderResponse(res) ||
+        attempt >= retryOptions.retries
+      ) {
+        break;
       }
-    );
+
+      const delay = retryDelayMs(attempt, res, retryOptions);
+      console.log(
+        `Retry attempt ${attempt}. Status code: ${res.status}. Waiting ${delay}ms`,
+      );
+      await sleep(delay);
+    }
   } catch (e) {
     console.warn(
-      `Retried ${retryOptions.retries} times but still failed. Error: ${e}`
+      `Retried ${retryOptions.retries} times but still failed. Error: ${e}`,
     );
   }
 

valhalla/jawn/src/lib/proxy/__tests__/ProviderClient.test.ts

@@ -0,0 +1,64 @@
+import { Headers, Response } from "node-fetch";
+import {
+  exponentialDelayMs,
+  isRetryableProviderResponse,
+  retryAfterMs,
+  retryDelayMs,
+} from "../ProviderClient";
+
+const retryOptions = {
+  retries: 3,
+  factor: 2,
+  minTimeout: 100,
+  maxTimeout: 1_000,
+};
+
+describe("ProviderClient retry helpers", () => {
+  it("parses numeric Retry-After headers as seconds", () => {
+    const headers = new Headers({ "retry-after": "2.5" });
+
+    expect(retryAfterMs(headers)).toBe(2_500);
+  });
+
+  it("parses date Retry-After headers as a delay from now", () => {
+    jest.useFakeTimers().setSystemTime(new Date("2026-05-26T00:00:00.000Z"));
+    const headers = new Headers({
+      "retry-after": "Tue, 26 May 2026 00:00:03 GMT",
+    });
+
+    expect(retryAfterMs(headers)).toBe(3_000);
+
+    jest.useRealTimers();
+  });
+
+  it("uses Retry-After as a minimum delay for 429 responses", () => {
+    const response = new Response("", {
+      status: 429,
+      headers: { "retry-after": "2" },
+    });
+
+    expect(retryDelayMs(1, response, retryOptions)).toBe(2_000);
+  });
+
+  it("falls back to exponential delay when Retry-After is absent", () => {
+    const response = new Response("", { status: 500 });
+
+    expect(exponentialDelayMs(3, retryOptions)).toBe(400);
+    expect(retryDelayMs(3, response, retryOptions)).toBe(400);
+  });
+
+  it("only retries provider rate-limit and gateway-style transient failures", () => {
+    expect(isRetryableProviderResponse(new Response("", { status: 429 }))).toBe(
+      true,
+    );
+    expect(isRetryableProviderResponse(new Response("", { status: 500 }))).toBe(
+      true,
+    );
+    expect(isRetryableProviderResponse(new Response("", { status: 522 }))).toBe(
+      true,
+    );
+    expect(isRetryableProviderResponse(new Response("", { status: 400 }))).toBe(
+      false,
+    );
+  });
+});