fix: retry transient failures for internal BTQL requests

.changeset/retry-internal-btql-requests.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+Retry transient failures for internal, read-only BTQL requests used to fetch datasets, objects, spans, and prompt versions.

js/src/logger.test.ts

@@ -717,6 +717,64 @@ test("dataset fetch forwards _internal_btql filter arrays to btql", async () =>
   }
 });
 
+test("dataset fetch retries an internally generated BTQL request after a transient 500", async () => {
+  vi.useFakeTimers();
+  const state = await _exportsForTestingOnly.simulateLoginForTests();
+
+  try {
+    vi.spyOn(state, "login").mockResolvedValue(state as any);
+    vi.spyOn(state.appConn(), "post_json").mockResolvedValue({
+      project: {
+        id: "00000000-0000-0000-0000-000000000001",
+        name: "test-project",
+      },
+      dataset: {
+        id: "00000000-0000-0000-0000-000000000002",
+        name: "test-dataset",
+      },
+    });
+
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(
+        new Response("Internal server error", {
+          status: 500,
+          statusText: "Internal Server Error",
+        }),
+      )
+      .mockResolvedValueOnce(
+        new Response(JSON.stringify({ data: [] }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        }),
+      );
+    state.setFetch(fetchMock as unknown as typeof globalThis.fetch);
+
+    const dataset = initDataset({
+      project: "test-project",
+      dataset: "test-dataset",
+      state,
+    });
+
+    const fetchedData = dataset.fetchedData();
+    await vi.advanceTimersByTimeAsync(0);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+
+    await vi.advanceTimersByTimeAsync(1_200);
+    await expect(fetchedData).resolves.toEqual([]);
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    expect(fetchMock).toHaveBeenNthCalledWith(
+      1,
+      expect.stringMatching(/\/btql$/),
+      expect.objectContaining({ method: "POST" }),
+    );
+  } finally {
+    vi.useRealTimers();
+    _exportsForTestingOnly.simulateLogoutForTests();
+    vi.restoreAllMocks();
+  }
+});
+
 test("initDataset applies bt eval internal BTQL runtime value to eval data", async () => {
   const state = await _exportsForTestingOnly.simulateLoginForTests();
   const previousInternalBtql = globalThis.__bt_eval_internal_btql;
@@ -2101,6 +2159,64 @@ test("simulateLoginForTests and simulateLogoutForTests", async () => {
   }
 });
 
+describe("HTTPConnection POST retries", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+    _exportsForTestingOnly.simulateLogoutForTests();
+  });
+
+  test("does not retry ordinary POST requests", async () => {
+    const state = await _exportsForTestingOnly.simulateLoginForTests();
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(
+        new Response("Internal server error", {
+          status: 500,
+          statusText: "Internal Server Error",
+        }),
+      )
+      .mockResolvedValueOnce(new Response(null, { status: 200 }));
+    state.setFetch(fetchMock as unknown as typeof globalThis.fetch);
+
+    await expect(state.apiConn().post("write", {})).rejects.toThrow(
+      "500: Internal Server Error",
+    );
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  test("does not retry a non-transient response", async () => {
+    const state = await _exportsForTestingOnly.simulateLoginForTests();
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(
+        new Response("Invalid query", {
+          status: 400,
+          statusText: "Bad Request",
+        }),
+      )
+      .mockResolvedValueOnce(new Response(null, { status: 200 }));
+    state.setFetch(fetchMock as unknown as typeof globalThis.fetch);
+
+    await expect(
+      state.apiConn().post("btql", {}, undefined, 3),
+    ).rejects.toThrow("400: Bad Request");
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  test("does not retry an aborted POST request", async () => {
+    const state = await _exportsForTestingOnly.simulateLoginForTests();
+    const controller = new AbortController();
+    controller.abort();
+    const fetchMock = vi.fn().mockRejectedValue(controller.signal.reason);
+    state.setFetch(fetchMock as unknown as typeof globalThis.fetch);
+
+    await expect(
+      state.apiConn().post("btql", {}, { signal: controller.signal }, 3),
+    ).rejects.toBe(controller.signal.reason);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+});
+
 describe("HTTPConnection get_json retries", () => {
   afterEach(() => {
     vi.useRealTimers();

js/src/logger.ts

@@ -1320,6 +1320,7 @@ class HTTPConnection {
     path: string,
     params?: Record<string, unknown> | string,
     config?: RequestInit,
+    retries: number = 0,
   ) {
     const { headers, ...rest } = config || {};
     // On platforms like Cloudflare, we lose "this" when we make an async call,
@@ -1328,25 +1329,50 @@ class HTTPConnection {
     const this_base_url = this.base_url;
     const this_headers = this.headers;
 
-    return await checkResponse(
-      await this_fetch(_urljoin(this_base_url, path), {
-        method: "POST",
-        headers: {
-          Accept: "application/json",
-          "Content-Type": "application/json",
-          ...this_headers,
-          ...headers,
-        },
-        body:
-          typeof params === "string"
-            ? params
-            : params
-              ? JSON.stringify(params)
-              : undefined,
-        keepalive: true,
-        ...rest,
-      }),
-    );
+    const tries = retries + 1;
+    for (let i = 0; i < tries; i++) {
+      try {
+        return await checkResponse(
+          await this_fetch(_urljoin(this_base_url, path), {
+            method: "POST",
+            headers: {
+              Accept: "application/json",
+              "Content-Type": "application/json",
+              ...this_headers,
+              ...headers,
+            },
+            body:
+              typeof params === "string"
+                ? params
+                : params
+                  ? JSON.stringify(params)
+                  : undefined,
+            keepalive: true,
+            ...rest,
+          }),
+        );
+      } catch (error) {
+        if (config?.signal?.aborted) {
+          throw getAbortReason(config.signal);
+        }
+        if (i === tries - 1 || !isRetryableHTTPError(error)) {
+          throw error;
+        }
+
+        debugLogger.debug(
+          `Retrying API request ${path} after ${formatHTTPError(error)}`,
+        );
+        const sleepTimeMs =
+          HTTP_RETRY_BASE_SLEEP_TIME_S * 1000 * 2 ** i +
+          Math.random() * HTTP_RETRY_JITTER_MS;
+        debugLogger.info(
+          `Sleeping for ${sleepTimeMs}ms before retrying API request`,
+        );
+        await waitForRetry(sleepTimeMs, config?.signal);
+      }
+    }
+
+    throw new Error("Unexpected retry state");
   }
 
   async get_json(
@@ -2914,6 +2940,57 @@ export class TestBackgroundLogger implements BackgroundLogger {
 
 const BACKGROUND_LOGGER_BASE_SLEEP_TIME_S = 1.0;
 const HTTP_RETRY_BASE_SLEEP_TIME_S = 1.0;
+const HTTP_RETRY_JITTER_MS = 200;
+const BTQL_HTTP_RETRIES = 3;
+const RETRYABLE_HTTP_STATUS_CODES = new Set([500, 502, 503, 504]);
+
+function isRetryableHTTPError(error: unknown): boolean {
+  return (
+    !(error instanceof FailedHTTPResponse) ||
+    RETRYABLE_HTTP_STATUS_CODES.has(error.status)
+  );
+}
+
+function formatHTTPError(error: unknown): string {
+  if (error instanceof FailedHTTPResponse) {
+    return `${error.status} ${error.text}`;
+  }
+  return error instanceof Error ? error.message : String(error);
+}
+
+function getAbortReason(signal: AbortSignal): unknown {
+  return signal.reason ?? new Error("Request aborted");
+}
+
+async function waitForRetry(
+  delayMs: number,
+  signal?: AbortSignal | null,
+): Promise<void> {
+  if (!signal) {
+    await new Promise((resolve) => setTimeout(resolve, delayMs));
+    return;
+  }
+
+  if (signal.aborted) {
+    throw getAbortReason(signal);
+  }
+
+  await new Promise<void>((resolve, reject) => {
+    const onAbort = () => {
+      clearTimeout(timeout);
+      signal.removeEventListener("abort", onAbort);
+      reject(getAbortReason(signal));
+    };
+    const timeout = setTimeout(() => {
+      signal.removeEventListener("abort", onAbort);
+      resolve();
+    }, delayMs);
+    signal.addEventListener("abort", onAbort, { once: true });
+    if (signal.aborted) {
+      onAbort();
+    }
+  });
+}
 
 // We should only have one instance of this object per state object in
 // 'BraintrustState._bgLogger'. Be careful about spawning multiple
@@ -6819,6 +6896,7 @@ export class ObjectFetcher<RecordType> implements AsyncIterable<
             : {}),
         },
         { headers: { "Accept-Encoding": "gzip" } },
+        BTQL_HTTP_RETRIES,
       );
       const respJson = await resp.json();
       const mutate = this.mutateRecord;
@@ -9414,6 +9492,7 @@ export async function getPromptVersions(
       brainstore_realtime: true,
     },
     { headers: { "Accept-Encoding": "gzip" } },
+    BTQL_HTTP_RETRIES,
   );
 
   if (!response.ok) {