fix(cli) - Move logging into CodeAssistServer (#5781)

Co-authored-by: Shi Shu <shii@google.com>

packages/core/src/code_assist/converter.test.ts

@@ -9,8 +9,10 @@ import {
   toGenerateContentRequest,
   fromGenerateContentResponse,
   CaGenerateContentResponse,
+  toContents,
 } from './converter.js';
 import {
+  ContentListUnion,
   GenerateContentParameters,
   GenerateContentResponse,
   FinishReason,
@@ -295,4 +297,57 @@ describe('converter', () => {
       );
     });
   });
+
+  describe('toContents', () => {
+    it('should handle Content', () => {
+      const content: ContentListUnion = {
+        role: 'user',
+        parts: [{ text: 'hello' }],
+      };
+      expect(toContents(content)).toEqual([
+        { role: 'user', parts: [{ text: 'hello' }] },
+      ]);
+    });
+
+    it('should handle array of Contents', () => {
+      const contents: ContentListUnion = [
+        { role: 'user', parts: [{ text: 'hello' }] },
+        { role: 'model', parts: [{ text: 'hi' }] },
+      ];
+      expect(toContents(contents)).toEqual([
+        { role: 'user', parts: [{ text: 'hello' }] },
+        { role: 'model', parts: [{ text: 'hi' }] },
+      ]);
+    });
+
+    it('should handle Part', () => {
+      const part: ContentListUnion = { text: 'a part' };
+      expect(toContents(part)).toEqual([
+        { role: 'user', parts: [{ text: 'a part' }] },
+      ]);
+    });
+
+    it('should handle array of Parts', () => {
+      const parts = [{ text: 'part 1' }, 'part 2'];
+      expect(toContents(parts)).toEqual([
+        { role: 'user', parts: [{ text: 'part 1' }] },
+        { role: 'user', parts: [{ text: 'part 2' }] },
+      ]);
+    });
+
+    it('should handle string', () => {
+      const str: ContentListUnion = 'a string';
+      expect(toContents(str)).toEqual([
+        { role: 'user', parts: [{ text: 'a string' }] },
+      ]);
+    });
+
+    it('should handle array of strings', () => {
+      const strings: ContentListUnion = ['string 1', 'string 2'];
+      expect(toContents(strings)).toEqual([
+        { role: 'user', parts: [{ text: 'string 1' }] },
+        { role: 'user', parts: [{ text: 'string 2' }] },
+      ]);
+    });
+  });
 });

packages/core/src/code_assist/converter.ts

@@ -157,7 +157,7 @@ function toVertexGenerateContentRequest(
   };
 }
 
