From f114469057db18b5d8efac8659554423623df2b8 Mon Sep 17 00:00:00 2001
From: reanon <>
Date: Wed, 23 Jul 2025 23:19:54 +0200
Subject: [PATCH] google is idiotic

---
 .../key-management/google-ai/checker.ts       | 38 ++++++++++++++++++-
 .../key-management/google-ai/provider.ts      | 28 ++++++++++++--
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/src/shared/key-management/google-ai/checker.ts b/src/shared/key-management/google-ai/checker.ts
index 01abd4a..fce82f8 100644
--- a/src/shared/key-management/google-ai/checker.ts
+++ b/src/shared/key-management/google-ai/checker.ts
@@ -15,6 +15,8 @@ const GENERATE_CONTENT_URL =
 const PRO_MODEL_ID = "gemini-2.5-pro";
 const GENERATE_PRO_CONTENT_URL =
   `https://generativelanguage.googleapis.com/v1beta/models/${PRO_MODEL_ID}:generateContent?key=%KEY%`;
+const IMAGEN_BILLING_TEST_URL =
+  "https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate-002:predict?key=%KEY%";
 
 type ListModelsResponse = {
   models: {
@@ -53,6 +55,9 @@ export class GoogleAIKeyChecker extends KeyCheckerBase<GoogleAIKey> {
     // Always test flash model access (existing behaviour)
     await this.testGenerateContent(key);
 
+    // Test if billing is enabled for this key
+    const billingEnabled = await this.testBillingEnabled(key);
+
     // If key claims to support gemini-pro, perform a second layer test with a pro model.
     let effectiveFamilies = [...provisionedModels];
     if (effectiveFamilies.includes("gemini-pro")) {
@@ -66,10 +71,10 @@ export class GoogleAIKeyChecker extends KeyCheckerBase<GoogleAIKey> {
       }
     }
 
-    const updates = { modelFamilies: effectiveFamilies };
+    const updates = { modelFamilies: effectiveFamilies, billingEnabled };
     this.updateKey(key.hash, updates);
     this.log.info(
-      { key: key.hash, models: effectiveFamilies, ids: key.modelIds?.length },
+      { key: key.hash, models: effectiveFamilies, ids: key.modelIds?.length, billingEnabled },
       "Checked key."
     );
   }
@@ -134,6 +139,35 @@ export class GoogleAIKeyChecker extends KeyCheckerBase<GoogleAIKey> {
     }
   }
 
+  private async testBillingEnabled(key: GoogleAIKey): Promise<boolean> {
+    const payload = {
+      instances: [{ prompt: "" }]
+    };
+    try {
+      const response = await axios.post(
+        IMAGEN_BILLING_TEST_URL.replace("%KEY%", key.key),
+        payload,
+        { validateStatus: () => true } // Accept all status codes
+      );
+      
+      if (response.status === 400) {
+        const errorMessage = response.data?.error?.message || "";
+        // If the error message contains the billing requirement, billing is NOT enabled
+        if (errorMessage.includes("Imagen API is only accessible to billed users at this time")) {
+          return false;
+        }
+        // Other 400 errors indicate billing IS enabled (following Python logic)
+        return true;
+      }
+      
+      // For other status codes, assume no billing (conservative approach)
+      return false;
+    } catch (error: any) {
+      // Network errors or other issues - assume no billing
+      return false;
+    }
+  }
+
   protected handleAxiosError(key: GoogleAIKey, error: AxiosError): void {
     if (error.response && GoogleAIKeyChecker.errorIsGoogleAIError(error)) {
       const httpStatus = error.response.status;
diff --git a/src/shared/key-management/google-ai/provider.ts b/src/shared/key-management/google-ai/provider.ts
index c88a218..2d58792 100644
--- a/src/shared/key-management/google-ai/provider.ts
+++ b/src/shared/key-management/google-ai/provider.ts
@@ -32,6 +32,8 @@ export interface GoogleAIKey extends Key {
   isOverQuota?: boolean;
   /** Model families that are over quota and need to be excluded. */
   overQuotaFamilies?: GoogleAIModelFamily[];
+  /** Whether this key has billing enabled (required for preview models). */
+  billingEnabled?: boolean;
 }
 
 /**
@@ -46,6 +48,13 @@ const RATE_LIMIT_LOCKOUT = 2000;
  */
 const KEY_REUSE_DELAY = 500;
 
+/**
+ * Determines if a model is a preview model that requires billing-enabled keys.
+ */
+function isPreviewModel(model: string): boolean {
+  return model.includes("-preview");
+}
+
 export class GoogleAIKeyProvider implements KeyProvider<GoogleAIKey> {
   readonly service = "google-ai";
 
@@ -84,6 +93,7 @@ export class GoogleAIKeyProvider implements KeyProvider<GoogleAIKey> {
         tokenUsage: {}, // Initialize new tokenUsage field
         modelIds: [],
         overQuotaFamilies: [],
+        billingEnabled: false, // Will be determined during key checking
       };
       this.keys.push(newKey);
     }
@@ -103,11 +113,23 @@ export class GoogleAIKeyProvider implements KeyProvider<GoogleAIKey> {
 
   public get(model: string) {
     const neededFamily = getGoogleAIModelFamily(model);
-    const availableKeys = this.keys.filter(
+    let availableKeys = this.keys.filter(
       (k) => !k.isDisabled && k.modelFamilies.includes(neededFamily)
     );
-    if (availableKeys.length === 0) {
-      throw new PaymentRequiredError("No Google AI keys available");
+    
+    // For preview models, only use billing-enabled keys
+    if (isPreviewModel(model)) {
+      availableKeys = availableKeys.filter((k) => k.billingEnabled === true);
+      if (availableKeys.length === 0) {
+        throw new PaymentRequiredError(
+          "No billing-enabled Google AI keys available for preview models"
+        );
+      }
+    } else {
+      // For standard models, use any available key
+      if (availableKeys.length === 0) {
+        throw new PaymentRequiredError("No Google AI keys available");
+      }
     }
 
     const keysByPriority = prioritizeKeys(availableKeys);