feat(cliproxy): add proactive token refresh for Gemini quota (match AGY pattern)

- Add proactive refresh: refresh token 5min before expiry (not just after)
- Add retry on 401: if API returns 401, refresh token and retry once
- Extract fetchWithAuthData helper to reduce code duplication
- Update UI to show better error messages (no more "via CLI" text)
- Falls back gracefully if proactive refresh fails but token not yet expired
This commit is contained in:
kaitranntt
2026-01-29 23:09:47 -05:00
parent 7947a7ac89
commit 606bb72723
3 changed files with 103 additions and 44 deletions
+91 -37
View File
@@ -10,6 +10,7 @@ import * as path from 'node:path';
import { getAuthDir } from './config-generator';
import { getProviderAccounts, getPausedDir } from './account-manager';
import { sanitizeEmail, isTokenExpired } from './auth-utils';
import { refreshGeminiToken } from './auth/gemini-token-refresh';
import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types';
/** Google Cloud Code API endpoints */
@@ -322,46 +323,14 @@ function buildGeminiCliBuckets(rawBuckets: RawGeminiCliBucket[]): GeminiCliBucke
}
/**
* Fetch quota for a single Gemini CLI account
*
* @param accountId - Account identifier (email)
* @param verbose - Show detailed diagnostics
* @returns Quota result with buckets and percentages
* Internal helper: Fetch quota with validated auth data
* Extracted to support auto-refresh retry logic
*/
export async function fetchGeminiCliQuota(
async function fetchWithAuthData(
authData: GeminiCliAuthData,
accountId: string,
verbose = false
verbose: boolean
): Promise<GeminiCliQuotaResult> {
if (verbose) console.error(`[i] Fetching Gemini CLI quota for ${accountId}...`);
const authData = readGeminiCliAuthData(accountId);
if (!authData) {
const error = 'Auth file not found for Gemini account';
if (verbose) console.error(`[!] Error: ${error}`);
return {
success: false,
buckets: [],
projectId: null,
lastUpdated: Date.now(),
error,
accountId,
};
}
if (authData.isExpired) {
const error = 'Token expired - re-authenticate with ccs cliproxy auth gemini';
if (verbose) console.error(`[!] Error: ${error}`);
return {
success: false,
buckets: [],
projectId: null,
lastUpdated: Date.now(),
error,
accountId,
needsReauth: true,
};
}
if (!authData.projectId) {
const error = 'Cannot resolve project ID from auth file';
if (verbose) console.error(`[!] Error: ${error}`);
@@ -474,6 +443,91 @@ export async function fetchGeminiCliQuota(
}
}
/**
* Fetch quota for a single Gemini CLI account
*
* @param accountId - Account identifier (email)
* @param verbose - Show detailed diagnostics
* @returns Quota result with buckets and percentages
*/
export async function fetchGeminiCliQuota(
accountId: string,
verbose = false
): Promise<GeminiCliQuotaResult> {
if (verbose) console.error(`[i] Fetching Gemini CLI quota for ${accountId}...`);
let authData = readGeminiCliAuthData(accountId);
if (!authData) {
const error = 'Auth file not found for Gemini account';
if (verbose) console.error(`[!] Error: ${error}`);
return {
success: false,
buckets: [],
projectId: null,
lastUpdated: Date.now(),
error,
accountId,
};
}
// Proactive refresh: refresh if expired OR expiring within 5 minutes
const REFRESH_LEAD_TIME_MS = 5 * 60 * 1000;
const shouldRefresh =
authData.isExpired ||
!authData.expiresAt ||
new Date(authData.expiresAt).getTime() - Date.now() < REFRESH_LEAD_TIME_MS;
if (shouldRefresh) {
if (verbose)
console.error(
authData.isExpired
? '[i] Token expired, refreshing...'
: '[i] Token expiring soon, proactive refresh...'
);
const refreshResult = await refreshGeminiToken();
if (refreshResult.success) {
if (verbose) console.error('[i] Token refreshed successfully');
// Re-read auth data after successful refresh
const refreshedAuthData = readGeminiCliAuthData(accountId);
if (refreshedAuthData) {
authData = refreshedAuthData;
}
} else if (authData.isExpired) {
// Only fail if token is actually expired (not just expiring soon)
const error = refreshResult.error || 'Token refresh failed';
if (verbose) console.error(`[!] Refresh failed: ${error}`);
return {
success: false,
buckets: [],
projectId: null,
lastUpdated: Date.now(),
error,
accountId,
needsReauth: true,
};
}
// If proactive refresh fails but token isn't expired yet, continue with existing token
}
// First attempt with current token
const result = await fetchWithAuthData(authData, accountId, verbose);
// If 401 error and we haven't refreshed yet, try refresh and retry
if (result.needsReauth && result.error?.includes('expired')) {
if (verbose) console.error('[i] Got 401, attempting refresh and retry...');
const refreshResult = await refreshGeminiToken();
if (refreshResult.success) {
const refreshedAuthData = readGeminiCliAuthData(accountId);
if (refreshedAuthData) {
return await fetchWithAuthData(refreshedAuthData, accountId, verbose);
}
}
}
return result;
}
/**
* Fetch quota for all Gemini CLI accounts
*
@@ -258,8 +258,12 @@ export function AccountCard({
<span>Reauth needed</span>
</div>
</TooltipTrigger>
<TooltipContent side="top" className="max-w-[180px]">
<p className="text-xs">Token expired. Re-authenticate via CLI.</p>
<TooltipContent side="top" className="max-w-[200px]">
<p className="text-xs">
{quota.error?.includes('No refresh token')
? 'Remove and re-add account'
: quota.error || 'Auto-refresh failed'}
</p>
</TooltipContent>
</Tooltip>
</TooltipProvider>
@@ -366,12 +366,13 @@ export function AccountItem({
</Badge>
</div>
</TooltipTrigger>
<TooltipContent side="bottom" className="max-w-[200px]">
<TooltipContent side="bottom" className="max-w-[220px]">
<p className="text-xs">
Token expired. Re-authenticate via CLI:{' '}
<code className="font-mono text-[10px]">
ccs cliproxy auth {account.provider}
</code>
{quota.error?.includes('No refresh token')
? 'No refresh token available. Remove and re-add account to fix.'
: quota.error?.includes('refresh') || quota.error?.includes('Invalid')
? `Auto-refresh failed: ${quota.error}`
: `Token issue: ${quota.error || 'Re-authenticate required'}`}
</p>
</TooltipContent>
</Tooltip>