diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 4c98827e..ceab46b1 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -138,7 +138,9 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { getShowMemoryUsage: vi.fn(() => opts.showMemoryUsage ?? false), getAccessibility: vi.fn(() => opts.accessibility ?? {}), getProjectRoot: vi.fn(() => opts.targetDir), - getGeminiClient: vi.fn(() => ({})), + getGeminiClient: vi.fn(() => ({ + getUserTier: vi.fn(), + })), getCheckpointingEnabled: vi.fn(() => opts.checkpointing ?? true), getAllGeminiMdFilenames: vi.fn(() => ['GEMINI.md']), setFlashFallbackHandler: vi.fn(), @@ -639,6 +641,7 @@ describe('App UI', () => { mockConfig.getGeminiClient.mockReturnValue({ isInitialized: vi.fn(() => true), + getUserTier: vi.fn(), } as unknown as GeminiClient); const { unmount, rerender } = render( diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 027665f1..c6e6bd43 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -206,26 +206,11 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { // Sync user tier from config when authentication changes useEffect(() => { - const syncUserTier = async () => { - try { - const configUserTier = await config.getUserTier(); - if (configUserTier !== userTier) { - setUserTier(configUserTier); - } - } catch (error) { - // Silently fail - this is not critical functionality - // Only log in debug mode to avoid cluttering the console - if (config.getDebugMode()) { - console.debug('Failed to sync user tier:', error); - } - } - }; - // Only sync when not currently authenticating if (!isAuthenticating) { - syncUserTier(); + setUserTier(config.getGeminiClient()?.getUserTier()); } - }, [config, userTier, isAuthenticating]); + }, [config, isAuthenticating]); const { isEditorDialogOpen, diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts index 23dfe403..f9c6a7a3 100644 --- a/packages/core/src/code_assist/codeAssist.ts +++ b/packages/core/src/code_assist/codeAssist.ts @@ -21,8 +21,14 @@ export async function createCodeAssistContentGenerator( authType === AuthType.CLOUD_SHELL ) { const authClient = await getOauthClient(authType, config); - const projectId = await setupUser(authClient); - return new CodeAssistServer(authClient, projectId, httpOptions, sessionId); + const userData = await setupUser(authClient); + return new CodeAssistServer( + authClient, + userData.projectId, + httpOptions, + sessionId, + userData.userTier, + ); } throw new Error(`Unsupported authType: ${authType}`); diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index fe8661f1..a6cf6af7 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -32,23 +32,6 @@ import { toCountTokenRequest, toGenerateContentRequest, } from './converter.js'; -import { Readable } from 'node:stream'; - -interface ErrorData { - error?: { - message?: string; - }; -} - -interface GaxiosResponse { - status: number; - data: unknown; -} - -interface StreamError extends Error { - status?: number; - response?: GaxiosResponse; -} /** HTTP options to be used in each of the requests. */ export interface HttpOptions { @@ -60,13 +43,12 @@ export const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com'; export const CODE_ASSIST_API_VERSION = 'v1internal'; export class CodeAssistServer implements ContentGenerator { - private userTier: UserTierId | undefined = undefined; - constructor( readonly client: OAuth2Client, readonly projectId?: string, readonly httpOptions: HttpOptions = {}, readonly sessionId?: string, + readonly userTier?: UserTierId, ) {} async generateContentStream( @@ -196,45 +178,8 @@ export class CodeAssistServer implements ContentGenerator { }); return (async function* (): AsyncGenerator { - // Convert ReadableStream to Node.js stream if needed - let nodeStream: NodeJS.ReadableStream; - - if (res.data instanceof ReadableStream) { - // Convert Web ReadableStream to Node.js Readable stream - // eslint-disable-next-line @typescript-eslint/no-explicit-any - nodeStream = Readable.fromWeb(res.data as any); - } else if ( - res.data && - typeof (res.data as NodeJS.ReadableStream).on === 'function' - ) { - // Already a Node.js stream - nodeStream = res.data as NodeJS.ReadableStream; - } else { - // If res.data is not a stream, it might be an error response - // Try to extract error information from the response - let errorMessage = - 'Response data is not a readable stream. This may indicate a server error or quota issue.'; - - if (res.data && typeof res.data === 'object') { - // Check if this is an error response with error details - const errorData = res.data as ErrorData; - if (errorData.error?.message) { - errorMessage = errorData.error.message; - } else if (typeof errorData === 'string') { - errorMessage = errorData; - } - } - - // Create an error that looks like a quota error if it contains quota information - const error: StreamError = new Error(errorMessage); - // Add status and response properties so it can be properly handled by retry logic - error.status = res.status; - error.response = res; - throw error; - } - const rl = readline.createInterface({ - input: nodeStream, + input: res.data as NodeJS.ReadableStream, crlfDelay: Infinity, // Recognizes '\r\n' and '\n' as line breaks }); @@ -256,40 +201,6 @@ export class CodeAssistServer implements ContentGenerator { })(); } - async getTier(): Promise { - if (this.userTier === undefined) { - await this.detectUserTier(); - } - return this.userTier; - } - - private async detectUserTier(): Promise { - try { - // Reset user tier when detection runs - this.userTier = undefined; - - // Only attempt tier detection if we have a project ID - if (this.projectId) { - const loadRes = await this.loadCodeAssist({ - cloudaicompanionProject: this.projectId, - metadata: { - ideType: 'IDE_UNSPECIFIED', - platform: 'PLATFORM_UNSPECIFIED', - pluginType: 'GEMINI', - duetProject: this.projectId, - }, - }); - if (loadRes.currentTier) { - this.userTier = loadRes.currentTier.id; - } - } - } catch (error) { - // Silently fail - this is not critical functionality - // We'll default to FREE tier behavior if tier detection fails - console.debug('User tier detection failed:', error); - } - } - getMethodUrl(method: string): string { const endpoint = process.env.CODE_ASSIST_ENDPOINT ?? CODE_ASSIST_ENDPOINT; return `${endpoint}/${CODE_ASSIST_API_VERSION}:${method}`; diff --git a/packages/core/src/code_assist/setup.test.ts b/packages/core/src/code_assist/setup.test.ts index 479abae0..6db5fd88 100644 --- a/packages/core/src/code_assist/setup.test.ts +++ b/packages/core/src/code_assist/setup.test.ts @@ -65,7 +65,10 @@ describe('setupUser', () => { expect.any(Object), undefined, ); - expect(projectId).toBe('server-project'); + expect(projectId).toEqual({ + projectId: 'server-project', + userTier: 'standard-tier', + }); }); it('should throw ProjectIdRequiredError when no project ID is available', async () => { diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts index 3c7b81b0..8831d24b 100644 --- a/packages/core/src/code_assist/setup.ts +++ b/packages/core/src/code_assist/setup.ts @@ -22,12 +22,17 @@ export class ProjectIdRequiredError extends Error { } } +export interface UserData { + projectId: string; + userTier: UserTierId; +} + /** * * @param projectId the user's project id, if any * @returns the user's actual project id */ -export async function setupUser(client: OAuth2Client): Promise { +export async function setupUser(client: OAuth2Client): Promise { let projectId = process.env.GOOGLE_CLOUD_PROJECT || undefined; const caServer = new CodeAssistServer(client, projectId); @@ -64,7 +69,10 @@ export async function setupUser(client: OAuth2Client): Promise { await new Promise((f) => setTimeout(f, 5000)); lroRes = await caServer.onboardUser(onboardReq); } - return lroRes.response?.cloudaicompanionProject?.id || ''; + return { + projectId: lroRes.response?.cloudaicompanionProject?.id || '', + userTier: tier.id, + }; } function getOnboardTier(res: LoadCodeAssistResponse): GeminiUserTier { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 9528f648..6a3a18b6 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -11,7 +11,6 @@ import { ContentGeneratorConfig, createContentGeneratorConfig, } from '../core/contentGenerator.js'; -import { UserTierId } from '../code_assist/types.js'; import { ToolRegistry } from '../tools/tool-registry.js'; import { LSTool } from '../tools/ls.js'; import { ReadFileTool } from '../tools/read-file.js'; @@ -359,14 +358,6 @@ export class Config { return this.quotaErrorOccurred; } - async getUserTier(): Promise { - if (!this.geminiClient) { - return undefined; - } - const generator = this.geminiClient.getContentGenerator(); - return await generator.getTier?.(); - } - getEmbeddingModel(): string { return this.embeddingModel; } diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 969e1da6..61195e2f 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -22,6 +22,7 @@ import { ChatCompressionInfo, } from './turn.js'; import { Config } from '../config/config.js'; +import { UserTierId } from '../code_assist/types.js'; import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js'; import { ReadManyFilesTool } from '../tools/read-many-files.js'; import { getResponseText } from '../utils/generateContentResponseUtilities.js'; @@ -130,6 +131,10 @@ export class GeminiClient { return this.contentGenerator; } + getUserTier(): UserTierId | undefined { + return this.contentGenerator?.userTier; + } + async addHistory(content: Content) { this.getChat().addHistory(content); } diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index b381de8e..721a16d7 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -35,7 +35,7 @@ export interface ContentGenerator { embedContent(request: EmbedContentParameters): Promise; - getTier?(): Promise; + userTier?: UserTierId; } export enum AuthType { diff --git a/packages/core/src/utils/quotaErrorDetection.ts b/packages/core/src/utils/quotaErrorDetection.ts index b07309cd..6fe9b312 100644 --- a/packages/core/src/utils/quotaErrorDetection.ts +++ b/packages/core/src/utils/quotaErrorDetection.ts @@ -68,10 +68,6 @@ export function isProQuotaExceededError(error: unknown): boolean { }; }; if (gaxiosError.response && gaxiosError.response.data) { - console.log( - '[DEBUG] isProQuotaExceededError - checking response data:', - gaxiosError.response.data, - ); if (typeof gaxiosError.response.data === 'string') { return checkMessage(gaxiosError.response.data); } @@ -87,11 +83,6 @@ export function isProQuotaExceededError(error: unknown): boolean { } } } - - console.log( - '[DEBUG] isProQuotaExceededError - no matching error format for:', - error, - ); return false; }