feat: Retry model provider throttling

control-plane/src/modules/models/index.test.ts

@@ -53,6 +53,7 @@ describe("buildModel", () => {
       });
     });
 
+
     it("should not retry other errors", async () => {
       const error = new Error("");
       mockCreate.mockImplementationOnce(() => {
@@ -77,5 +78,48 @@ describe("buildModel", () => {
         identifier: "claude-3-haiku",
       });
     });
+
+    it.skip("should throw after exhausting retries", async () => {
+      mockCreate.mockImplementation(() => {
+        throw new RetryableError("");
+      });
+
+      const model = buildModel({
+        identifier: "claude-3-haiku",
+      });
+
+      await expect(
+        () => model.call({
+          messages: [],
+        })
+      ).rejects.toThrow(RetryableError);
+
+      expect(getRouting).toHaveBeenCalledTimes(6);
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 0,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 1,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 2,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 3,
+        identifier: "claude-3-haiku",
+      });
+
+      expect(getRouting).toHaveBeenCalledWith({
+        index: 4,
+        identifier: "claude-3-haiku",
+      });
+    }, 60_000);
   });
 });

control-plane/src/utilities/errors.test.ts

@@ -0,0 +1,14 @@
+import { isRetryableError, RetryableError } from "./errors";
+import { RateLimitError } from "@anthropic-ai/sdk";
+
+describe("isRetryableError", () => {
+  it("should return true for retryable errors", async () => {
+    expect(isRetryableError(new Error())).toBe(false);
+
+    expect(isRetryableError(new RateLimitError(429, new Error(), "", {}))).toBe(true);
+    expect(isRetryableError(new RetryableError(""))).toBe(true);
+
+
+    expect(isRetryableError(new Error("429 Too many requests, please wait before trying again."))).toBe(true);
+  });
+})

control-plane/src/utilities/errors.ts

@@ -1,22 +1,36 @@
+import { RateLimitError } from "@anthropic-ai/sdk";
+
 export class RetryableError extends Error {
   constructor(message: string) {
     super(message);
     this.name = "RetryableError";
   }
 }
 
+const retryableErrors = [RetryableError, RateLimitError]
 const retryableErrorMessages = [
-  "Connection terminated due to connection timeout",
+  // DB Connection Errors
+  "connection terminated due to connection timeout",
   "timeout exceeded when trying to connect",
-  "Connection terminated unexpectedly",
-  "account does not have an agreement to this model",
+  "connection terminated unexpectedly",
+  // DB Connection Pool Exhaustion
+  "remaining connection slots are reserved for roles with the SUPERUSER attribute",
+  "too many clients already",
+  // Bedrock Errors
+  "503 bedrock is unable to process your request",
+  "429 too many requests"
 ];
 
 export const isRetryableError = (error: unknown) => {
-  if (error instanceof Error && retryableErrorMessages.includes(error.message))
+  if (error instanceof Error && retryableErrorMessages.find((message) => error.message.toLowerCase().includes(message))) {
+    return true
+  }
+
+  if (error instanceof Error && retryableErrors.find((type) => error instanceof type)) {
     return true;
+  }
 
-  return error instanceof RetryableError;
+  return false;
 };
 
 export class AuthenticationError extends Error {