Merge pull request #21286 from ahmedhamidawan/add_rate_limiter_to_client_api_requests

[25.1] Add a rate limiter to the API client factory

Author: mvdbeek

client/src/api/client/index.ts

@@ -1,5 +1,6 @@
 import createClient from "openapi-fetch";
 
+import { createRateLimiterMiddleware } from "@/api/client/rateLimiter";
 import type { GalaxyApiPaths } from "@/api/schema";
 import { getAppRoot } from "@/onload/loadConfig";
 
@@ -9,7 +10,19 @@ function getBaseUrl() {
 }
 
 function apiClientFactory() {
-    return createClient<GalaxyApiPaths>({ baseUrl: getBaseUrl() });
+    const client = createClient<GalaxyApiPaths>({ baseUrl: getBaseUrl() });
+
+    // TODO: Adjust based on server limits (maybe this goes in Galaxy config?)
+    client.use(
+        createRateLimiterMiddleware({
+            maxRequests: 100,
+            windowMs: 3000,
+            retryDelay: 1000,
+            maxRetries: 3,
+        }),
+    );
+
+    return client;
 }
 
 export type GalaxyApiClient = ReturnType<typeof apiClientFactory>;

client/src/api/client/rateLimiter.test.ts

@@ -0,0 +1,139 @@
+import type { MessageException } from "@/api";
+import { GalaxyApi } from "@/api/client";
+import { useServerMock } from "@/api/client/__mocks__";
+
+import { DEFAULT_CONFIG } from "./rateLimiter";
+
+const { server, http } = useServerMock();
+
+/** Spy to count number of times a 429 response is returned */
+const mock429ResponseSpy = jest.fn();
+
+describe("Rate Limiter Middleware", () => {
+    let consoleWarnSpy: jest.SpyInstance;
+    let consoleErrorSpy: jest.SpyInstance;
+
+    /** Helper to verify that there is a 429 response without retries */
+    function ensure429AndNoRetries(response: Response, error: MessageException | undefined) {
+        // Verify the request failed as expected
+        expect(response.status).toBe(429);
+        expect(error).toBeDefined();
+        expect(error?.err_code).toBe(429);
+
+        // Verify no retry behavior occurred by confirming console warns/errors were not called
+        expect(consoleWarnSpy).not.toHaveBeenCalled();
+        expect(consoleErrorSpy).not.toHaveBeenCalled();
+
+        // Verify the mock 429 response was called only once (no retries)
+        expect(mock429ResponseSpy).toHaveBeenCalledTimes(1);
+    }
+
+    beforeEach(() => {
+        consoleWarnSpy = jest.spyOn(console, "warn").mockImplementation();
+        consoleErrorSpy = jest.spyOn(console, "error").mockImplementation();
+    });
+    afterEach(() => {
+        consoleWarnSpy.mockRestore();
+        consoleErrorSpy.mockRestore();
+        mock429ResponseSpy.mockReset();
+    });
+
+    it("should retry 429 GET responses", async () => {
+        // Set up a mock GET endpoint that always returns 429
+        server.use(
+            http.get("/api/histories/{history_id}", ({ response }) => {
+                mock429ResponseSpy();
+                return response("4XX").json({ err_code: 429, err_msg: "Too Many Requests" }, { status: 429 });
+            }),
+        );
+
+        const { error, response } = await GalaxyApi().GET("/api/histories/{history_id}", {
+            params: {
+                path: { history_id: "test" },
+            },
+        });
+
+        // Verify the request failed as expected
+        expect(response.status).toBe(429);
+        expect(error).toBeDefined();
+        expect(error?.err_code).toBe(429);
+
+        // Verify retry behavior occurred by confirming console warns/errors
+        expect(consoleWarnSpy).toHaveBeenCalledWith(
+            expect.stringContaining(`Received 429 from server, waiting ${DEFAULT_CONFIG.retryDelay}ms before retry`),
+        );
+
+        for (let i = 1; i <= DEFAULT_CONFIG.maxRetries; i++) {
+            expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining(`Retry ${i} also received 429`));
+        }
+
+        expect(consoleErrorSpy).toHaveBeenCalledWith(
+            expect.stringContaining(`Max retries reached for request to ${response.url}`),
+        );
+
+        // Verify the mock 429 response was called the first time and then for each retry
+        expect(mock429ResponseSpy).toHaveBeenCalledTimes(DEFAULT_CONFIG.maxRetries + 1);
+    });
+
+    it("should not retry 429 POST responses", async () => {
+        // Set up a mock POST endpoint that always returns 429
+        server.use(
+            http.post("/api/chat", ({ response }) => {
+                mock429ResponseSpy();
+                return response("4XX").json({ err_code: 429, err_msg: "Too Many Requests" }, { status: 429 });
+            }),
+        );
+
+        const { error, response } = await GalaxyApi().POST("/api/chat", {
+            params: {
+                query: { job_id: "test" },
+            },
+            body: {
+                query: "test message",
+                context: "test",
+            },
+        });
+
+        ensure429AndNoRetries(response, error);
+    });
+
+    it("should not retry 429 DELETE responses", async () => {
+        // Set up a mock DELETE endpoint that always returns 429
+        server.use(
+            http.delete("/api/datasets/{dataset_id}", ({ response }) => {
+                mock429ResponseSpy();
+                return response("4XX").json({ err_code: 429, err_msg: "Too Many Requests" }, { status: 429 });
+            }),
+        );
+
+        const { error, response } = await GalaxyApi().DELETE("/api/datasets/{dataset_id}", {
+            params: {
+                path: { dataset_id: "test_id" },
+                query: { purge: true },
+            },
+        });
+
+        ensure429AndNoRetries(response, error);
+    });
+
+    it("should not retry 429 PUT responses", async () => {
+        // Set up a mock PUT endpoint that always returns 429
+        server.use(
+            http.put("/api/datasets/{dataset_id}", ({ response }) => {
+                mock429ResponseSpy();
+                return response("4XX").json({ err_code: 429, err_msg: "Too Many Requests" }, { status: 429 });
+            }),
+        );
+
+        const { error, response } = await GalaxyApi().PUT("/api/datasets/{dataset_id}", {
+            params: {
+                path: { dataset_id: "test_id" },
+            },
+            body: {
+                deleted: false,
+            },
+        });
+
+        ensure429AndNoRetries(response, error);
+    });
+});

