feat: Implement rate limit waiting instead of erroring

Author: ericc-ch

src/lib/rate-limit.ts

@@ -1,10 +1,11 @@
+import { sleep } from "bun"
 import consola from "consola"
 
 import type { State } from "./state"
 
 import { HTTPError } from "./http-error"
 
-export function checkRateLimit(state: State) {
+export async function checkRateLimit(state: State) {
   if (state.rateLimitSeconds === undefined) return
 
   const now = Date.now()
@@ -21,10 +22,22 @@ export function checkRateLimit(state: State) {
     return
   }
 
-  const waitTimeSeconds = Math.round(state.rateLimitSeconds - elapsedSeconds)
+  const waitTimeSeconds = Math.ceil(state.rateLimitSeconds - elapsedSeconds)
+
+  if (!state.rateLimitWait) {
+    consola.warn(
+      `Rate limit exceeded. Need to wait ${waitTimeSeconds} more seconds.`,
+    )
+    throw new HTTPError("Rate limit exceeded", Response.json({ status: 429 }))
+  }
+
+  const waitTimeMs = waitTimeSeconds * 1000
   consola.warn(
-    `Rate limit exceeded. Need to wait ${waitTimeSeconds} more seconds.`,
+    `Rate limit reached. Waiting ${waitTimeSeconds} seconds before proceeding...`,
   )
-
-  throw new HTTPError("Rate limit exceeded", Response.json({ status: 429 }))
+  await sleep(waitTimeMs)
+  // eslint-disable-next-line require-atomic-updates
+  state.lastRequestTimestamp = now
+  consola.info("Rate limit wait completed, proceeding with request")
+  return
 }

src/lib/state.ts

@@ -9,6 +9,7 @@ export interface State {
   vsCodeVersion?: string
 
   manualApprove: boolean
+  rateLimitWait: boolean
 
   // Rate limiting configuration
   rateLimitSeconds?: number
@@ -18,4 +19,5 @@ export interface State {
 export const state: State = {
   accountType: "individual",
   manualApprove: false,
+  rateLimitWait: false,
 }

src/main.ts

@@ -17,6 +17,7 @@ interface RunServerOptions {
   business: boolean
   manual: boolean
   rateLimit: number | undefined
+  rateLimitWait: boolean
 }
 
 export async function runServer(options: RunServerOptions): Promise<void> {
@@ -32,6 +33,7 @@ export async function runServer(options: RunServerOptions): Promise<void> {
 
   state.manualApprove = options.manual
   state.rateLimitSeconds = options.rateLimit
+  state.rateLimitWait = options.rateLimitWait
 
   await ensurePaths()
   await cacheVSCodeVersion()
@@ -77,6 +79,13 @@ const main = defineCommand({
       type: "string",
       description: "Rate limit in seconds between requests",
     },
+    wait: {
+      alias: "w",
+      type: "boolean",
+      default: false,
+      description:
+        "Wait instead of error when rate limit is hit. Has no effect if rate limit is not set",
+    },
   },
   run({ args }) {
     const rateLimitRaw = args["rate-limit"]
@@ -85,14 +94,14 @@ const main = defineCommand({
       rateLimitRaw === undefined ? undefined : Number.parseInt(rateLimitRaw, 10)
 
     const port = Number.parseInt(args.port, 10)
-    // const rateLimit = Number.parseInt(args["rate-limit"], 10)
 
     return runServer({
       port,
       verbose: args.verbose,
       business: args.business,
       manual: args.manual,
       rateLimit,
+      rateLimitWait: Boolean(args["rate-limit-wait"]),
     })
   },
 })

src/routes/chat-completions/handler.ts

@@ -15,7 +15,7 @@ import {
 } from "~/services/copilot/create-chat-completions"
 
 export async function handleCompletion(c: Context) {
-  checkRateLimit(state)
+  await checkRateLimit(state)
 
   let payload = await c.req.json<ChatCompletionsPayload>()