diff --git a/.changeset/retry-internal-btql-requests.md b/.changeset/retry-internal-btql-requests.md new file mode 100644 index 000000000..eaa4c29e0 --- /dev/null +++ b/.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. diff --git a/js/src/logger.test.ts b/js/src/logger.test.ts index a7516207a..86a3a1a5c 100644 --- a/js/src/logger.test.ts +++ b/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(); diff --git a/js/src/logger.ts b/js/src/logger.ts index 6da93a9ae..a22cb30f9 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -1320,6 +1320,7 @@ class HTTPConnection { path: string, params?: Record | 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 { + if (!signal) { + await new Promise((resolve) => setTimeout(resolve, delayMs)); + return; + } + + if (signal.aborted) { + throw getAbortReason(signal); + } + + await new Promise((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 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) {