fix: add privacy settings hook and tests (#6360)

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,
+    });
+  });
+});

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');

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,

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';