client/src/api/client/rateLimiter.ts

@@ -0,0 +1,94 @@
+import type { Middleware } from "openapi-fetch";
+
+interface RateLimitConfig {
+    /** Maximum requests per window */
+    maxRequests?: number;
+    /** Time the requestwindow lasts in milliseconds */
+    windowMs?: number;
+    /** Delay between retries on 429 */
+    retryDelay?: number;
+    /** Maximum retry attempts */
+    maxRetries?: number;
+}
+
+export const DEFAULT_CONFIG: Required<RateLimitConfig> = {
+    maxRequests: 100,
+    windowMs: 60000,
+    retryDelay: 1000,
+    maxRetries: 3,
+};
+
+/**
+ * Rate limiting middleware to control the rate of API requests.
+ *
+ * Uses a timed window and only allows a maximum number of requests per that window.
+ */
+export function createRateLimiterMiddleware(config: RateLimitConfig = {}): Middleware {
+    const cfg = { ...DEFAULT_CONFIG, ...config };
+
+    let requestCount = 0;
+    let windowStart = Date.now();
+
+    /** Resets the request window if the time has elapsed */
+    function resetWindowIfNeeded() {
+        const now = Date.now();
+        if (now - windowStart >= cfg.windowMs) {
+            requestCount = 0;
+            windowStart = now;
+        }
+    }
+
+    /** Places a request in the rate limiter queue */
+    async function placeRequestInQueue(): Promise<void> {
+        resetWindowIfNeeded();
+
+        if (requestCount < cfg.maxRequests) {
+            requestCount++;
+            return;
+        }
+
+        // Rate limit exceeded, wait for next window
+        const waitTime = cfg.windowMs - (Date.now() - windowStart);
+
+        await new Promise((resolve) => setTimeout(resolve, waitTime));
+
+        // After waiting, try again (this will reset the window)
+        return placeRequestInQueue();
+    }
+
+    const middleware: Middleware = {
+        async onRequest({ request }) {
+            await placeRequestInQueue();
+            return request;
+        },
+
+        async onResponse({ response: res, request }) {
+            // Handle 429 Too Many Requests from server
+            if (res.status === 429 && request.method === "GET") {
+                const retryAfter = res.headers.get("Retry-After");
+                const delay = retryAfter ? parseInt(retryAfter) * 1000 : cfg.retryDelay;
+
+                console.warn(`Received 429 from server, waiting ${delay}ms before retry`);
+
+                let retries = 0;
+                while (retries < cfg.maxRetries) {
+                    retries++;
+                    await new Promise((resolve) => setTimeout(resolve, delay));
+
+                    // A tricky thing here is that we will bypass the middleware chain on retry
+                    const retryResponse = await fetch(request);
+                    if (retryResponse.status !== 429) {
+                        return retryResponse;
+                    }
+                    console.warn(`Retry ${retries} also received 429, retrying...`);
+                }
+
+                console.error(`Max retries reached for request to ${request.url}`);
+            }
+
+            return res;
+        },
+    };
+
+    return middleware;
+}