Update 3 files

- /src/shared/key-management/xai/checker.ts
- /src/shared/key-management/key-pool.ts
- /src/service-info.ts

src/service-info.ts

@@ -481,6 +481,9 @@ function getInfoForFamily(family: ModelFamily): BaseFamilyInfo {
         break;
       case "deepseek":
         info.overQuotaKeys = familyStats.get(`${family}__overQuota`) || 0;
+      case "xai":
+        info.overQuotaKeys = familyStats.get(`${family}__overQuota`) || 0;
+        break;
     }
   }
 
@@ -505,4 +508,4 @@ function getQueueInformation(partition: ModelFamily) {
     proomptersInQueue: getQueueLength(partition),
     estimatedQueueTime: waitMs > 2000 ? waitTime : "no wait",
   };
-}
+}

src/shared/key-management/key-pool.ts

@@ -16,7 +16,7 @@ import { MistralAIKeyProvider } from "./mistral-ai/provider";
 import { DeepseekKeyProvider } from "./deepseek/provider";
 import { XaiKeyProvider } from "./xai/provider";
 
-type AllowedPartial = OpenAIKeyUpdate | AnthropicKeyUpdate | Partial<GcpKey>;
+type AllowedPartial = OpenAIKeyUpdate | AnthropicKeyUpdate | Partial<GcpKey> | Partial<XaiKey>;
 
 export class KeyPool {
   private keyProviders: KeyProvider[] = [];
@@ -74,7 +74,8 @@ export class KeyPool {
     if (
       service instanceof OpenAIKeyProvider ||
       service instanceof AnthropicKeyProvider ||
-      service instanceof DeepseekKeyProvider
+      service instanceof DeepseekKeyProvider ||
+      service instanceof XaiKeyProvider
     ) {
       service.update(key.hash, { isOverQuota: reason === "quota" });
     }
@@ -197,4 +198,4 @@ export class KeyPool {
     );
     this.recheckJobs.openai = job;
   }
-}
+}

src/shared/key-management/xai/checker.ts

@@ -7,23 +7,36 @@ const CHECK_TIMEOUT = 10000;
 export class XaiKeyChecker {
   private log = logger.child({ module: "key-checker", service: "xai" });
 
-  constructor(private readonly update: (hash: string, key: Partial<XaiKey>) => void) {}
+  constructor(private readonly update: (hash: string, key: Partial<XaiKey>) => void) {
+    this.log.info("XaiKeyChecker initialized");
+  }
 
   public async checkKey(key: XaiKey): Promise<void> {
+    this.log.info({ hash: key.hash }, "Starting key validation check");
     try {
       const result = await this.validateKey(key);
       this.handleCheckResult(key, result);
     } catch (error) {
-      this.log.warn(
-        { error, hash: key.hash },
-        "Failed to check key status"
-      );
+      if (error instanceof Error) {
+        this.log.warn(
+          { error: error.message, stack: error.stack, hash: key.hash },
+          "Failed to check key status"
+        );
+      } else {
+        this.log.warn(
+          { error, hash: key.hash },
+          "Failed to check key status with unknown error"
+        );
+      }
     }
   }
 
   private async validateKey(key: XaiKey): Promise<"valid" | "invalid" | "quota"> {
     const controller = new AbortController();
-    const timeout = setTimeout(() => controller.abort(), CHECK_TIMEOUT);
+    const timeout = setTimeout(() => {
+      controller.abort();
+      this.log.warn({ hash: key.hash }, "Key validation timed out after " + CHECK_TIMEOUT + "ms");
+    }, CHECK_TIMEOUT);
 
     try {
       const response = await fetch("https://api.x.ai/v1/chat/completions", {
@@ -40,32 +53,27 @@ export class XaiKeyChecker {
         signal: controller.signal,
       });
 
-      const rateLimit = {
-        limit: parseInt(response.headers.get("x-ratelimit-limit") || "200"),
-        remaining: parseInt(response.headers.get("x-ratelimit-remaining") || "199"),
-      };
-
       switch (response.status) {
         case 400:
-          this.log.debug(
-            { key: key.hash, rateLimit },
-            "Key check successful, updating rate limit info"
-          );
           return "valid";
         case 401:
           return "invalid";
-        case 402:
-          return "quota";
+        case 403:
+          return "invalid";
         case 429:
-          this.log.warn({ key: key.hash }, "Key is rate limited");
-          return "valid";
+          return "quota";
         default:
           this.log.warn(
             { status: response.status, hash: key.hash },
             "Unexpected status code while checking key"
           );
-          return "valid";
+          return "invalid";
       }
+    } catch (error) {
+      if (error instanceof Error && error.name === 'AbortError') {
+        this.log.warn({ hash: key.hash }, "Key validation aborted");
+      }
+      throw error;
     } finally {
       clearTimeout(timeout);
     }
@@ -77,13 +85,14 @@ export class XaiKeyChecker {
   ): void {
     switch (result) {
       case "valid":
+        this.log.info({ hash: key.hash }, "Key is valid and enabled");
         this.update(key.hash, {
           isDisabled: false,
           lastChecked: Date.now(),
         });
         break;
       case "invalid":
-        this.log.warn({ hash: key.hash }, "Key is invalid");
+        this.log.warn({ hash: key.hash }, "Key is invalid, marking as revoked");
         this.update(key.hash, {
           isDisabled: true,
           isRevoked: true,
@@ -91,7 +100,7 @@ export class XaiKeyChecker {
         });
         break;
       case "quota":
-        this.log.warn({ hash: key.hash }, "Key has exceeded its quota");
+        this.log.warn({ hash: key.hash }, "Key has exceeded its quota, disabling");
         this.update(key.hash, {
           isDisabled: true,
           isOverQuota: true,
@@ -102,4 +111,4 @@ export class XaiKeyChecker {
         assertNever(result);
     }
   }
-}
+}