fixes broken AWS rate limit backoff

src/proxy/aws.ts

@@ -3,6 +3,7 @@ import { createProxyMiddleware } from "http-proxy-middleware";
 import { v4 } from "uuid";
 import { config } from "../config";
 import { logger } from "../logger";
+import { keyPool } from "../shared/key-management";
 import { createQueueMiddleware } from "./queue";
 import { ipLimiter } from "./rate-limit";
 import { handleProxyError } from "./middleware/common";
@@ -134,6 +135,7 @@ const awsProxy = createQueueMiddleware(
     on: {
       proxyReq: createOnProxyReqHandler({
         pipeline: [
+          (_, req) => keyPool.throttle(req.key!),
           applyQuotaLimits,
           // Credentials are added by signAwsRequest preprocessor
           languageFilter,

src/proxy/middleware/request/add-key.ts

@@ -59,6 +59,7 @@ export const addKey: ProxyRequestMiddleware = (proxyReq, req) => {
     }
   }
 
+  keyPool.throttle(assignedKey);
   req.key = assignedKey;
   req.log.info(
     {

src/proxy/middleware/response/handle-streamed-response.ts

@@ -4,7 +4,7 @@ import { promisify } from "util";
 import {
   buildFakeSse,
   copySseResponseHeaders,
-  initializeSseStream
+  initializeSseStream,
 } from "../../../shared/streaming";
 import { enqueue } from "../../queue";
 import { decodeResponseBody, RawResponseBodyHandler, RetryableError } from ".";
@@ -83,7 +83,7 @@ export const handleStreamedResponse: RawResponseBodyHandler = async (
     return aggregator.getFinalResponse();
   } catch (err) {
     if (err instanceof RetryableError) {
-      keyPool.markRateLimited(req.key!)
+      keyPool.markRateLimited(req.key!);
       req.log.warn(
         { key: req.key!.hash, retryCount: req.retryCount },
         `Re-enqueueing request due to retryable error during streaming response.`