feat(vscode): show status of rate limited (#3635)

* feat(vscode): show status of rate limited * [autofix.ci] apply automated fixes * update * [autofix.ci] apply automated fixes * update * update * update * update * [autofix.ci] apply automated fixes * Update clients/tabby-agent/src/protocol.ts Co-authored-by: Zhiming Ma <[email protected]> * Update clients/tabby-agent/src/status.ts Co-authored-by: Zhiming Ma <[email protected]> * Update clients/tabby-agent/src/status.ts Co-authored-by: Zhiming Ma <[email protected]> * Update clients/vscode/src/commands/commandPalette.ts Co-authored-by: Zhiming Ma <[email protected]> * [autofix.ci] apply automated fixes * update * [autofix.ci] apply automated fixes * update --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Zhiming Ma <[email protected]>
TabbyML · Dec 31, 2024 · 1b09dcd · 1b09dcd
1 parent bafd011
commit 1b09dcd
Show file tree

Hide file tree

Showing 6 changed files with 49 additions and 3 deletions.
diff --git a/clients/tabby-agent/src/http/tabbyApiClient.ts b/clients/tabby-agent/src/http/tabbyApiClient.ts
@@ -26,6 +26,7 @@ import {
   isUnauthorizedError,
   isCanceledError,
   isTimeoutError,
+  isRateLimitExceededError,
 } from "../utils/error";
 import { RequestStats } from "./statistics";
 
@@ -48,6 +49,7 @@ export class TabbyApiClient extends EventEmitter {
 
   private readonly completionRequestStats = new RequestStats();
   private completionResponseIssue: "highTimeoutRate" | "slowResponseTime" | undefined = undefined;
+  private rateLimitExceeded: boolean = false;
 
   private connectionErrorMessage: string | undefined = undefined;
   private serverHealth: TabbyApiComponents["schemas"]["HealthState"] | undefined = undefined;
@@ -176,6 +178,14 @@ export class TabbyApiClient extends EventEmitter {
     }
   }
 
+  private updateIsRateLimitExceeded(isRateLimitExceeded: boolean) {
+    if (this.rateLimitExceeded != isRateLimitExceeded) {
+      this.logger.debug(`updateIsRateLimitExceeded: ${isRateLimitExceeded}`);
+      this.rateLimitExceeded = isRateLimitExceeded;
+      this.emit("isRateLimitExceededUpdated", isRateLimitExceeded);
+    }
+  }
+
   getCompletionRequestStats(): RequestStats {
     return this.completionRequestStats;
   }
@@ -214,6 +224,10 @@ export class TabbyApiClient extends EventEmitter {
     return !!this.completionResponseIssue;
   }
 
+  isRateLimitExceeded(): boolean {
+    return this.rateLimitExceeded;
+  }
+
   getServerHealth(): TabbyApiComponents["schemas"]["HealthState"] | undefined {
     return this.serverHealth;
   }
@@ -370,6 +384,7 @@ export class TabbyApiClient extends EventEmitter {
       }
       this.logger.trace(`Completion response data: [${requestId}]`, response.data);
       statsData.latency = performance.now() - requestStartedAt;
+      this.updateIsRateLimitExceeded(false);
       return response.data;
     } catch (error) {
       this.updateIsFetchingCompletion(false);
@@ -382,17 +397,24 @@ export class TabbyApiClient extends EventEmitter {
       } else if (isUnauthorizedError(error)) {
         this.logger.debug(`Completion request failed due to unauthorized. [${requestId}]`);
         statsData.notAvailable = true;
+        this.updateIsRateLimitExceeded(false);
         this.connect(); // schedule a reconnection
+      } else if (isRateLimitExceededError(error)) {
+        this.logger.debug(`Completion request failed due to rate limiting. [${requestId}]`);
+        statsData.notAvailable = true;
+        this.updateIsRateLimitExceeded(true);
       } else {
         this.logger.error(`Completion request failed. [${requestId}]`, error);
         statsData.notAvailable = true;
+        this.updateIsRateLimitExceeded(false);
         this.connect(); // schedule a reconnection
       }
       throw error; // rethrow error
     } finally {
       if (!statsData.notAvailable) {
         stats?.addRequestStatsEntry(statsData);
       }
+
       if (!statsData.notAvailable && !statsData.canceled) {
         this.completionRequestStats.add(statsData.latency);
         const statsResult = this.completionRequestStats.stats();

diff --git a/clients/tabby-agent/src/protocol.ts b/clients/tabby-agent/src/protocol.ts
@@ -849,7 +849,8 @@ export type StatusInfo = {
     | "readyForAutoTrigger"
     | "readyForManualTrigger"
     | "fetching"
-    | "completionResponseSlow";
+    | "completionResponseSlow"
+    | "rateLimitExceeded";
   tooltip?: string;
   /**
    * The health information of the server if available.

diff --git a/clients/tabby-agent/src/status.ts b/clients/tabby-agent/src/status.ts
@@ -70,6 +70,9 @@ export class StatusProvider extends EventEmitter implements Feature {
     this.tabbyApiClient.on("hasCompletionResponseTimeIssueUpdated", async () => {
       this.notify();
     });
+    this.tabbyApiClient.on("isRateLimitExceededUpdated", async () => {
+      this.notify();
+    });
 
     this.configurations.on(
       "clientProvidedConfigUpdated",
@@ -203,7 +206,12 @@ export class StatusProvider extends EventEmitter implements Feature {
       case "ready":
         {
           const ignored = this.dataStore.data.statusIgnoredIssues ?? [];
-          if (this.tabbyApiClient.hasCompletionResponseTimeIssue() && !ignored.includes("completionResponseSlow")) {
+          if (this.tabbyApiClient.isRateLimitExceeded()) {
+            statusInfo = { status: "rateLimitExceeded" };
+          } else if (
+            this.tabbyApiClient.hasCompletionResponseTimeIssue() &&
+            !ignored.includes("completionResponseSlow")
+          ) {
             statusInfo = { status: "completionResponseSlow" };
           } else if (this.tabbyApiClient.isFetchingCompletion()) {
             statusInfo = { status: "fetching" };
@@ -264,6 +272,9 @@ export class StatusProvider extends EventEmitter implements Feature {
       case "completionResponseSlow":
         statusInfo.tooltip = "Tabby: Slow Completion Response Detected";
         break;
+      case "rateLimitExceeded":
+        statusInfo.tooltip = "Tabby: Too Many Requests";
+        break;
       default:
         break;
     }

diff --git a/clients/tabby-agent/src/utils/error.ts b/clients/tabby-agent/src/utils/error.ts
@@ -35,6 +35,10 @@ export function isUnauthorizedError(error: any) {
   return error instanceof HttpError && [401, 403].includes(error.status);
 }
 
+export function isRateLimitExceededError(error: any) {
+  return error instanceof HttpError && error.status === 429;
+}
+
 export function errorToString(error: Error) {
   let message = error.message || error.toString();
   if (error.cause instanceof Error) {

diff --git a/clients/vscode/src/StatusBarItem.ts b/clients/vscode/src/StatusBarItem.ts
@@ -93,7 +93,8 @@ export class StatusBarItem {
             this.setTooltip(statusInfo.tooltip);
             break;
           }
-          case "completionResponseSlow": {
+          case "completionResponseSlow":
+          case "rateLimitExceeded": {
             this.setColorWarning();
             this.setIcon(iconWarning);
             this.setTooltip(statusInfo.tooltip);

diff --git a/clients/vscode/src/commands/commandPalette.ts b/clients/vscode/src/commands/commandPalette.ts
@@ -209,6 +209,13 @@ export class CommandPalette {
               },
             };
           }
+          case "rateLimitExceeded": {
+            return {
+              label: `${STATUS_PREFIX}Too Many Requests`,
+              description: "Request limit exceeded",
+              command: "tabby.outputPanel.focus",
+            };
+          }
           default: {
             return {
               label: `${STATUS_PREFIX}Unknown Status`,