fix(quota): avoid blocking gemini cli quota refresh on code assist

2026-06-16 21:03:58 +08:00 · 2026-03-17 00:46:32 +08:00
parent e83f99ba95
commit 6d91cd0c06
7 changed files with 350 additions and 53 deletions
@@ -20,13 +20,17 @@ import type {
  CodexUsageWindow,
  CodexQuotaWindow,
  CodexUsagePayload,
+  GeminiCliCodeAssistPayload,
+  GeminiCliCredits,
  GeminiCliParsedBucket,
  GeminiCliQuotaBucketState,
  GeminiCliQuotaState,
+  GeminiCliUserTier,
  KimiQuotaRow,
  KimiQuotaState,
 } from '@/types';
 import { apiCallApi, authFilesApi, getApiCallErrorMessage } from '@/services/api';
+import { useQuotaStore } from '@/stores';
 import {
  ANTIGRAVITY_QUOTA_URLS,
  ANTIGRAVITY_REQUEST_HEADERS,
@@ -37,6 +41,7 @@ import {
  CODEX_USAGE_URL,
  CODEX_REQUEST_HEADERS,
  GEMINI_CLI_QUOTA_URL,
+  GEMINI_CLI_CODE_ASSIST_URL,
  GEMINI_CLI_REQUEST_HEADERS,
  KIMI_USAGE_URL,
  KIMI_REQUEST_HEADERS,
@@ -49,6 +54,7 @@ import {
  parseClaudeUsagePayload,
  parseCodexUsagePayload,
  parseGeminiCliQuotaPayload,
+  parseGeminiCliCodeAssistPayload,
  parseKimiUsagePayload,
  resolveCodexChatgptAccountId,
  resolveCodexPlanType,
@@ -78,6 +84,11 @@ type QuotaUpdater<T> = T | ((prev: T) => T);
 type QuotaType = 'antigravity' | 'claude' | 'codex' | 'gemini-cli' | 'kimi';

 const DEFAULT_ANTIGRAVITY_PROJECT_ID = 'bamboo-precept-lgxtn';
+const geminiCliSupplementaryRequestIds = new Map<string, number>();
+const geminiCliSupplementaryCache = new Map<
+  string,
+  { requestId: number; tierLabel: string | null; creditBalance: number | null }
+>();

 export interface QuotaStore {
  antigravityQuota: Record<string, AntigravityQuotaState>;
@@ -427,10 +438,161 @@ const fetchCodexQuota = async (
  return { planType: planTypeFromUsage ?? planTypeFromFile, windows };
 };

+const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI';
+
+const GEMINI_CLI_TIER_LABELS: Record<string, string> = {
+  'free-tier': 'tier_free',
+  'legacy-tier': 'tier_legacy',
+  'standard-tier': 'tier_standard',
+};
+
+const resolveGeminiCliTierLabel = (
+  payload: GeminiCliCodeAssistPayload | null,
+  t: TFunction
+): string | null => {
+  if (!payload) return null;
+  const currentTier: GeminiCliUserTier | null | undefined =
+    payload.currentTier ?? payload.current_tier;
+  const paidTier: GeminiCliUserTier | null | undefined =
+    payload.paidTier ?? payload.paid_tier;
+  const tierId = normalizeStringValue(paidTier?.id) ?? normalizeStringValue(currentTier?.id);
+  if (!tierId) return null;
+  const labelKey = GEMINI_CLI_TIER_LABELS[tierId];
+  return labelKey ? t(`gemini_cli_quota.${labelKey}`) : tierId;
+};
+
+const resolveGeminiCliCreditBalance = (
+  payload: GeminiCliCodeAssistPayload | null
+): number | null => {
+  if (!payload) return null;
+  const paidTier: GeminiCliUserTier | null | undefined =
+    payload.paidTier ?? payload.paid_tier;
+  const currentTier: GeminiCliUserTier | null | undefined =
+    payload.currentTier ?? payload.current_tier;
+  const tier = paidTier ?? currentTier;
+  if (!tier) return null;
+  const credits: GeminiCliCredits[] =
+    tier.availableCredits ?? tier.available_credits ?? [];
+  let total = 0;
+  let found = false;
+  for (const credit of credits) {
+    const creditType = normalizeStringValue(credit.creditType ?? credit.credit_type);
+    if (creditType !== GEMINI_CLI_G1_CREDIT_TYPE) continue;
+    const amount = normalizeNumberValue(credit.creditAmount ?? credit.credit_amount);
+    if (amount !== null) {
+      total += amount;
+      found = true;
+    }
+  }
+  return found ? total : null;
+};
+
+const fetchGeminiCliCodeAssist = async (
+  authIndex: string,
+  projectId: string,
+  t: TFunction
+): Promise<{ tierLabel: string | null; creditBalance: number | null }> => {
+  try {
+    const result = await apiCallApi.request({
+      authIndex,
+      method: 'POST',
+      url: GEMINI_CLI_CODE_ASSIST_URL,
+      header: { ...GEMINI_CLI_REQUEST_HEADERS },
+      data: JSON.stringify({
+        cloudaicompanionProject: projectId,
+        metadata: {
+          ideType: 'IDE_UNSPECIFIED',
+          platform: 'PLATFORM_UNSPECIFIED',
+          pluginType: 'GEMINI',
+          duetProject: projectId,
+        },
+      }),
+    });
+
+    if (result.statusCode < 200 || result.statusCode >= 300) {
+      return { tierLabel: null, creditBalance: null };
+    }
+
+    const payload = parseGeminiCliCodeAssistPayload(result.body ?? result.bodyText);
+    return {
+      tierLabel: resolveGeminiCliTierLabel(payload, t),
+      creditBalance: resolveGeminiCliCreditBalance(payload),
+    };
+  } catch {
+    return { tierLabel: null, creditBalance: null };
+  }
+};
+
+const readGeminiCliSupplementarySnapshot = (
+  fileName: string,
+  requestId: number
+): { tierLabel: string | null; creditBalance: number | null } => {
+  const cached = geminiCliSupplementaryCache.get(fileName);
+  if (!cached || cached.requestId !== requestId) {
+    return { tierLabel: null, creditBalance: null };
+  }
+
+  return {
+    tierLabel: cached.tierLabel,
+    creditBalance: cached.creditBalance,
+  };
+};
+
+const scheduleGeminiCliSupplementaryRefresh = (
+  fileName: string,
+  authIndex: string,
+  projectId: string,
+  t: TFunction
+): number => {
+  const requestId = (geminiCliSupplementaryRequestIds.get(fileName) ?? 0) + 1;
+  geminiCliSupplementaryRequestIds.set(fileName, requestId);
+  geminiCliSupplementaryCache.delete(fileName);
+
+  void (async () => {
+    const supplementary = await fetchGeminiCliCodeAssist(authIndex, projectId, t);
+    if (geminiCliSupplementaryRequestIds.get(fileName) !== requestId) {
+      return;
+    }
+
+    geminiCliSupplementaryCache.set(fileName, { requestId, ...supplementary });
+
+    useQuotaStore.getState().setGeminiCliQuota((prev) => {
+      const current = prev[fileName];
+      if (!current || current.status !== 'success') {
+        return prev;
+      }
+
+      if (
+        current.tierLabel === supplementary.tierLabel &&
+        current.creditBalance === supplementary.creditBalance
+      ) {
+        return prev;
+      }
+
+      return {
+        ...prev,
+        [fileName]: {
+          ...current,
+          tierLabel: supplementary.tierLabel,
+          creditBalance: supplementary.creditBalance,
+        },
+      };
+    });
+  })();
+
+  return requestId;
+};
+
 const fetchGeminiCliQuota = async (
  file: AuthFileItem,
  t: TFunction
-): Promise<GeminiCliQuotaBucketState[]> => {
+): Promise<{
+  fileName: string;
+  supplementaryRequestId: number;
+  buckets: GeminiCliQuotaBucketState[];
+  tierLabel: string | null;
+  creditBalance: number | null;
+}> => {
  const rawAuthIndex = file['auth_index'] ?? file.authIndex;
  const authIndex = normalizeAuthIndex(rawAuthIndex);
  if (!authIndex) {
@@ -442,21 +604,19 @@ const fetchGeminiCliQuota = async (
    throw new Error(t('gemini_cli_quota.missing_project_id'));
  }

-  const result = await apiCallApi.request({
+  const quotaResponse = await apiCallApi.request({
    authIndex,
    method: 'POST',
    url: GEMINI_CLI_QUOTA_URL,
    header: { ...GEMINI_CLI_REQUEST_HEADERS },
    data: JSON.stringify({ project: projectId }),
  });
-
-  if (result.statusCode < 200 || result.statusCode >= 300) {
-    throw createStatusError(getApiCallErrorMessage(result), result.statusCode);
+  if (quotaResponse.statusCode < 200 || quotaResponse.statusCode >= 300) {
+    throw createStatusError(getApiCallErrorMessage(quotaResponse), quotaResponse.statusCode);
  }

-  const payload = parseGeminiCliQuotaPayload(result.body ?? result.bodyText);
+  const payload = parseGeminiCliQuotaPayload(quotaResponse.body ?? quotaResponse.bodyText);
  const buckets = Array.isArray(payload?.buckets) ? payload?.buckets : [];
-  if (buckets.length === 0) return [];

  const parsedBuckets = buckets
    .map((bucket) => {
@@ -487,7 +647,25 @@ const fetchGeminiCliQuota = async (
    })
    .filter((bucket): bucket is GeminiCliParsedBucket => bucket !== null);

-  return buildGeminiCliQuotaBuckets(parsedBuckets);
+  const builtBuckets = buildGeminiCliQuotaBuckets(parsedBuckets);
+  const supplementaryRequestId = scheduleGeminiCliSupplementaryRefresh(
+    file.name,
+    authIndex,
+    projectId,
+    t
+  );
+  const supplementarySnapshot = readGeminiCliSupplementarySnapshot(
+    file.name,
+    supplementaryRequestId
+  );
+
+  return {
+    fileName: file.name,
+    supplementaryRequestId,
+    buckets: builtBuckets,
+    tierLabel: supplementarySnapshot.tierLabel,
+    creditBalance: supplementarySnapshot.creditBalance,
+  };
 };

 const renderAntigravityItems = (
@@ -605,50 +783,86 @@ const renderGeminiCliItems = (
  helpers: QuotaRenderHelpers
 ): ReactNode => {
  const { styles: styleMap, QuotaProgressBar } = helpers;
-  const { createElement: h } = React;
+  const { createElement: h, Fragment } = React;
  const buckets = quota.buckets ?? [];
+  const tierLabel = quota.tierLabel ?? null;
+  const creditBalance = quota.creditBalance ?? null;
+  const nodes: ReactNode[] = [];

-  if (buckets.length === 0) {
-    return h('div', { className: styleMap.quotaMessage }, t('gemini_cli_quota.empty_buckets'));
-  }
-
-  return buckets.map((bucket) => {
-    const fraction = bucket.remainingFraction;
-    const clamped = fraction === null ? null : Math.max(0, Math.min(1, fraction));
-    const percent = clamped === null ? null : Math.round(clamped * 100);
-    const percentLabel = percent === null ? '--' : `${percent}%`;
-    const remainingAmountLabel =
-      bucket.remainingAmount === null || bucket.remainingAmount === undefined
-        ? null
-        : t('gemini_cli_quota.remaining_amount', {
-            count: bucket.remainingAmount,
-          });
-    const titleBase =
-      bucket.modelIds && bucket.modelIds.length > 0 ? bucket.modelIds.join(', ') : bucket.label;
-    const title = bucket.tokenType ? `${titleBase} (${bucket.tokenType})` : titleBase;
-
-    const resetLabel = formatQuotaResetTime(bucket.resetTime);
-
-    return h(
-      'div',
-      { key: bucket.id, className: styleMap.quotaRow },
+  if (tierLabel) {
+    nodes.push(
      h(
        'div',
-        { className: styleMap.quotaRowHeader },
-        h('span', { className: styleMap.quotaModel, title }, bucket.label),
+        { key: 'tier', className: styleMap.codexPlan },
+        h('span', { className: styleMap.codexPlanLabel }, t('gemini_cli_quota.tier_label')),
+        h('span', { className: styleMap.codexPlanValue }, tierLabel)
+      )
+    );
+  }
+
+  if (creditBalance !== null) {
+    nodes.push(
+      h(
+        'div',
+        { key: 'credits', className: styleMap.codexPlan },
+        h('span', { className: styleMap.codexPlanLabel }, t('gemini_cli_quota.credit_label')),
+        h(
+          'span',
+          { className: styleMap.codexPlanValue },
+          t('gemini_cli_quota.credit_amount', { count: creditBalance })
+        )
+      )
+    );
+  }
+
+  if (buckets.length === 0) {
+    nodes.push(
+      h('div', { key: 'empty', className: styleMap.quotaMessage }, t('gemini_cli_quota.empty_buckets'))
+    );
+    return h(Fragment, null, ...nodes);
+  }
+
+  nodes.push(
+    ...buckets.map((bucket) => {
+      const fraction = bucket.remainingFraction;
+      const clamped = fraction === null ? null : Math.max(0, Math.min(1, fraction));
+      const percent = clamped === null ? null : Math.round(clamped * 100);
+      const percentLabel = percent === null ? '--' : `${percent}%`;
+      const remainingAmountLabel =
+        bucket.remainingAmount === null || bucket.remainingAmount === undefined
+          ? null
+          : t('gemini_cli_quota.remaining_amount', {
+              count: bucket.remainingAmount,
+            });
+      const titleBase =
+        bucket.modelIds && bucket.modelIds.length > 0 ? bucket.modelIds.join(', ') : bucket.label;
+      const title = bucket.tokenType ? `${titleBase} (${bucket.tokenType})` : titleBase;
+
+      const resetLabel = formatQuotaResetTime(bucket.resetTime);
+
+      return h(
+        'div',
+        { key: bucket.id, className: styleMap.quotaRow },
        h(
          'div',
-          { className: styleMap.quotaMeta },
-          h('span', { className: styleMap.quotaPercent }, percentLabel),
-          remainingAmountLabel
-            ? h('span', { className: styleMap.quotaAmount }, remainingAmountLabel)
-            : null,
-          h('span', { className: styleMap.quotaReset }, resetLabel)
-        )
-      ),
-      h(QuotaProgressBar, { percent, highThreshold: 60, mediumThreshold: 20 })
-    );
-  });
+          { className: styleMap.quotaRowHeader },
+          h('span', { className: styleMap.quotaModel, title }, bucket.label),
+          h(
+            'div',
+            { className: styleMap.quotaMeta },
+            h('span', { className: styleMap.quotaPercent }, percentLabel),
+            remainingAmountLabel
+              ? h('span', { className: styleMap.quotaAmount }, remainingAmountLabel)
+              : null,
+            h('span', { className: styleMap.quotaReset }, resetLabel)
+          )
+        ),
+        h(QuotaProgressBar, { percent, highThreshold: 60, mediumThreshold: 20 })
+      );
+    })
+  );
+
+  return h(Fragment, null, ...nodes);
 };

 const buildClaudeQuotaWindows = (
@@ -927,7 +1141,16 @@ export const CODEX_CONFIG: QuotaConfig<
  renderQuotaItems: renderCodexItems,
 };

-export const GEMINI_CLI_CONFIG: QuotaConfig<GeminiCliQuotaState, GeminiCliQuotaBucketState[]> = {
+export const GEMINI_CLI_CONFIG: QuotaConfig<
+  GeminiCliQuotaState,
+  {
+    fileName: string;
+    supplementaryRequestId: number;
+    buckets: GeminiCliQuotaBucketState[];
+    tierLabel: string | null;
+    creditBalance: number | null;
+  }
+> = {
  type: 'gemini-cli',
  i18nPrefix: 'gemini_cli_quota',
  cardIdleMessageKey: 'quota_management.card_idle_hint',
@@ -936,8 +1159,20 @@ export const GEMINI_CLI_CONFIG: QuotaConfig<GeminiCliQuotaState, GeminiCliQuotaB
  fetchQuota: fetchGeminiCliQuota,
  storeSelector: (state) => state.geminiCliQuota,
  storeSetter: 'setGeminiCliQuota',
-  buildLoadingState: () => ({ status: 'loading', buckets: [] }),
-  buildSuccessState: (buckets) => ({ status: 'success', buckets }),
+  buildLoadingState: () => ({ status: 'loading', buckets: [], tierLabel: null, creditBalance: null }),
+  buildSuccessState: (data) => {
+    const supplementarySnapshot = readGeminiCliSupplementarySnapshot(
+      data.fileName,
+      data.supplementaryRequestId
+    );
+
+    return {
+      status: 'success',
+      buckets: data.buckets,
+      tierLabel: supplementarySnapshot.tierLabel ?? data.tierLabel,
+      creditBalance: supplementarySnapshot.creditBalance ?? data.creditBalance,
+    };
+  },
  buildErrorState: (message, status) => ({
    status: 'error',
    buckets: [],
@@ -644,7 +644,13 @@
    "empty_buckets": "No quota data available",
    "refresh_button": "Refresh Quota",
    "fetch_all": "Fetch All",
-    "remaining_amount": "Remaining {{count}}"
+    "remaining_amount": "Remaining {{count}}",
+    "tier_label": "Tier",
+    "tier_free": "Free",
+    "tier_legacy": "Legacy",
+    "tier_standard": "Standard",
+    "credit_label": "Google One AI Credits",
+    "credit_amount": "{{count}} credits"
  },
  "kimi_quota": {
    "title": "Kimi Quota",
@@ -647,7 +647,13 @@
    "empty_buckets": "Данные по квоте отсутствуют",
    "refresh_button": "Обновить квоту",
    "fetch_all": "Получить все",
-    "remaining_amount": "Осталось {{count}}"
+    "remaining_amount": "Осталось {{count}}",
+    "tier_label": "Уровень",
+    "tier_free": "Бесплатный",
+    "tier_legacy": "Устаревший",
+    "tier_standard": "Стандартный",
+    "credit_label": "Google One AI кредиты",
+    "credit_amount": "{{count}} кредитов"
  },
  "kimi_quota": {
    "title": "Квота Kimi",
@@ -644,7 +644,13 @@
    "empty_buckets": "暂无额度数据",
    "refresh_button": "刷新额度",
    "fetch_all": "获取全部",
-    "remaining_amount": "剩余 {{count}}"
+    "remaining_amount": "剩余 {{count}}",
+    "tier_label": "层级",
+    "tier_free": "免费版",
+    "tier_legacy": "旧版",
+    "tier_standard": "标准版",
+    "credit_label": "Google One AI 积分",
+    "credit_amount": "{{count}} 积分"
  },
  "kimi_quota": {
    "title": "Kimi 额度",
@@ -25,6 +25,28 @@ export interface GeminiCliQuotaPayload {
  buckets?: GeminiCliQuotaBucket[];
 }

+export interface GeminiCliCredits {
+  creditType?: string;
+  credit_type?: string;
+  creditAmount?: string | number;
+  credit_amount?: string | number;
+}
+
+export interface GeminiCliUserTier {
+  id?: string;
+  name?: string;
+  description?: string;
+  availableCredits?: GeminiCliCredits[];
+  available_credits?: GeminiCliCredits[];
+}
+
+export interface GeminiCliCodeAssistPayload {
+  currentTier?: GeminiCliUserTier | null;
+  current_tier?: GeminiCliUserTier | null;
+  paidTier?: GeminiCliUserTier | null;
+  paid_tier?: GeminiCliUserTier | null;
+}
+
 export interface AntigravityQuotaInfo {
  displayName?: string;
  quotaInfo?: {
@@ -200,6 +222,8 @@ export interface GeminiCliQuotaBucketState {
 export interface GeminiCliQuotaState {
  status: 'idle' | 'loading' | 'success' | 'error';
  buckets: GeminiCliQuotaBucketState[];
+  tierLabel?: string | null;
+  creditBalance?: number | null;
  error?: string;
  errorStatus?: number;
 }
@@ -117,6 +117,9 @@ export const ANTIGRAVITY_QUOTA_GROUPS: AntigravityQuotaGroupDefinition[] = [
 export const GEMINI_CLI_QUOTA_URL =
  'https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota';

+export const GEMINI_CLI_CODE_ASSIST_URL =
+  'https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist';
+
 export const GEMINI_CLI_REQUEST_HEADERS = {
  Authorization: 'Bearer $TOKEN$',
  'Content-Type': 'application/json',
@@ -2,7 +2,7 @@
 * Normalization and parsing functions for quota data.
 */

-import type { ClaudeUsagePayload, CodexUsagePayload, GeminiCliQuotaPayload, KimiUsagePayload } from '@/types';
+import type { ClaudeUsagePayload, CodexUsagePayload, GeminiCliCodeAssistPayload, GeminiCliQuotaPayload, KimiUsagePayload } from '@/types';
 import { normalizeAuthIndex } from '@/utils/usage';

 const GEMINI_CLI_MODEL_SUFFIX = '_vertex';
@@ -191,6 +191,23 @@ export function parseGeminiCliQuotaPayload(payload: unknown): GeminiCliQuotaPayl
  return null;
 }

+export function parseGeminiCliCodeAssistPayload(payload: unknown): GeminiCliCodeAssistPayload | null {
+  if (payload === undefined || payload === null) return null;
+  if (typeof payload === 'string') {
+    const trimmed = payload.trim();
+    if (!trimmed) return null;
+    try {
+      return JSON.parse(trimmed) as GeminiCliCodeAssistPayload;
+    } catch {
+      return null;
+    }
+  }
+  if (typeof payload === 'object') {
+    return payload as GeminiCliCodeAssistPayload;
+  }
+  return null;
+}
+
 export function parseKimiUsagePayload(payload: unknown): KimiUsagePayload | null {
  if (payload === undefined || payload === null) return null;
  if (typeof payload === 'string') {