abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-20 16:57:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Auth First Run (#1207)

Browse Source 
Co-authored-by: Tommaso Sciortino <sciortino@gmail.com>
Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
This commit is contained in:
matt korwel
2025-06-19 16:52:22 -07:00
committed by GitHub
parent c48fcaa8c3
commit 04518b52c0
37 changed files with 636 additions and 349 deletions
Download Patch File Download Diff File Expand all files Collapse all files

154
packages/cli/src/utils/modelCheck.test.ts
View File

@@ -1,154 +0,0 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { getEffectiveModel } from './modelCheck.js';
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
} from '@gemini-cli/core';
// Mock global fetch
global.fetch = vi.fn();
// Mock AbortController
const mockAbort = vi.fn();
global.AbortController = vi.fn(() => ({
signal: { aborted: false }, // Start with not aborted
abort: mockAbort,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
})) as any;
describe('getEffectiveModel', () => {
const apiKey = 'test-api-key';
beforeEach(() => {
vi.useFakeTimers();
vi.clearAllMocks();
// Reset signal for each test if AbortController mock is more complex
global.AbortController = vi.fn(() => ({
signal: { aborted: false },
abort: mockAbort,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
})) as any;
});
afterEach(() => {
vi.restoreAllMocks();
vi.useRealTimers();
});
describe('when currentConfiguredModel is not DEFAULT_GEMINI_MODEL', () => {
it('should return the currentConfiguredModel without fetching', async () => {
const customModel = 'custom-model-name';
const result = await getEffectiveModel(apiKey, customModel);
expect(result).toEqual(customModel);
expect(fetch).not.toHaveBeenCalled();
});
});
describe('when currentConfiguredModel is DEFAULT_GEMINI_MODEL', () => {
it('should switch to DEFAULT_GEMINI_FLASH_MODEL if fetch returns 429', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({
ok: false,
status: 429,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_FLASH_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
expect(fetch).toHaveBeenCalledWith(
`https://generativelanguage.googleapis.com/v1beta/models/${DEFAULT_GEMINI_MODEL}:generateContent?key=${apiKey}`,
expect.any(Object),
);
});
it('should return DEFAULT_GEMINI_MODEL if fetch returns 200', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({
ok: true,
status: 200,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch returns a non-429 error status (e.g., 500)', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({
ok: false,
status: 500,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch throws a network error', async () => {
(fetch as vi.Mock).mockRejectedValueOnce(new Error('Network error'));
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch times out', async () => {
// Simulate AbortController's signal changing and fetch throwing AbortError
const abortControllerInstance = {
signal: { aborted: false }, // mutable signal
abort: vi.fn(() => {
abortControllerInstance.signal.aborted = true; // Use abortControllerInstance
}),
};
(global.AbortController as vi.Mock).mockImplementationOnce(
() => abortControllerInstance,
);
(fetch as vi.Mock).mockImplementationOnce(
async ({ signal }: { signal: AbortSignal }) => {
// Simulate the timeout advancing and abort being called
vi.advanceTimersByTime(2000);
if (signal.aborted) {
throw new DOMException('Aborted', 'AbortError');
}
// Should not reach here in a timeout scenario
return { ok: true, status: 200 };
},
);
const resultPromise = getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
// Ensure timers are advanced to trigger the timeout within getEffectiveModel
await vi.advanceTimersToNextTimerAsync(); // Or advanceTimersByTime(2000) if more precise control is needed
const result = await resultPromise;
expect(mockAbort).toHaveBeenCalledTimes(0); // setTimeout calls controller.abort(), not our direct mockAbort
expect(abortControllerInstance.abort).toHaveBeenCalledTimes(1);
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should correctly pass API key and model in the fetch request', async () => {
(fetch as vi.Mock).mockResolvedValueOnce({ ok: true, status: 200 });
const specificApiKey = 'specific-key-for-this-test';
await getEffectiveModel(specificApiKey, DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledWith(
`https://generativelanguage.googleapis.com/v1beta/models/${DEFAULT_GEMINI_MODEL}:generateContent?key=${specificApiKey}`,
expect.objectContaining({
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
contents: [{ parts: [{ text: 'test' }] }],
generationConfig: {
maxOutputTokens: 1,
temperature: 0,
topK: 1,
thinkingConfig: { thinkingBudget: 0, includeThoughts: false },
},
}),
}),
);
});
});
});

68
packages/cli/src/utils/modelCheck.ts
View File

@@ -1,68 +0,0 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
} from '@gemini-cli/core';
/**
* Checks if the default "pro" model is rate-limited and returns a fallback "flash"
* model if necessary. This function is designed to be silent.
* @param apiKey The API key to use for the check.
* @param currentConfiguredModel The model currently configured in settings.
* @returns An object indicating the model to use, whether a switch occurred,
* and the original model if a switch happened.
*/
export async function getEffectiveModel(
apiKey: string,
currentConfiguredModel: string,
): Promise<string> {
if (currentConfiguredModel !== DEFAULT_GEMINI_MODEL) {
// Only check if the user is trying to use the specific pro model we want to fallback from.
return currentConfiguredModel;
}
const modelToTest = DEFAULT_GEMINI_MODEL;
const fallbackModel = DEFAULT_GEMINI_FLASH_MODEL;
const endpoint = `https://generativelanguage.googleapis.com/v1beta/models/${modelToTest}:generateContent?key=${apiKey}`;
const body = JSON.stringify({
contents: [{ parts: [{ text: 'test' }] }],
generationConfig: {
maxOutputTokens: 1,
temperature: 0,
topK: 1,
thinkingConfig: { thinkingBudget: 0, includeThoughts: false },
},
});
const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), 2000); // 500ms timeout for the request
try {
const response = await fetch(endpoint, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body,
signal: controller.signal,
});
clearTimeout(timeoutId);
if (response.status === 429) {
console.log(
`[INFO] Your configured model (${modelToTest}) was temporarily unavailable. Switched to ${fallbackModel} for this session.`,
);
return fallbackModel;
}
// For any other case (success, other error codes), we stick to the original model.
return currentConfiguredModel;
} catch (_error) {
clearTimeout(timeoutId);
// On timeout or any other fetch error, stick to the original model.
return currentConfiguredModel;
}
}