mirror of
https://github.com/tiennm99/ccs.git
synced 2026-10-13 04:14:57 +00:00
fix(cliproxy): avoid duplicate gemini management retries
This commit is contained in:
1 parent
9dc6374851
commit
637a591dca
2 files changed
+109
-8
No files matched your search
@@ -114,6 +114,10 @@ interface ManagedResponse {
|
|||||||
viaManagement: boolean;
|
viaManagement: boolean;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function getRemainingTimeoutMs(deadlineMs: number): number {
|
||||||
|
return Math.max(1, deadlineMs - Date.now());
|
||||||
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Extract project ID from account field
|
* Extract project ID from account field
|
||||||
* Input: "user@example.com (cloudaicompanion-abc-123)"
|
* Input: "user@example.com (cloudaicompanion-abc-123)"
|
||||||
@@ -235,10 +239,13 @@ function isGeminiAuthFileForAccount(file: ManagementAuthFile, accountId: string)
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
async function findManagedGeminiAuthIndex(accountId: string): Promise<string | number | null> {
|
async function findManagedGeminiAuthIndex(
|
||||||
|
accountId: string,
|
||||||
|
timeoutMs: number
|
||||||
|
): Promise<string | number | null> {
|
||||||
const target = getProxyTarget();
|
const target = getProxyTarget();
|
||||||
const controller = new AbortController();
|
const controller = new AbortController();
|
||||||
const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS);
|
const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), {
|
const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), {
|
||||||
@@ -263,16 +270,17 @@ async function findManagedGeminiAuthIndex(accountId: string): Promise<string | n
|
|||||||
async function performManagedGeminiRequest(
|
async function performManagedGeminiRequest(
|
||||||
accountId: string,
|
accountId: string,
|
||||||
url: string,
|
url: string,
|
||||||
body: string
|
body: string,
|
||||||
|
timeoutMs: number
|
||||||
): Promise<ManagedResponse | null> {
|
): Promise<ManagedResponse | null> {
|
||||||
const authIndex = await findManagedGeminiAuthIndex(accountId);
|
const authIndex = await findManagedGeminiAuthIndex(accountId, timeoutMs);
|
||||||
if (authIndex === null || authIndex === undefined) {
|
if (authIndex === null || authIndex === undefined) {
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
const target = getProxyTarget();
|
const target = getProxyTarget();
|
||||||
const controller = new AbortController();
|
const controller = new AbortController();
|
||||||
const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS);
|
const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), {
|
const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), {
|
||||||
@@ -319,15 +327,24 @@ async function performGeminiCliRequest(
|
|||||||
body: string,
|
body: string,
|
||||||
preferManagement = false
|
preferManagement = false
|
||||||
): Promise<ManagedResponse> {
|
): Promise<ManagedResponse> {
|
||||||
|
const deadlineMs = Date.now() + MANAGEMENT_API_TIMEOUT_MS;
|
||||||
|
let managementAttempted = false;
|
||||||
|
|
||||||
if (preferManagement) {
|
if (preferManagement) {
|
||||||
const managedResult = await performManagedGeminiRequest(accountId, url, body);
|
managementAttempted = true;
|
||||||
|
const managedResult = await performManagedGeminiRequest(
|
||||||
|
accountId,
|
||||||
|
url,
|
||||||
|
body,
|
||||||
|
getRemainingTimeoutMs(deadlineMs)
|
||||||
|
);
|
||||||
if (managedResult) {
|
if (managedResult) {
|
||||||
return managedResult;
|
return managedResult;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
const controller = new AbortController();
|
const controller = new AbortController();
|
||||||
const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS);
|
const timeoutId = setTimeout(() => controller.abort(), getRemainingTimeoutMs(deadlineMs));
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const response = await fetch(url, {
|
const response = await fetch(url, {
|
||||||
@@ -346,7 +363,16 @@ async function performGeminiCliRequest(
|
|||||||
return directResult;
|
return directResult;
|
||||||
}
|
}
|
||||||
|
|
||||||
const managedResult = await performManagedGeminiRequest(accountId, url, body);
|
if (managementAttempted) {
|
||||||
|
return directResult;
|
||||||
|
}
|
||||||
|
|
||||||
|
const managedResult = await performManagedGeminiRequest(
|
||||||
|
accountId,
|
||||||
|
url,
|
||||||
|
body,
|
||||||
|
getRemainingTimeoutMs(deadlineMs)
|
||||||
|
);
|
||||||
return managedResult ?? directResult;
|
return managedResult ?? directResult;
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
clearTimeout(timeoutId);
|
clearTimeout(timeoutId);
|
||||||
|
|||||||
@@ -714,6 +714,7 @@ describe('Gemini CLI Quota Fetcher', () => {
|
|||||||
expect(managedLookupRequest.url).toBe(MANAGEMENT_AUTH_FILES_URL);
|
expect(managedLookupRequest.url).toBe(MANAGEMENT_AUTH_FILES_URL);
|
||||||
expect(managedQuotaRequest.url).toBe(MANAGEMENT_API_CALL_URL);
|
expect(managedQuotaRequest.url).toBe(MANAGEMENT_API_CALL_URL);
|
||||||
expect(managedQuotaRequest.body).toContain('"auth_index":"target-auth-index"');
|
expect(managedQuotaRequest.body).toContain('"auth_index":"target-auth-index"');
|
||||||
|
expect(managedQuotaRequest.body).toContain('"Authorization":"Bearer $TOKEN$"');
|
||||||
expect(managedQuotaRequest.body).not.toContain('default-refresh-token');
|
expect(managedQuotaRequest.body).not.toContain('default-refresh-token');
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -824,6 +825,80 @@ describe('Gemini CLI Quota Fetcher', () => {
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it('does not retry the management API twice when the preferred managed path already failed', async () => {
|
||||||
|
writeGeminiToken(
|
||||||
|
{
|
||||||
|
type: 'gemini',
|
||||||
|
email: 'managed-failure@example.com',
|
||||||
|
project_id: 'managed-failure-project',
|
||||||
|
token: {
|
||||||
|
access_token: 'expired-token',
|
||||||
|
refresh_token: 'refresh-token',
|
||||||
|
expiry: Date.now() - 1000,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
'gemini-managed-failure.json'
|
||||||
|
);
|
||||||
|
|
||||||
|
mockFetch([
|
||||||
|
{
|
||||||
|
url: GEMINI_CODE_ASSIST_URL,
|
||||||
|
method: 'POST',
|
||||||
|
status: 503,
|
||||||
|
response: { error: { message: 'supplementary unavailable' } },
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
const originalFetch = globalThis.fetch;
|
||||||
|
let directQuotaAttempt = 0;
|
||||||
|
let managedLookupAttempt = 0;
|
||||||
|
let managedRequestAttempt = 0;
|
||||||
|
globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
|
||||||
|
const url =
|
||||||
|
typeof input === 'string' ? input : input instanceof URL ? input.href : input.url;
|
||||||
|
|
||||||
|
if (url === MANAGEMENT_AUTH_FILES_URL) {
|
||||||
|
managedLookupAttempt += 1;
|
||||||
|
return new Response('lookup unavailable', { status: 503 });
|
||||||
|
}
|
||||||
|
|
||||||
|
if (url === MANAGEMENT_API_CALL_URL) {
|
||||||
|
managedRequestAttempt += 1;
|
||||||
|
return new Response('should not be called', { status: 500 });
|
||||||
|
}
|
||||||
|
|
||||||
|
if (url === GEMINI_QUOTA_URL) {
|
||||||
|
directQuotaAttempt += 1;
|
||||||
|
return new Response(
|
||||||
|
JSON.stringify({
|
||||||
|
error: {
|
||||||
|
message: 'Session expired',
|
||||||
|
status: 'UNAUTHENTICATED',
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
{
|
||||||
|
status: 401,
|
||||||
|
headers: { 'Content-Type': 'application/json' },
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
return originalFetch(input, init);
|
||||||
|
}) as typeof fetch;
|
||||||
|
|
||||||
|
try {
|
||||||
|
const result = await fetchGeminiCliQuota('managed-failure@example.com');
|
||||||
|
|
||||||
|
expect(result.success).toBe(false);
|
||||||
|
expect(result.needsReauth).toBe(true);
|
||||||
|
expect(directQuotaAttempt).toBe(1);
|
||||||
|
expect(managedLookupAttempt).toBe(1);
|
||||||
|
expect(managedRequestAttempt).toBe(0);
|
||||||
|
} finally {
|
||||||
|
globalThis.fetch = originalFetch;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
it('classifies model capacity exhaustion separately from generic rate limits', async () => {
|
it('classifies model capacity exhaustion separately from generic rate limits', async () => {
|
||||||
writeActiveGeminiAccount('capacity@example.com');
|
writeActiveGeminiAccount('capacity@example.com');
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user