Explorar o código

go: rate limit metadata

Frank hai 3 meses
pai
achega
72ec05d0be

+ 10 - 1
packages/console/app/src/routes/zen/util/error.ts

@@ -13,4 +13,13 @@ class LimitError extends Error {
 }
 export class RateLimitError extends LimitError {}
 export class FreeUsageLimitError extends LimitError {}
-export class SubscriptionUsageLimitError extends LimitError {}
+
+class SubscriptionUsageLimitError extends LimitError {
+  workspace: string
+  constructor(message: string, workspace: string, retryAfter?: number) {
+    super(message, retryAfter)
+    this.workspace = workspace
+  }
+}
+export class GoUsageLimitError extends SubscriptionUsageLimitError {}
+export class BlackUsageLimitError extends SubscriptionUsageLimitError {}

+ 22 - 8
packages/console/app/src/routes/zen/util/handler.ts

@@ -23,7 +23,8 @@ import {
   ModelError,
   RateLimitError,
   FreeUsageLimitError,
-  SubscriptionUsageLimitError,
+  GoUsageLimitError,
+  BlackUsageLimitError,
 } from "./error"
 import {
   buildCostChunk,
@@ -395,7 +396,8 @@ export async function handler(
     if (
       error instanceof RateLimitError ||
       error instanceof FreeUsageLimitError ||
-      error instanceof SubscriptionUsageLimitError
+      error instanceof GoUsageLimitError ||
+      error instanceof BlackUsageLimitError
     ) {
       const headers = new Headers()
       if (error.retryAfter) {
@@ -404,7 +406,14 @@ export async function handler(
       return new Response(
         JSON.stringify({
           type: "error",
-          error: { type: error.constructor.name, message: error.message },
+          error: {
+            type: error.constructor.name,
+            message: error.message,
+          },
+          metadata:
+            error instanceof GoUsageLimitError || error instanceof BlackUsageLimitError
+              ? { workspace: error.workspace }
+              : {},
         }),
         { status: 429, headers },
       )
@@ -693,10 +702,11 @@ export async function handler(
             timeUpdated: sub.timeFixedUpdated,
           })
           if (result.status === "rate-limited")
-            throw new SubscriptionUsageLimitError(
+            throw new BlackUsageLimitError(
               t("zen.api.error.subscriptionQuotaExceeded", {
                 retryIn: formatRetryTime(result.resetInSec),
               }),
+              authInfo.workspaceID,
               result.resetInSec,
             )
         }
@@ -711,10 +721,11 @@ export async function handler(
             timeUpdated: sub.timeRollingUpdated,
           })
           if (result.status === "rate-limited")
-            throw new SubscriptionUsageLimitError(
+            throw new BlackUsageLimitError(
               t("zen.api.error.subscriptionQuotaExceeded", {
                 retryIn: formatRetryTime(result.resetInSec),
               }),
+              authInfo.workspaceID,
               result.resetInSec,
             )
         }
@@ -739,8 +750,9 @@ export async function handler(
             timeUpdated: sub.timeWeeklyUpdated,
           })
           if (result.status === "rate-limited")
-            throw new SubscriptionUsageLimitError(
+            throw new GoUsageLimitError(
               t("zen.api.error.subscriptionQuotaExceededUseFreeModels"),
+              authInfo.workspaceID,
               result.resetInSec,
             )
         }
@@ -754,8 +766,9 @@ export async function handler(
             timeSubscribed: sub.timeCreated,
           })
           if (result.status === "rate-limited")
-            throw new SubscriptionUsageLimitError(
+            throw new GoUsageLimitError(
               t("zen.api.error.subscriptionQuotaExceededUseFreeModels"),
+              authInfo.workspaceID,
               result.resetInSec,
             )
         }
@@ -769,8 +782,9 @@ export async function handler(
             timeUpdated: sub.timeRollingUpdated,
           })
           if (result.status === "rate-limited")
-            throw new SubscriptionUsageLimitError(
+            throw new GoUsageLimitError(
               t("zen.api.error.subscriptionQuotaExceededUseFreeModels"),
+              authInfo.workspaceID,
               result.resetInSec,
             )
         }