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

Co-authored-by: Shi Shu <[email protected]>

3 files changed, 61 insertions, 2 deletions

diff --git a/packages/core/src/code_assist/converter.test.ts b/packages/core/src/code_assist/converter.test.ts
index 3d3a8ef3..ab040fe6 100644
--- a/packages/core/src/code_assist/converter.test.ts
+++ b/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' }] },
+      ]);
+    });
+  });
 });
diff --git a/packages/core/src/code_assist/converter.ts b/packages/core/src/code_assist/converter.ts
index ffd471da..a71c70da 100644
--- a/packages/core/src/code_assist/converter.ts
+++ b/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);
diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts
index 3fc1891f..bc3978fb 100644
--- a/packages/core/src/code_assist/server.test.ts
+++ b/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(