-function toContents(contents: ContentListUnion): Content[] {
+export function toContents(contents: ContentListUnion): Content[] {
   if (Array.isArray(contents)) {
     // it's a Content[] or a PartsUnion[]
     return contents.map(toContent);

packages/core/src/code_assist/server.test.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi } from 'vitest';
+import { beforeEach, describe, it, expect, vi } from 'vitest';
 import { CodeAssistServer } from './server.js';
 import { OAuth2Client } from 'google-auth-library';
 import { UserTierId } from './types.js';
@@ -12,6 +12,10 @@ import { UserTierId } from './types.js';
 vi.mock('google-auth-library');
 
 describe('CodeAssistServer', () => {
+  beforeEach(() => {
+    vi.resetAllMocks();
+  });
+
   it('should be able to be constructed', () => {
     const auth = new OAuth2Client();
     const server = new CodeAssistServer(

packages/core/src/core/contentGenerator.test.ts

@@ -9,10 +9,12 @@ import {
   createContentGenerator,
   AuthType,
   createContentGeneratorConfig,
+  ContentGenerator,
 } from './contentGenerator.js';
 import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
 import { GoogleGenAI } from '@google/genai';
 import { Config } from '../config/config.js';
+import { LoggingContentGenerator } from './loggingContentGenerator.js';
 
 vi.mock('../code_assist/codeAssist.js');
 vi.mock('@google/genai');
@@ -21,7 +23,7 @@ const mockConfig = {} as unknown as Config;
 
 describe('createContentGenerator', () => {
   it('should create a CodeAssistContentGenerator', async () => {
-    const mockGenerator = {} as unknown;
+    const mockGenerator = {} as unknown as ContentGenerator;
     vi.mocked(createCodeAssistContentGenerator).mockResolvedValue(
       mockGenerator as never,
     );
@@ -33,13 +35,15 @@ describe('createContentGenerator', () => {
       mockConfig,
     );
     expect(createCodeAssistContentGenerator).toHaveBeenCalled();
-    expect(generator).toBe(mockGenerator);
+    expect(generator).toEqual(
+      new LoggingContentGenerator(mockGenerator, mockConfig),
+    );
   });
 
   it('should create a GoogleGenAI content generator', async () => {
     const mockGenerator = {
       models: {},
-    } as unknown;
+    } as unknown as GoogleGenAI;
     vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
     const generator = await createContentGenerator(
       {
@@ -58,7 +62,12 @@ describe('createContentGenerator', () => {
         },
       },
     });
-    expect(generator).toBe((mockGenerator as GoogleGenAI).models);
+    expect(generator).toEqual(
+      new LoggingContentGenerator(
+        (mockGenerator as GoogleGenAI).models,
+        mockConfig,
+      ),
+    );
   });
 });
 

packages/core/src/core/contentGenerator.ts

@@ -18,6 +18,7 @@ import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
 import { Config } from '../config/config.js';
 import { getEffectiveModel } from './modelCheck.js';
 import { UserTierId } from '../code_assist/types.js';
+import { LoggingContentGenerator } from './loggingContentGenerator.js';
 
 /**
  * Interface abstracting the core functionalities for generating content and counting tokens.
@@ -121,11 +122,14 @@ export async function createContentGenerator(
     config.authType === AuthType.LOGIN_WITH_GOOGLE ||
     config.authType === AuthType.CLOUD_SHELL
   ) {
-    return createCodeAssistContentGenerator(
+    return new LoggingContentGenerator(
+      await createCodeAssistContentGenerator(
         httpOptions,
         config.authType,
         gcConfig,
         sessionId,
+      ),
+      gcConfig,
     );
   }
 
@@ -138,10 +142,8 @@ export async function createContentGenerator(
       vertexai: config.vertexai,
       httpOptions,
     });
-
-    return googleGenAI.models;
+    return new LoggingContentGenerator(googleGenAI.models, gcConfig);
   }
-
   throw new Error(
     `Error creating contentGenerator: Unsupported authType: ${config.authType}`,
   );

packages/core/src/core/geminiChat.ts

@@ -21,16 +21,8 @@ import { retryWithBackoff } from '../utils/retry.js';
 import { isFunctionResponse } from '../utils/messageInspectors.js';
 import { ContentGenerator, AuthType } from './contentGenerator.js';
 import { Config } from '../config/config.js';
-import {
-  logApiRequest,
-  logApiResponse,
-  logApiError,
-} from '../telemetry/loggers.js';
-import {
-  ApiErrorEvent,
-  ApiRequestEvent,
-  ApiResponseEvent,
-} from '../telemetry/types.js';
+import { logApiResponse, logApiError } from '../telemetry/loggers.js';
+import { ApiErrorEvent, ApiResponseEvent } from '../telemetry/types.js';
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { hasCycleInSchema } from '../tools/tools.js';
 import { StructuredError } from './turn.js';
@@ -139,22 +131,6 @@ export class GeminiChat {
     validateHistory(history);
   }
 
-  private _getRequestTextFromContents(contents: Content[]): string {
-    return JSON.stringify(contents);
-  }
-
-  private async _logApiRequest(
-    contents: Content[],
-    model: string,
-    prompt_id: string,
-  ): Promise<void> {
-    const requestText = this._getRequestTextFromContents(contents);
-    logApiRequest(
-      this.config,
-      new ApiRequestEvent(model, prompt_id, requestText),
-    );
-  }
-
   private async _logApiResponse(
     durationMs: number,
     prompt_id: string,
@@ -273,8 +249,6 @@ export class GeminiChat {
     const userContent = createUserContent(params.message);
     const requestContents = this.getHistory(true).concat(userContent);
 
-    this._logApiRequest(requestContents, this.config.getModel(), prompt_id);
-
     const startTime = Date.now();
     let response: GenerateContentResponse;
 
@@ -386,7 +360,6 @@ export class GeminiChat {
     await this.sendPromise;
     const userContent = createUserContent(params.message);
     const requestContents = this.getHistory(true).concat(userContent);
-    this._logApiRequest(requestContents, this.config.getModel(), prompt_id);
 
     const startTime = Date.now();
 

packages/core/src/core/loggingContentGenerator.ts

@@ -0,0 +1,68 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  Content,
+  CountTokensParameters,
+  CountTokensResponse,
+  EmbedContentParameters,
+  EmbedContentResponse,
+  GenerateContentParameters,
+  GenerateContentResponse,
+} from '@google/genai';
+import { ApiRequestEvent } from '../telemetry/types.js';
+import { Config } from '../config/config.js';
+import { logApiRequest } from '../telemetry/loggers.js';
+import { ContentGenerator } from './contentGenerator.js';
+import { toContents } from '../code_assist/converter.js';
+
+/**
+ * A decorator that wraps a ContentGenerator to add logging to API calls.
+ */
+export class LoggingContentGenerator implements ContentGenerator {
+  constructor(
+    private readonly wrapped: ContentGenerator,
+    private readonly config: Config,
+  ) {}
+
+  private logApiRequest(
+    contents: Content[],
+    model: string,
+    promptId: string,
+  ): void {
+    const requestText = JSON.stringify(contents);
+    logApiRequest(
+      this.config,
+      new ApiRequestEvent(model, promptId, requestText),
+    );
+  }
+
+  async generateContent(
+    req: GenerateContentParameters,
+    userPromptId: string,
+  ): Promise<GenerateContentResponse> {
+    this.logApiRequest(toContents(req.contents), req.model, userPromptId);
+    return this.wrapped.generateContent(req, userPromptId);
+  }
+
+  async generateContentStream(
+    req: GenerateContentParameters,
+    userPromptId: string,
+  ): Promise<AsyncGenerator<GenerateContentResponse>> {
+    this.logApiRequest(toContents(req.contents), req.model, userPromptId);
+    return this.wrapped.generateContentStream(req, userPromptId);
+  }
+
+  async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {
+    return this.wrapped.countTokens(req);
+  }
+
+  async embedContent(
+    req: EmbedContentParameters,
+  ): Promise<EmbedContentResponse> {
+    return this.wrapped.embedContent(req);
+  }
+}

packages/core/src/core/nonInteractiveToolExecutor.ts

@@ -14,7 +14,7 @@ import {
 } from '../index.js';
 import { Config } from '../config/config.js';
 import { convertToFunctionResponse } from './coreToolScheduler.js';
-import { ToolCallDecision } from '../telemetry/types.js';
+import { ToolCallDecision } from '../telemetry/tool-call-decision.js';
 
 /**
  * Executes a single tool call non-interactively.

packages/core/src/telemetry/loggers.test.ts

@@ -35,11 +35,11 @@ import {
   logToolCall,
   logFlashFallback,
 } from './loggers.js';
+import { ToolCallDecision } from './tool-call-decision.js';
 import {
   ApiRequestEvent,
   ApiResponseEvent,
   StartSessionEvent,
-  ToolCallDecision,
   ToolCallEvent,
   UserPromptEvent,
   FlashFallbackEvent,

packages/core/src/telemetry/tool-call-decision.ts

@@ -0,0 +1,32 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { ToolConfirmationOutcome } from '../tools/tools.js';
+
+export enum ToolCallDecision {
+  ACCEPT = 'accept',
+  REJECT = 'reject',
+  MODIFY = 'modify',
+  AUTO_ACCEPT = 'auto_accept',
+}
+
+export function getDecisionFromOutcome(
+  outcome: ToolConfirmationOutcome,
+): ToolCallDecision {
+  switch (outcome) {
+    case ToolConfirmationOutcome.ProceedOnce:
+      return ToolCallDecision.ACCEPT;
+    case ToolConfirmationOutcome.ProceedAlways:
+    case ToolConfirmationOutcome.ProceedAlwaysServer:
+    case ToolConfirmationOutcome.ProceedAlwaysTool:
+      return ToolCallDecision.AUTO_ACCEPT;
+    case ToolConfirmationOutcome.ModifyWithEditor:
+      return ToolCallDecision.MODIFY;
+    case ToolConfirmationOutcome.Cancel:
+    default:
+      return ToolCallDecision.REJECT;
+  }
+}

packages/core/src/telemetry/types.ts

@@ -7,33 +7,11 @@
 import { GenerateContentResponseUsageMetadata } from '@google/genai';
 import { Config } from '../config/config.js';
 import { CompletedToolCall } from '../core/coreToolScheduler.js';
-import { ToolConfirmationOutcome } from '../tools/tools.js';
 import { AuthType } from '../core/contentGenerator.js';
-
-export enum ToolCallDecision {
-  ACCEPT = 'accept',
-  REJECT = 'reject',
-  MODIFY = 'modify',
-  AUTO_ACCEPT = 'auto_accept',
-}
-
-export function getDecisionFromOutcome(
-  outcome: ToolConfirmationOutcome,
-): ToolCallDecision {
-  switch (outcome) {
-    case ToolConfirmationOutcome.ProceedOnce:
-      return ToolCallDecision.ACCEPT;
-    case ToolConfirmationOutcome.ProceedAlways:
-    case ToolConfirmationOutcome.ProceedAlwaysServer:
-    case ToolConfirmationOutcome.ProceedAlwaysTool:
-      return ToolCallDecision.AUTO_ACCEPT;
-    case ToolConfirmationOutcome.ModifyWithEditor:
-      return ToolCallDecision.MODIFY;
-    case ToolConfirmationOutcome.Cancel:
-    default:
-      return ToolCallDecision.REJECT;
-  }
-}
+import {
+  getDecisionFromOutcome,
+  ToolCallDecision,
+} from './tool-call-decision.js';
 
 export class StartSessionEvent {
   'event.name': 'cli_config';

packages/core/src/telemetry/uiTelemetry.test.ts

@@ -6,12 +6,8 @@
 
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { UiTelemetryService } from './uiTelemetry.js';
-import {
-  ApiErrorEvent,
-  ApiResponseEvent,
-  ToolCallEvent,
-  ToolCallDecision,
-} from './types.js';
+import { ToolCallDecision } from './tool-call-decision.js';
+import { ApiErrorEvent, ApiResponseEvent, ToolCallEvent } from './types.js';
 import {
   EVENT_API_ERROR,
   EVENT_API_RESPONSE,

packages/core/src/telemetry/uiTelemetry.ts

@@ -11,12 +11,8 @@ import {
   EVENT_TOOL_CALL,
 } from './constants.js';
 
-import {
-  ApiErrorEvent,
-  ApiResponseEvent,
-  ToolCallEvent,
-  ToolCallDecision,
-} from './types.js';
+import { ToolCallDecision } from './tool-call-decision.js';
+import { ApiErrorEvent, ApiResponseEvent, ToolCallEvent } from './types.js';
 
 export type UiEvent =
   | (ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE })