chore(usage telemetry): Freshen up Clearcut logging (#6013)

Co-authored-by: christine betts <[email protected]>
Co-authored-by: Jacob Richman <[email protected]>
Co-authored-by: matt korwel <[email protected]>

1 files changed, 154 insertions, 99 deletions

diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts
index f955eb5a..96129ad3 100644
--- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts
+++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts
@@ -4,33 +4,49 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { vi, describe, it, expect, beforeEach, afterEach } from 'vitest';
-import * as https from 'https';
-import { ClientRequest, IncomingMessage } from 'http';
-import { Readable, Writable } from 'stream';
-
 import {
-  ClearcutLogger,
-  LogResponse,
-  LogEventEntry,
-} from './clearcut-logger.js';
-import { Config } from '../../config/config.js';
+  vi,
+  describe,
+  it,
+  expect,
+  afterEach,
+  beforeAll,
+  afterAll,
+} from 'vitest';
+
+import { ClearcutLogger, LogEventEntry, TEST_ONLY } from './clearcut-logger.js';
+import { ConfigParameters } from '../../config/config.js';
 import * as userAccount from '../../utils/user_account.js';
 import * as userId from '../../utils/user_id.js';
+import { EventMetadataKey } from './event-metadata-key.js';
+import { makeFakeConfig } from '../../test-utils/config.js';
+import { http, HttpResponse } from 'msw';
+import { server } from '../../mocks/msw.js';
 
-// Mock dependencies
-vi.mock('https-proxy-agent');
-vi.mock('https');
 vi.mock('../../utils/user_account');
 vi.mock('../../utils/user_id');
 
-const mockHttps = vi.mocked(https);
 const mockUserAccount = vi.mocked(userAccount);
 const mockUserId = vi.mocked(userId);
 
+// TODO(richieforeman): Consider moving this to test setup globally.
+beforeAll(() => {
+  server.listen({});
+});
+
+afterEach(() => {
+  server.resetHandlers();
+});
+
+afterAll(() => {
+  server.close();
+});
+
 describe('ClearcutLogger', () => {
-  let mockConfig: Config;
-  let logger: ClearcutLogger | undefined;
+  const NEXT_WAIT_MS = 1234;
+  const CLEARCUT_URL = 'https://play.googleapis.com/log';
+  const MOCK_DATE = new Date('2025-01-02T00:00:00.000Z');
+  const EXAMPLE_RESPONSE = `["${NEXT_WAIT_MS}",null,[[["ANDROID_BACKUP",0],["BATTERY_STATS",0],["SMART_SETUP",0],["TRON",0]],-3334737594024971225],[]]`;
 
   // A helper to get the internal events array for testing
   const getEvents = (l: ClearcutLogger): LogEventEntry[][] =>
@@ -38,32 +54,37 @@ describe('ClearcutLogger', () => {
 
   const getEventsSize = (l: ClearcutLogger): number => l['events'].size;
 
-  const getMaxEvents = (l: ClearcutLogger): number => l['max_events'];
-
-  const getMaxRetryEvents = (l: ClearcutLogger): number =>
-    l['max_retry_events'];
-
   const requeueFailedEvents = (l: ClearcutLogger, events: LogEventEntry[][]) =>
     l['requeueFailedEvents'](events);
 
-  beforeEach(() => {
+  function setup({
+    config = {} as Partial<ConfigParameters>,
+    lifetimeGoogleAccounts = 1,
+    cachedGoogleAccount = '[email protected]',
+    installationId = 'test-installation-id',
+  } = {}) {
+    server.resetHandlers(
+      http.post(CLEARCUT_URL, () => HttpResponse.text(EXAMPLE_RESPONSE)),
+    );
+
     vi.useFakeTimers();
-    vi.setSystemTime(new Date());
+    vi.setSystemTime(MOCK_DATE);
 
-    mockConfig = {
-      getUsageStatisticsEnabled: vi.fn().mockReturnValue(true),
-      getDebugMode: vi.fn().mockReturnValue(false),
-      getSessionId: vi.fn().mockReturnValue('test-session-id'),
-      getProxy: vi.fn().mockReturnValue(undefined),
-    } as unknown as Config;
+    const loggerConfig = makeFakeConfig({
+      ...config,
+    });
+    ClearcutLogger.clearInstance();
 
-    mockUserAccount.getCachedGoogleAccount.mockReturnValue('[email protected]');
-    mockUserAccount.getLifetimeGoogleAccounts.mockReturnValue(1);
-    mockUserId.getInstallationId.mockReturnValue('test-installation-id');
+    mockUserAccount.getCachedGoogleAccount.mockReturnValue(cachedGoogleAccount);
+    mockUserAccount.getLifetimeGoogleAccounts.mockReturnValue(
+      lifetimeGoogleAccounts,
+    );
+    mockUserId.getInstallationId.mockReturnValue(installationId);
 
-    logger = ClearcutLogger.getInstance(mockConfig);
-    expect(logger).toBeDefined();
-  });
+    const logger = ClearcutLogger.getInstance(loggerConfig);
+
+    return { logger, loggerConfig };
+  }
 
   afterEach(() => {
     ClearcutLogger.clearInstance();
@@ -71,109 +92,131 @@ describe('ClearcutLogger', () => {
     vi.restoreAllMocks();
   });
 
-  it('should not return an instance if usage statistics are disabled', () => {
-    ClearcutLogger.clearInstance();
-    vi.spyOn(mockConfig, 'getUsageStatisticsEnabled').mockReturnValue(false);
-    const disabledLogger = ClearcutLogger.getInstance(mockConfig);
-    expect(disabledLogger).toBeUndefined();
+  describe('getInstance', () => {
+    it.each([
+      { usageStatisticsEnabled: false, expectedValue: undefined },
+      {
+        usageStatisticsEnabled: true,
+        expectedValue: expect.any(ClearcutLogger),
+      },
+    ])(
+      'returns an instance if usage statistics are enabled',
+      ({ usageStatisticsEnabled, expectedValue }) => {
+        ClearcutLogger.clearInstance();
+        const { logger } = setup({
+          config: {
+            usageStatisticsEnabled,
+          },
+        });
+        expect(logger).toEqual(expectedValue);
+      },
+    );
+
+    it('is a singleton', () => {
+      ClearcutLogger.clearInstance();
+      const { loggerConfig } = setup();
+      const logger1 = ClearcutLogger.getInstance(loggerConfig);
+      const logger2 = ClearcutLogger.getInstance(loggerConfig);
+      expect(logger1).toBe(logger2);
+    });
+  });
+
+  describe('createLogEvent', () => {
+    it('logs the total number of google accounts', () => {
+      const { logger } = setup({
+        lifetimeGoogleAccounts: 9001,
+      });
+
+      const event = logger?.createLogEvent('abc', []);
+
+      expect(event?.event_metadata[0][0]).toEqual({
+        gemini_cli_key: EventMetadataKey.GEMINI_CLI_GOOGLE_ACCOUNTS_COUNT,
+        value: '9001',
+      });
+    });
+
+    it('logs the current surface', () => {
+      const { logger } = setup({});
+
+      const event = logger?.createLogEvent('abc', []);
+
+      expect(event?.event_metadata[0][1]).toEqual({
+        gemini_cli_key: EventMetadataKey.GEMINI_CLI_SURFACE,
+        value: 'SURFACE_NOT_SET',
+      });
+    });
   });
 
   describe('enqueueLogEvent', () => {
     it('should add events to the queue', () => {
+      const { logger } = setup();
       logger!.enqueueLogEvent({ test: 'event1' });
       expect(getEventsSize(logger!)).toBe(1);
     });
 
     it('should evict the oldest event when the queue is full', () => {
-      const maxEvents = getMaxEvents(logger!);
+      const { logger } = setup();
 
-      for (let i = 0; i < maxEvents; i++) {
+      for (let i = 0; i < TEST_ONLY.MAX_EVENTS; i++) {
         logger!.enqueueLogEvent({ event_id: i });
       }
 
-      expect(getEventsSize(logger!)).toBe(maxEvents);
+      expect(getEventsSize(logger!)).toBe(TEST_ONLY.MAX_EVENTS);
       const firstEvent = JSON.parse(
         getEvents(logger!)[0][0].source_extension_json,
       );
       expect(firstEvent.event_id).toBe(0);
 
       // This should push out the first event
-      logger!.enqueueLogEvent({ event_id: maxEvents });
+      logger!.enqueueLogEvent({ event_id: TEST_ONLY.MAX_EVENTS });
 
-      expect(getEventsSize(logger!)).toBe(maxEvents);
+      expect(getEventsSize(logger!)).toBe(TEST_ONLY.MAX_EVENTS);
       const newFirstEvent = JSON.parse(
         getEvents(logger!)[0][0].source_extension_json,
       );
       expect(newFirstEvent.event_id).toBe(1);
       const lastEvent = JSON.parse(
-        getEvents(logger!)[maxEvents - 1][0].source_extension_json,
+        getEvents(logger!)[TEST_ONLY.MAX_EVENTS - 1][0].source_extension_json,
       );
-      expect(lastEvent.event_id).toBe(maxEvents);
+      expect(lastEvent.event_id).toBe(TEST_ONLY.MAX_EVENTS);
     });
   });
 
   describe('flushToClearcut', () => {
-    let mockRequest: Writable;
-    let mockResponse: Readable & Partial<IncomingMessage>;
-
-    beforeEach(() => {
-      mockRequest = new Writable({
-        write(chunk, encoding, callback) {
-          callback();
+    it('allows for usage with a configured proxy agent', async () => {
+      const { logger } = setup({
+        config: {
+          proxy: 'http://mycoolproxy.whatever.com:3128',
         },
       });
-      vi.spyOn(mockRequest, 'on');
-      vi.spyOn(mockRequest, 'end').mockReturnThis();
-      vi.spyOn(mockRequest, 'destroy').mockReturnThis();
 
-      mockResponse = new Readable({ read() {} }) as Readable &
-        Partial<IncomingMessage>;
+      logger!.enqueueLogEvent({ event_id: 1 });
 
-      mockHttps.request.mockImplementation(
-        (
-          _options: string | https.RequestOptions | URL,
-          ...args: unknown[]
-        ): ClientRequest => {
-          const callback = args.find((arg) => typeof arg === 'function') as
-            | ((res: IncomingMessage) => void)
-            | undefined;
+      const response = await logger!.flushToClearcut();
 
-          if (callback) {
-            callback(mockResponse as IncomingMessage);
-          }
-          return mockRequest as ClientRequest;
-        },
-      );
+      expect(response.nextRequestWaitMs).toBe(NEXT_WAIT_MS);
     });
 
     it('should clear events on successful flush', async () => {
-      mockResponse.statusCode = 200;
-      const mockResponseBody = { nextRequestWaitMs: 1000 };
-      // Encoded protobuf for {nextRequestWaitMs: 1000} which is `08 E8 07`
-      const encodedResponse = Buffer.from([8, 232, 7]);
+      const { logger } = setup();
 
       logger!.enqueueLogEvent({ event_id: 1 });
-      const flushPromise = logger!.flushToClearcut();
+      const response = await logger!.flushToClearcut();
 
-      mockResponse.push(encodedResponse);
-      mockResponse.push(null); // End the stream
-
-      const response: LogResponse = await flushPromise;
-
-      expect(getEventsSize(logger!)).toBe(0);
-      expect(response.nextRequestWaitMs).toBe(
-        mockResponseBody.nextRequestWaitMs,
-      );
+      expect(getEvents(logger!)).toEqual([]);
+      expect(response.nextRequestWaitMs).toBe(NEXT_WAIT_MS);
     });
 
     it('should handle a network error and requeue events', async () => {
+      const { logger } = setup();
+
+      server.resetHandlers(http.post(CLEARCUT_URL, () => HttpResponse.error()));
       logger!.enqueueLogEvent({ event_id: 1 });
       logger!.enqueueLogEvent({ event_id: 2 });
       expect(getEventsSize(logger!)).toBe(2);
 
-      const flushPromise = logger!.flushToClearcut();
-      mockRequest.emit('error', new Error('Network error'));
-      await flushPromise;
+      const x = logger!.flushToClearcut();
+      await x;
 
       expect(getEventsSize(logger!)).toBe(2);
       const events = getEvents(logger!);
@@ -181,18 +224,28 @@ describe('ClearcutLogger', () => {
     });
 
     it('should handle an HTTP error and requeue events', async () => {
-      mockResponse.statusCode = 500;
-      mockResponse.statusMessage = 'Internal Server Error';
+      const { logger } = setup();
+
+      server.resetHandlers(
+        http.post(
+          CLEARCUT_URL,
+          () =>
+            new HttpResponse(
+              { 'the system is down': true },
+              {
+                status: 500,
+              },
+            ),
+        ),
+      );
 
       logger!.enqueueLogEvent({ event_id: 1 });
       logger!.enqueueLogEvent({ event_id: 2 });
-      expect(getEventsSize(logger!)).toBe(2);
 
-      const flushPromise = logger!.flushToClearcut();
-      mockResponse.emit('end'); // End the response to trigger promise resolution
-      await flushPromise;
+      expect(getEvents(logger!).length).toBe(2);
+      await logger!.flushToClearcut();
 
-      expect(getEventsSize(logger!)).toBe(2);
+      expect(getEvents(logger!).length).toBe(2);
       const events = getEvents(logger!);
       expect(JSON.parse(events[0][0].source_extension_json).event_id).toBe(1);
     });
@@ -200,7 +253,8 @@ describe('ClearcutLogger', () => {
 
   describe('requeueFailedEvents logic', () => {
     it('should limit the number of requeued events to max_retry_events', () => {
-      const maxRetryEvents = getMaxRetryEvents(logger!);
+      const { logger } = setup();
+      const maxRetryEvents = TEST_ONLY.MAX_RETRY_EVENTS;
       const eventsToLogCount = maxRetryEvents + 5;
       const eventsToSend: LogEventEntry[][] = [];
       for (let i = 0; i < eventsToLogCount; i++) {
@@ -225,7 +279,8 @@ describe('ClearcutLogger', () => {
     });
 
     it('should not requeue more events than available space in the queue', () => {
-      const maxEvents = getMaxEvents(logger!);
+      const { logger } = setup();
+      const maxEvents = TEST_ONLY.MAX_EVENTS;
       const spaceToLeave = 5;
       const initialEventCount = maxEvents - spaceToLeave;
       for (let i = 0; i < initialEventCount; i++) {