feat: retry failed aichat requests (#110)

Author: ChristopherChudzicki

package.json

@@ -67,6 +67,7 @@
     "@mui/utils": "^6.1.6",
     "better-react-mathjax": "^2.3.0",
     "classnames": "^2.5.1",
+    "fetch-retry": "^6.0.0",
     "lodash": "^4.17.21",
     "react-markdown": "^10.0.0",
     "rehype-mathjax": "^7.1.0",

src/components/AiChat/AiChatContext.tsx

@@ -2,14 +2,15 @@ import * as React from "react"
 import { useChat, UseChatHelpers } from "@ai-sdk/react"
 import type { RequestOpts, AiChatMessage, AiChatContextProps } from "./types"
 import { useMemo, createContext } from "react"
+import retryingFetch from "../../utils/retryingFetch"
 
 const identity = <T,>(x: T): T => x
 
 const getFetcher: (requestOpts: RequestOpts) => typeof fetch =
   (requestOpts: RequestOpts) => async (url, opts) => {
     if (typeof opts?.body !== "string") {
       console.error("Unexpected body type.")
-      return window.fetch(url, opts)
+      return retryingFetch(url, opts)
     }
     const messages: AiChatMessage[] = JSON.parse(opts?.body).messages
     const transformBody: RequestOpts["transformBody"] =
@@ -24,7 +25,7 @@ const getFetcher: (requestOpts: RequestOpts) => typeof fetch =
         ...requestOpts.fetchOpts?.headers,
       },
     }
-    return fetch(url, options)
+    return retryingFetch(url, options)
   }
 
 /**

src/components/AiChat/types.ts

@@ -100,6 +100,11 @@ type AiChatDisplayProps = {
    * Defaults to false.
    */
   useMathJax?: boolean
+  /**
+   * If true, the chat input will be autofocused on load.
+   * Defaults to true.
+   */
+  autofocus?: boolean
 } & RefAttributes<HTMLDivElement>
 
 type AiChatProps = AiChatContextProps & AiChatDisplayProps

src/utils/retryingFetch.test.ts

@@ -0,0 +1,51 @@
+import retryingFetch from "./retryingFetch"
+import { http, HttpResponse } from "msw"
+import { setupServer } from "msw/node"
+
+const counter = jest.fn() // use jest.fn as counter because it resets on each test
+const NETWORK_SUCCESS_URL = "http://localhost:3456/success"
+const NETWORK_ERROR_URL = "http://localhost:3456/error"
+const server = setupServer(
+  http.get(NETWORK_SUCCESS_URL, async ({ request }) => {
+    counter()
+    const url = new URL(request.url)
+    const status = +(url.searchParams.get("status") ?? 200)
+    return HttpResponse.text(`Status ${status}`, { status })
+  }),
+  http.get(NETWORK_ERROR_URL, async () => {
+    counter()
+    return HttpResponse.error()
+  }),
+)
+beforeAll(() => server.listen())
+afterEach(() => server.resetHandlers())
+afterAll(() => server.close())
+
+describe("retryingFetch", () => {
+  beforeAll(() => {})
+
+  test.each([200, 201, 202, 367, 400, 401, 403])(
+    "should not retry on %s",
+    async (status) => {
+      const result = await retryingFetch(
+        `${NETWORK_SUCCESS_URL}?status=${status}`,
+      )
+      expect(await result.text()).toBe(`Status ${status}`)
+      expect(counter).toHaveBeenCalledTimes(1)
+    },
+  )
+
+  test.each([429, 500, 501, 502, 503])("should retry on %s", async (status) => {
+    const result = await retryingFetch(
+      `${NETWORK_SUCCESS_URL}?status=${status}`,
+    )
+    expect(await result.text()).toBe(`Status ${status}`)
+    expect(counter).toHaveBeenCalledTimes(4)
+  })
+
+  test("should retry on error", async () => {
+    const result = await retryingFetch(NETWORK_ERROR_URL).catch((err) => err)
+    expect(result).toBeInstanceOf(Error)
+    expect(counter).toHaveBeenCalledTimes(4)
+  })
+})

src/utils/retryingFetch.ts

@@ -0,0 +1,44 @@
+import fetchRetrier from "fetch-retry"
+
+const MAX_RETRIES = 3
+const NO_RETRY_ERROR_CODES = [400, 401, 403, 404, 405, 409, 422]
+const BASE_RETRY_DELAY = process.env.NODE_ENV === "test" ? 10 : 1000
+
+// This is a workaround to ensure retryingFetch uses MSW's fetch version in tests.
+// See https://github.com/jonbern/fetch-retry/issues/95#issuecomment-2613990480
+const fetch: typeof window.fetch = (...args) => window.fetch(...args)
+
+/**
+ * A wrapper around fetch that retries the request on certain error codes.
+ *
+ * By default:
+ *
+ * Requests will be retried if max retries not exceeded and :
+ *  - The status code is >= 400 AND NOT in the NO_RETRY_ERROR_CODES list,
+ *  - OR the request promise rejected (network error)
+ *
+ * The retry delay is exponential, 1s, 2s, 4s, 8s, ... maxing at 30s.
+ *
+ * Maximum retries is 3.
+ *
+ * NOTE:
+ *  - When NODE_ENV="test", the maximum retries is set 0 by default but can be
+ *  set via the TEST_ENV_MAX_RETRIES environment variable.
+ */
+const retryingFetch = fetchRetrier(fetch, {
+  retryDelay: (attempt, _error, _response) => {
+    return Math.min(BASE_RETRY_DELAY * 2 ** attempt, 30_000)
+  },
+  retryOn: (attempt, _error, response) => {
+    if (attempt >= MAX_RETRIES) {
+      return false
+    }
+    if (response) {
+      if (response.status < 400) return false
+      return !NO_RETRY_ERROR_CODES.includes(response.status)
+    }
+    return true
+  },
+})
+
+export default retryingFetch

yarn.lock

@@ -2834,6 +2834,7 @@ __metadata:
     eslint-plugin-react-hooks: "npm:^5.0.0"
     eslint-plugin-styled-components-a11y: "npm:^2.1.35"
     eslint-plugin-testing-library: "npm:^7.0.0"
+    fetch-retry: "npm:^6.0.0"
     jest: "npm:^29.7.0"
     jest-environment-jsdom: "npm:^29.5.0"
     jest-extended: "npm:^4.0.2"
@@ -9229,6 +9230,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"fetch-retry@npm:^6.0.0":
+  version: 6.0.0
+  resolution: "fetch-retry@npm:6.0.0"
+  checksum: 10/0c8d3082e2d76fff2df75adef6280bc854bc36fd3ef38506674f0216d0d819e2efd14da7477d3f1732415aea1d2cfde7cd3e1aeae46f45f2adbfc5133296e8de
+  languageName: node
+  linkType: hard
+
 "figures@npm:^2.0.0":
   version: 2.0.0
   resolution: "figures@npm:2.0.0"