fix: Refine model message consolidation for improved model interaction (#685)

2 files changed, 67 insertions, 30 deletions

diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts
index 11e222c9..8d434dd4 100644
--- a/packages/core/src/core/geminiChat.test.ts
+++ b/packages/core/src/core/geminiChat.test.ts
@@ -69,10 +69,7 @@ describe('GeminiChat', () => {
       expect(history.length).toBe(2);
       expect(history[0]).toEqual(userInput);
       expect(history[1].role).toBe('model');
-      expect(history[1].parts).toEqual([
-        { text: 'Model part 1' },
-        { text: 'Model part 2' },
-      ]);
+      expect(history[1].parts).toEqual([{ text: 'Model part 1Model part 2' }]);
     });
 
     it('should handle a mix of user and model roles in outputContents (though unusual)', () => {
@@ -101,11 +98,7 @@ describe('GeminiChat', () => {
       chat.recordHistory(userInput, modelOutputParts);
       const history = chat.getHistory();
       expect(history.length).toBe(2);
-      expect(history[1].parts).toEqual([
-        { text: 'M1' },
-        { text: 'M2' },
-        { text: 'M3' },
-      ]);
+      expect(history[1].parts).toEqual([{ text: 'M1M2M3' }]);
     });
 
     it('should not consolidate if roles are different between model outputs', () => {
@@ -177,8 +170,7 @@ describe('GeminiChat', () => {
       expect(finalHistory[2]).toEqual(secondUserInput);
       expect(finalHistory[3].role).toBe('model');
       expect(finalHistory[3].parts).toEqual([
-        { text: 'Second model response part 1' },
-        { text: 'Second model response part 2' },
+        { text: 'Second model response part 1Second model response part 2' },
       ]);
     });
 
@@ -218,8 +210,7 @@ describe('GeminiChat', () => {
       expect(history[2]).toEqual(currentUserInput);
       expect(history[3].role).toBe('model');
       expect(history[3].parts).toEqual([
-        { text: 'Part A of new answer.' },
-        { text: 'Part B of new answer.' },
+        { text: 'Part A of new answer.Part B of new answer.' },
       ]);
     });
 
@@ -235,6 +226,31 @@ describe('GeminiChat', () => {
       expect(history[1].parts).toEqual([]);
     });
 
+    it('should handle aggregating modelOutput', () => {
+      const modelOutputUndefinedParts: Content[] = [
+        { role: 'model', parts: [{ text: 'First model part' }] },
+        { role: 'model', parts: [{ text: 'Second model part' }] },
+        { role: 'model', parts: undefined as unknown as Part[] }, // Test undefined parts
+        { role: 'model', parts: [{ text: 'Third model part' }] },
+        { role: 'model', parts: [] }, // Test empty parts array
+      ];
+      // @ts-expect-error Accessing private method for testing purposes
+      chat.recordHistory(userInput, modelOutputUndefinedParts);
+      const history = chat.getHistory();
+      expect(history.length).toBe(5);
+      expect(history[0]).toEqual(userInput);
+      expect(history[1].role).toBe('model');
+      expect(history[1].parts).toEqual([
+        { text: 'First model partSecond model part' },
+      ]);
+      expect(history[2].role).toBe('model');
+      expect(history[2].parts).toBeUndefined();
+      expect(history[3].role).toBe('model');
+      expect(history[3].parts).toEqual([{ text: 'Third model part' }]);
+      expect(history[4].role).toBe('model');
+      expect(history[4].parts).toEqual([]);
+    });
+
     it('should handle modelOutput with parts being undefined or empty (if they pass initial every check)', () => {
       const modelOutputUndefinedParts: Content[] = [
         { role: 'model', parts: [{ text: 'Text part' }] },
@@ -244,10 +260,14 @@ describe('GeminiChat', () => {
       // @ts-expect-error Accessing private method for testing purposes
       chat.recordHistory(userInput, modelOutputUndefinedParts);
       const history = chat.getHistory();
-      expect(history.length).toBe(2);
+      expect(history.length).toBe(4); // userInput, model1 (text), model2 (undefined parts), model3 (empty parts)
+      expect(history[0]).toEqual(userInput);
       expect(history[1].role).toBe('model');
-      // The consolidation logic should handle undefined/empty parts by spreading `|| []`
       expect(history[1].parts).toEqual([{ text: 'Text part' }]);
+      expect(history[2].role).toBe('model');
+      expect(history[2].parts).toBeUndefined();
+      expect(history[3].role).toBe('model');
+      expect(history[3].parts).toEqual([]);
     });
 
     it('should correctly handle automaticFunctionCallingHistory', () => {
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index b34b6f35..25ea6e52 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -15,6 +15,7 @@ import {
   SendMessageParameters,
   GoogleGenAI,
   createUserContent,
+  Part,
 } from '@google/genai';
 import { retryWithBackoff } from '../utils/retry.js';
 import { isFunctionResponse } from '../utils/messageInspectors.js';
@@ -343,13 +344,13 @@ export class GeminiChat {
     for (const content of outputContents) {
       const lastContent =
         consolidatedOutputContents[consolidatedOutputContents.length - 1];
-      if (
-        lastContent &&
-        lastContent.role === 'model' &&
-        content.role === 'model' &&
-        lastContent.parts
-      ) {
-        lastContent.parts.push(...(content.parts || []));
+      if (this.isTextContent(lastContent) && this.isTextContent(content)) {
+        // If both current and last are text, combine their text into the lastContent's first part
+        // and append any other parts from the current content.
+        lastContent.parts[0].text += content.parts[0].text || '';
+        if (content.parts.length > 1) {
+          lastContent.parts.push(...content.parts.slice(1));
+        }
       } else {
         consolidatedOutputContents.push(content);
       }
@@ -357,24 +358,40 @@ export class GeminiChat {
 
     if (consolidatedOutputContents.length > 0) {
       const lastHistoryEntry = this.history[this.history.length - 1];
-      // Only merge if AFC history was NOT just added, to prevent merging with last AFC model turn.
       const canMergeWithLastHistory =
         !automaticFunctionCallingHistory ||
         automaticFunctionCallingHistory.length === 0;
 
       if (
         canMergeWithLastHistory &&
-        lastHistoryEntry &&
-        lastHistoryEntry.role === 'model' &&
-        lastHistoryEntry.parts &&
-        consolidatedOutputContents[0].role === 'model'
+        this.isTextContent(lastHistoryEntry) &&
+        this.isTextContent(consolidatedOutputContents[0])
       ) {
-        lastHistoryEntry.parts.push(
-          ...(consolidatedOutputContents[0].parts || []),
-        );
+        // If both current and last are text, combine their text into the lastHistoryEntry's first part
+        // and append any other parts from the current content.
+        lastHistoryEntry.parts[0].text +=
+          consolidatedOutputContents[0].parts[0].text || '';
+        if (consolidatedOutputContents[0].parts.length > 1) {
+          lastHistoryEntry.parts.push(
+            ...consolidatedOutputContents[0].parts.slice(1),
+          );
+        }
         consolidatedOutputContents.shift(); // Remove the first element as it's merged
       }
       this.history.push(...consolidatedOutputContents);
     }
   }
+
+  private isTextContent(
+    content: Content | undefined,
+  ): content is Content & { parts: [{ text: string }, ...Part[]] } {
+    return !!(
+      content &&
+      content.role === 'model' &&
+      content.parts &&
+      content.parts.length > 0 &&
+      typeof content.parts[0].text === 'string' &&
+      content.parts[0].text !== ''
+    );
+  }
 }