feat(vscode): show status of rate limited

TabbyML · Dec 30, 2024 · 3261760 · 3261760
1 parent 652afaf
commit 3261760
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 2 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,
+  isRateLimitedError,
 } 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 rateLimited = false;
 
   private connectionErrorMessage: string | undefined = undefined;
   private serverHealth: TabbyApiComponents["schemas"]["HealthState"] | undefined = undefined;
@@ -214,6 +216,10 @@ export class TabbyApiClient extends EventEmitter {
     return !!this.completionResponseIssue;
   }
 
+  isRateLimited(): boolean {
+    return this.rateLimited;
+  }
+
   getServerHealth(): TabbyApiComponents["schemas"]["HealthState"] | undefined {
     return this.serverHealth;
   }
@@ -334,6 +340,7 @@ export class TabbyApiClient extends EventEmitter {
     }
   }
 
+  // todo
   async fetchCompletion(
     request: TabbyApiComponents["schemas"]["CompletionRequest"],
     signal?: AbortSignal,
@@ -353,6 +360,7 @@ export class TabbyApiClient extends EventEmitter {
       canceled: false,
       timeout: false,
       notAvailable: false,
+      rateLimited: false
     };
 
     try {
@@ -383,6 +391,9 @@ export class TabbyApiClient extends EventEmitter {
         this.logger.debug(`Completion request failed due to unauthorized. [${requestId}]`);
         statsData.notAvailable = true;
         this.connect(); // schedule a reconnection
+      } else if (isRateLimitedError(error)) {
+        this.logger.debug(`Completion request failed due to rate limiting. [${requestId}]`);
+        statsData.rateLimited = true
       } else {
         this.logger.error(`Completion request failed. [${requestId}]`, error);
         statsData.notAvailable = true;

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"
+    | "rateLimited";
   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
@@ -205,6 +205,8 @@ export class StatusProvider extends EventEmitter implements Feature {
           const ignored = this.dataStore.data.statusIgnoredIssues ?? [];
           if (this.tabbyApiClient.hasCompletionResponseTimeIssue() && !ignored.includes("completionResponseSlow")) {
             statusInfo = { status: "completionResponseSlow" };
+          } else if (this.tabbyApiClient.isRateLimited()) {
+            statusInfo = { status: "rateLimited" }
           } else if (this.tabbyApiClient.isFetchingCompletion()) {
             statusInfo = { status: "fetching" };
           } else {
@@ -264,6 +266,9 @@ export class StatusProvider extends EventEmitter implements Feature {
       case "completionResponseSlow":
         statusInfo.tooltip = "Tabby: Slow Completion Response Detected";
         break;
+      case "rateLimited":
+        statusInfo.tooltip = "Tabby: Too many request";
+        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 isRateLimitedError(error: any) {
+  return error instanceof HttpError && [429].includes(error.status);
+}
+
 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 "rateLimited": {
             this.setColorWarning();
             this.setIcon(iconWarning);
             this.setTooltip(statusInfo.tooltip);