From 8563e46ade1e4d72651f35d05bcfbdbf839bb41c Mon Sep 17 00:00:00 2001
From: "N. Taylor Mullen" <ntaylormullen@google.com>
Date: Mon, 2 Jun 2025 16:32:45 -0700
Subject: [PATCH] React to Gemini API break - Thought Inclusion (#705)

---
 packages/core/src/core/geminiChat.test.ts | 74 +++++++++++++++++++++++
 packages/core/src/core/geminiChat.ts      | 17 ++++++
 2 files changed, 91 insertions(+)

diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts
index 8d434dd4..3a6fb10c 100644
--- a/packages/core/src/core/geminiChat.test.ts
+++ b/packages/core/src/core/geminiChat.test.ts
@@ -298,5 +298,79 @@ describe('GeminiChat', () => {
       expect(history[0]).toEqual(userInput);
       expect(history[1]).toEqual(modelOutput[0]);
     });
+
+    it('should skip "thought" content from modelOutput', () => {
+      const modelOutputWithThought: Content[] = [
+        { role: 'model', parts: [{ thought: true }, { text: 'Visible text' }] },
+        { role: 'model', parts: [{ text: 'Another visible text' }] },
+      ];
+      // @ts-expect-error Accessing private method for testing purposes
+      chat.recordHistory(userInput, modelOutputWithThought);
+      const history = chat.getHistory();
+      expect(history.length).toBe(2); // User input + consolidated model output
+      expect(history[0]).toEqual(userInput);
+      expect(history[1].role).toBe('model');
+      // The 'thought' part is skipped, 'Another visible text' becomes the first part.
+      expect(history[1].parts).toEqual([{ text: 'Another visible text' }]);
+    });
+
+    it('should skip "thought" content even if it is the only content', () => {
+      const modelOutputOnlyThought: Content[] = [
+        { role: 'model', parts: [{ thought: true }] },
+      ];
+      // @ts-expect-error Accessing private method for testing purposes
+      chat.recordHistory(userInput, modelOutputOnlyThought);
+      const history = chat.getHistory();
+      expect(history.length).toBe(1); // User input + default empty model part
+      expect(history[0]).toEqual(userInput);
+    });
+
+    it('should correctly consolidate text parts when a thought part is in between', () => {
+      const modelOutputMixed: Content[] = [
+        { role: 'model', parts: [{ text: 'Part 1.' }] },
+        {
+          role: 'model',
+          parts: [{ thought: true }, { text: 'Should be skipped' }],
+        },
+        { role: 'model', parts: [{ text: 'Part 2.' }] },
+      ];
+      // @ts-expect-error Accessing private method for testing purposes
+      chat.recordHistory(userInput, modelOutputMixed);
+      const history = chat.getHistory();
+      expect(history.length).toBe(2);
+      expect(history[0]).toEqual(userInput);
+      expect(history[1].role).toBe('model');
+      expect(history[1].parts).toEqual([{ text: 'Part 1.Part 2.' }]);
+    });
+
+    it('should handle multiple thought parts correctly', () => {
+      const modelOutputMultipleThoughts: Content[] = [
+        { role: 'model', parts: [{ thought: true }] },
+        { role: 'model', parts: [{ text: 'Visible 1' }] },
+        { role: 'model', parts: [{ thought: true }] },
+        { role: 'model', parts: [{ text: 'Visible 2' }] },
+      ];
+      // @ts-expect-error Accessing private method for testing purposes
+      chat.recordHistory(userInput, modelOutputMultipleThoughts);
+      const history = chat.getHistory();
+      expect(history.length).toBe(2);
+      expect(history[0]).toEqual(userInput);
+      expect(history[1].role).toBe('model');
+      expect(history[1].parts).toEqual([{ text: 'Visible 1Visible 2' }]);
+    });
+
+    it('should handle thought part at the end of outputContents', () => {
+      const modelOutputThoughtAtEnd: Content[] = [
+        { role: 'model', parts: [{ text: 'Visible text' }] },
+        { role: 'model', parts: [{ thought: true }] },
+      ];
+      // @ts-expect-error Accessing private method for testing purposes
+      chat.recordHistory(userInput, modelOutputThoughtAtEnd);
+      const history = chat.getHistory();
+      expect(history.length).toBe(2);
+      expect(history[0]).toEqual(userInput);
+      expect(history[1].role).toBe('model');
+      expect(history[1].parts).toEqual([{ text: 'Visible text' }]);
+    });
   });
 });
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index 25ea6e52..b4844499 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -342,6 +342,10 @@ export class GeminiChat {
     // Consolidate adjacent model roles in outputContents
     const consolidatedOutputContents: Content[] = [];
     for (const content of outputContents) {
+      if (this.isThoughtContent(content)) {
+        continue;
+      }
+
       const lastContent =
         consolidatedOutputContents[consolidatedOutputContents.length - 1];
       if (this.isTextContent(lastContent) && this.isTextContent(content)) {
@@ -394,4 +398,17 @@ export class GeminiChat {
       content.parts[0].text !== ''
     );
   }
+
+  private isThoughtContent(
+    content: Content | undefined,
+  ): content is Content & { parts: [{ thought: boolean }, ...Part[]] } {
+    return !!(
+      content &&
+      content.role === 'model' &&
+      content.parts &&
+      content.parts.length > 0 &&
+      typeof content.parts[0].thought === 'boolean' &&
+      content.parts[0].thought === true
+    );
+  }
 }