diff options
| -rw-r--r-- | packages/cli/src/ui/hooks/usePrivacySettings.test.ts | 242 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/usePrivacySettings.ts | 15 | ||||
| -rw-r--r-- | packages/core/src/core/loggingContentGenerator.ts | 14 | ||||
| -rw-r--r-- | packages/core/src/index.ts | 1 |
4 files changed, 261 insertions, 11 deletions
diff --git a/packages/cli/src/ui/hooks/usePrivacySettings.test.ts b/packages/cli/src/ui/hooks/usePrivacySettings.test.ts new file mode 100644 index 00000000..9b2a7c2c --- /dev/null +++ b/packages/cli/src/ui/hooks/usePrivacySettings.test.ts @@ -0,0 +1,242 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { renderHook, waitFor } from '@testing-library/react'; +import { + Config, + CodeAssistServer, + LoggingContentGenerator, + UserTierId, + GeminiClient, + ContentGenerator, +} from '@google/gemini-cli-core'; +import { OAuth2Client } from 'google-auth-library'; +import { usePrivacySettings } from './usePrivacySettings.js'; + +// Mock the dependencies +vi.mock('@google/gemini-cli-core', () => { + // Mock classes for instanceof checks + class MockCodeAssistServer { + projectId = 'test-project-id'; + loadCodeAssist = vi.fn(); + getCodeAssistGlobalUserSetting = vi.fn(); + setCodeAssistGlobalUserSetting = vi.fn(); + + constructor( + _client?: GeminiClient, + _projectId?: string, + _httpOptions?: Record<string, unknown>, + _sessionId?: string, + _userTier?: UserTierId, + ) {} + } + + class MockLoggingContentGenerator { + getWrapped = vi.fn(); + + constructor( + _wrapped?: ContentGenerator, + _config?: Record<string, unknown>, + ) {} + } + + return { + Config: vi.fn(), + CodeAssistServer: MockCodeAssistServer, + LoggingContentGenerator: MockLoggingContentGenerator, + GeminiClient: vi.fn(), + UserTierId: { + FREE: 'free-tier', + LEGACY: 'legacy-tier', + STANDARD: 'standard-tier', + }, + }; +}); + +describe('usePrivacySettings', () => { + let mockConfig: Config; + let mockClient: GeminiClient; + let mockCodeAssistServer: CodeAssistServer; + let mockLoggingContentGenerator: LoggingContentGenerator; + + beforeEach(() => { + vi.clearAllMocks(); + + // Create mock CodeAssistServer instance + mockCodeAssistServer = new CodeAssistServer( + null as unknown as OAuth2Client, + 'test-project-id', + ) as unknown as CodeAssistServer; + ( + mockCodeAssistServer.loadCodeAssist as ReturnType<typeof vi.fn> + ).mockResolvedValue({ + currentTier: { id: UserTierId.FREE }, + }); + ( + mockCodeAssistServer.getCodeAssistGlobalUserSetting as ReturnType< + typeof vi.fn + > + ).mockResolvedValue({ + freeTierDataCollectionOptin: true, + }); + ( + mockCodeAssistServer.setCodeAssistGlobalUserSetting as ReturnType< + typeof vi.fn + > + ).mockResolvedValue({ + freeTierDataCollectionOptin: false, + }); + + // Create mock LoggingContentGenerator that wraps the CodeAssistServer + mockLoggingContentGenerator = new LoggingContentGenerator( + mockCodeAssistServer, + null as unknown as Config, + ) as unknown as LoggingContentGenerator; + ( + mockLoggingContentGenerator.getWrapped as ReturnType<typeof vi.fn> + ).mockReturnValue(mockCodeAssistServer); + + // Create mock GeminiClient + mockClient = { + getContentGenerator: vi.fn().mockReturnValue(mockLoggingContentGenerator), + } as unknown as GeminiClient; + + // Create mock Config + mockConfig = { + getGeminiClient: vi.fn().mockReturnValue(mockClient), + } as unknown as Config; + }); + + it('should handle LoggingContentGenerator wrapper correctly and not throw "Oauth not being used" error', async () => { + const { result } = renderHook(() => usePrivacySettings(mockConfig)); + + // Initial state should be loading + expect(result.current.privacyState.isLoading).toBe(true); + expect(result.current.privacyState.error).toBeUndefined(); + + // Wait for the hook to complete + await waitFor(() => { + expect(result.current.privacyState.isLoading).toBe(false); + }); + + // Should not have the "Oauth not being used" error + expect(result.current.privacyState.error).toBeUndefined(); + expect(result.current.privacyState.isFreeTier).toBe(true); + expect(result.current.privacyState.dataCollectionOptIn).toBe(true); + + // Verify that getWrapped was called to unwrap the LoggingContentGenerator + expect(mockLoggingContentGenerator.getWrapped).toHaveBeenCalled(); + }); + + it('should work with direct CodeAssistServer (no wrapper)', async () => { + // Test case where the content generator is directly a CodeAssistServer + const directServer = new CodeAssistServer( + null as unknown as OAuth2Client, + 'test-project-id', + ) as unknown as CodeAssistServer; + (directServer.loadCodeAssist as ReturnType<typeof vi.fn>).mockResolvedValue( + { + currentTier: { id: UserTierId.FREE }, + }, + ); + ( + directServer.getCodeAssistGlobalUserSetting as ReturnType<typeof vi.fn> + ).mockResolvedValue({ + freeTierDataCollectionOptin: true, + }); + + mockClient.getContentGenerator = vi.fn().mockReturnValue(directServer); + + const { result } = renderHook(() => usePrivacySettings(mockConfig)); + + await waitFor(() => { + expect(result.current.privacyState.isLoading).toBe(false); + }); + + expect(result.current.privacyState.error).toBeUndefined(); + expect(result.current.privacyState.isFreeTier).toBe(true); + expect(result.current.privacyState.dataCollectionOptIn).toBe(true); + }); + + it('should handle paid tier users correctly', async () => { + // Mock paid tier response + ( + mockCodeAssistServer.loadCodeAssist as ReturnType<typeof vi.fn> + ).mockResolvedValue({ + currentTier: { id: UserTierId.STANDARD }, + }); + + const { result } = renderHook(() => usePrivacySettings(mockConfig)); + + await waitFor(() => { + expect(result.current.privacyState.isLoading).toBe(false); + }); + + expect(result.current.privacyState.error).toBeUndefined(); + expect(result.current.privacyState.isFreeTier).toBe(false); + expect(result.current.privacyState.dataCollectionOptIn).toBeUndefined(); + }); + + it('should throw error when content generator is not a CodeAssistServer', async () => { + // Mock a non-CodeAssistServer content generator + const mockOtherGenerator = { someOtherMethod: vi.fn() }; + ( + mockLoggingContentGenerator.getWrapped as ReturnType<typeof vi.fn> + ).mockReturnValue(mockOtherGenerator); + + const { result } = renderHook(() => usePrivacySettings(mockConfig)); + + await waitFor(() => { + expect(result.current.privacyState.isLoading).toBe(false); + }); + + expect(result.current.privacyState.error).toBe('Oauth not being used'); + }); + + it('should throw error when CodeAssistServer has no projectId', async () => { + // Mock CodeAssistServer without projectId + const mockServerNoProject = { + ...mockCodeAssistServer, + projectId: undefined, + }; + ( + mockLoggingContentGenerator.getWrapped as ReturnType<typeof vi.fn> + ).mockReturnValue(mockServerNoProject); + + const { result } = renderHook(() => usePrivacySettings(mockConfig)); + + await waitFor(() => { + expect(result.current.privacyState.isLoading).toBe(false); + }); + + expect(result.current.privacyState.error).toBe('Oauth not being used'); + }); + + it('should update data collection opt-in setting', async () => { + const { result } = renderHook(() => usePrivacySettings(mockConfig)); + + // Wait for initial load + await waitFor(() => { + expect(result.current.privacyState.isLoading).toBe(false); + }); + + // Update the setting + await result.current.updateDataCollectionOptIn(false); + + // Wait for update to complete + await waitFor(() => { + expect(result.current.privacyState.dataCollectionOptIn).toBe(false); + }); + + expect( + mockCodeAssistServer.setCodeAssistGlobalUserSetting, + ).toHaveBeenCalledWith({ + cloudaicompanionProject: 'test-project-id', + freeTierDataCollectionOptin: false, + }); + }); +}); diff --git a/packages/cli/src/ui/hooks/usePrivacySettings.ts b/packages/cli/src/ui/hooks/usePrivacySettings.ts index bc98649b..47a62588 100644 --- a/packages/cli/src/ui/hooks/usePrivacySettings.ts +++ b/packages/cli/src/ui/hooks/usePrivacySettings.ts @@ -5,7 +5,12 @@ */ import { useState, useEffect, useCallback } from 'react'; -import { Config, CodeAssistServer, UserTierId } from '@google/gemini-cli-core'; +import { + Config, + CodeAssistServer, + UserTierId, + LoggingContentGenerator, +} from '@google/gemini-cli-core'; export interface PrivacyState { isLoading: boolean; @@ -80,7 +85,13 @@ export const usePrivacySettings = (config: Config) => { }; function getCodeAssistServer(config: Config): CodeAssistServer { - const server = config.getGeminiClient().getContentGenerator(); + let server = config.getGeminiClient().getContentGenerator(); + + // Unwrap LoggingContentGenerator if present + if (server instanceof LoggingContentGenerator) { + server = server.getWrapped(); + } + // Neither of these cases should ever happen. if (!(server instanceof CodeAssistServer)) { throw new Error('Oauth not being used'); diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts index 13bd6918..2abe3dce 100644 --- a/packages/core/src/core/loggingContentGenerator.ts +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -27,20 +27,12 @@ import { } from '../telemetry/loggers.js'; import { ContentGenerator } from './contentGenerator.js'; import { toContents } from '../code_assist/converter.js'; +import { isStructuredError } from '../utils/quotaErrorDetection.js'; interface StructuredError { status: number; } -export function isStructuredError(error: unknown): error is StructuredError { - return ( - typeof error === 'object' && - error !== null && - 'status' in error && - typeof (error as StructuredError).status === 'number' - ); -} - /** * A decorator that wraps a ContentGenerator to add logging to API calls. */ @@ -50,6 +42,10 @@ export class LoggingContentGenerator implements ContentGenerator { private readonly config: Config, ) {} + getWrapped(): ContentGenerator { + return this.wrapped; + } + private logApiRequest( contents: Content[], model: string, diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 45f7e4ce..e8dbe947 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -10,6 +10,7 @@ export * from './config/config.js'; // Export Core Logic export * from './core/client.js'; export * from './core/contentGenerator.js'; +export * from './core/loggingContentGenerator.js'; export * from './core/geminiChat.js'; export * from './core/logger.js'; export * from './core/prompts.js'; |