feat(commands): Enable @file processing in TOML commands (#6716)

packages/core/src/utils/partUtils.test.ts

@@ -5,8 +5,13 @@
  */
 
 import { describe, it, expect } from 'vitest';
-import { partToString, getResponseText } from './partUtils.js';
-import type { GenerateContentResponse, Part } from '@google/genai';
+import {
+  partToString,
+  getResponseText,
+  flatMapTextParts,
+  appendToLastTextPart,
+} from './partUtils.js';
+import type { GenerateContentResponse, Part, PartUnion } from '@google/genai';
 
 const mockResponse = (
   parts?: Array<{ text?: string; functionCall?: unknown }>,
@@ -162,5 +167,135 @@ describe('partUtils', () => {
       const result = mockResponse([]);
       expect(getResponseText(result)).toBeNull();
     });
+
+    it('should return null if the first candidate has no content property', () => {
+      const response: GenerateContentResponse = {
+        candidates: [
+          {
+            index: 0,
+          },
+        ],
+        promptFeedback: { safetyRatings: [] },
+        text: undefined,
+        data: undefined,
+        functionCalls: undefined,
+        executableCode: undefined,
+        codeExecutionResult: undefined,
+      };
+      expect(getResponseText(response)).toBeNull();
+    });
+  });
+
+  describe('flatMapTextParts', () => {
+    // A simple async transform function that splits a string into character parts.
+    const splitCharsTransform = async (text: string): Promise<PartUnion[]> =>
+      text.split('').map((char) => ({ text: char }));
+
+    it('should return an empty array for empty input', async () => {
+      const result = await flatMapTextParts([], splitCharsTransform);
+      expect(result).toEqual([]);
+    });
+
+    it('should transform a simple string input', async () => {
+      const result = await flatMapTextParts('hi', splitCharsTransform);
+      expect(result).toEqual([{ text: 'h' }, { text: 'i' }]);
+    });
+
+    it('should transform a single text part object', async () => {
+      const result = await flatMapTextParts(
+        { text: 'cat' },
+        splitCharsTransform,
+      );
+      expect(result).toEqual([{ text: 'c' }, { text: 'a' }, { text: 't' }]);
+    });
+
+    it('should transform an array of text parts and flatten the result', async () => {
+      // A transform that duplicates the text to test the "flatMap" behavior.
+      const duplicateTransform = async (text: string): Promise<PartUnion[]> => [
+        { text: `${text}` },
+        { text: `${text}` },
+      ];
+      const parts = [{ text: 'a' }, { text: 'b' }];
+      const result = await flatMapTextParts(parts, duplicateTransform);
+      expect(result).toEqual([
+        { text: 'a' },
+        { text: 'a' },
+        { text: 'b' },
+        { text: 'b' },
+      ]);
+    });
+
+    it('should pass through non-text parts unmodified', async () => {
+      const nonTextPart: Part = { functionCall: { name: 'do_stuff' } };
+      const result = await flatMapTextParts(nonTextPart, splitCharsTransform);
+      expect(result).toEqual([nonTextPart]);
+    });
+
+    it('should handle a mix of text and non-text parts in an array', async () => {
+      const nonTextPart: Part = {
+        inlineData: { mimeType: 'image/jpeg', data: '' },
+      };
+      const parts: PartUnion[] = [{ text: 'go' }, nonTextPart, ' stop'];
+      const result = await flatMapTextParts(parts, splitCharsTransform);
+      expect(result).toEqual([
+        { text: 'g' },
+        { text: 'o' },
+        nonTextPart, // Should be passed through
+        { text: ' ' },
+        { text: 's' },
+        { text: 't' },
+        { text: 'o' },
+        { text: 'p' },
+      ]);
+    });
+
+    it('should handle a transform that returns an empty array', async () => {
+      const removeTransform = async (_text: string): Promise<PartUnion[]> => [];
+      const parts: PartUnion[] = [
+        { text: 'remove' },
+        { functionCall: { name: 'keep' } },
+      ];
+      const result = await flatMapTextParts(parts, removeTransform);
+      expect(result).toEqual([{ functionCall: { name: 'keep' } }]);
+    });
+  });
+
+  describe('appendToLastTextPart', () => {
+    it('should append to an empty prompt', () => {
+      const prompt: PartUnion[] = [];
+      const result = appendToLastTextPart(prompt, 'new text');
+      expect(result).toEqual([{ text: 'new text' }]);
+    });
+
+    it('should append to a prompt with a string as the last part', () => {
+      const prompt: PartUnion[] = ['first part'];
+      const result = appendToLastTextPart(prompt, 'new text');
+      expect(result).toEqual(['first part\n\nnew text']);
+    });
+
+    it('should append to a prompt with a text part object as the last part', () => {
+      const prompt: PartUnion[] = [{ text: 'first part' }];
+      const result = appendToLastTextPart(prompt, 'new text');
+      expect(result).toEqual([{ text: 'first part\n\nnew text' }]);
+    });
+
+    it('should append a new text part if the last part is not a text part', () => {
+      const nonTextPart: Part = { functionCall: { name: 'do_stuff' } };
+      const prompt: PartUnion[] = [nonTextPart];
+      const result = appendToLastTextPart(prompt, 'new text');
+      expect(result).toEqual([nonTextPart, { text: '\n\nnew text' }]);
+    });
+
+    it('should not append anything if the text to append is empty', () => {
+      const prompt: PartUnion[] = ['first part'];
+      const result = appendToLastTextPart(prompt, '');
+      expect(result).toEqual(['first part']);
+    });
+
+    it('should use a custom separator', () => {
+      const prompt: PartUnion[] = ['first part'];
+      const result = appendToLastTextPart(prompt, 'new text', '---');
+      expect(result).toEqual(['first part---new text']);
+    });
   });
 });