mirror of
https://github.com/tiennm99/ccs.git
synced 2026-07-16 06:16:37 +00:00
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:
@@ -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>
|
||||
|
||||
Reference in New Issue
Block a user