From 27a3da3e88edddb00a03dfb4de24aee14432dc34 Mon Sep 17 00:00:00 2001 From: amelidev Date: Tue, 7 Jul 2026 12:03:31 -0600 Subject: [PATCH] fix(core): strip thoughts from scrubbed history turns and resolve thought leakage (#27971) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Gal Zahavi <38544478+galz10@users.noreply.github.com> Co-authored-by: David Pierce --- packages/core/src/core/geminiChat.test.ts | 30 +++ .../core/src/utils/historyHardening.test.ts | 205 +++++++++++++++++- packages/core/src/utils/historyHardening.ts | 115 ++++++++-- 3 files changed, 335 insertions(+), 15 deletions(-) diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index b6f9ef9886..07fe800f3b 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -10,6 +10,7 @@ import { ThinkingLevel, type Content, type GenerateContentResponse, + type Part, } from '@google/genai'; import type { ContentGenerator } from '../core/contentGenerator.js'; import { @@ -2253,6 +2254,35 @@ describe('GeminiChat', () => { }); }); + describe('thought leakage in getHistoryTurns', () => { + it('should completely filter out thought parts from getHistoryTurns when context management is enabled', () => { + vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(true); + + chat.setHistory([ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + { + role: 'model', + parts: [ + { text: 'internal monologue', thought: true } as unknown as Part, + { text: 'actual conversational response' }, + ], + }, + ]); + + const turns = chat.getHistoryTurns(true); + + expect(turns).toHaveLength(2); + const modelTurn = turns[1]; + expect(modelTurn.content.parts).toHaveLength(1); + expect(modelTurn.content.parts![0]).toEqual({ + text: 'actual conversational response', + }); + }); + }); + describe('ensureActiveLoopHasThoughtSignatures', () => { it('should add thoughtSignature to the first functionCall in each model turn of the active loop', () => { const chat = new GeminiChat(mockConfig, '', [], []); diff --git a/packages/core/src/utils/historyHardening.test.ts b/packages/core/src/utils/historyHardening.test.ts index face667132..90577e1e4b 100644 --- a/packages/core/src/utils/historyHardening.test.ts +++ b/packages/core/src/utils/historyHardening.test.ts @@ -8,10 +8,12 @@ import { describe, it, expect } from 'vitest'; import { hardenHistory, SYNTHETIC_THOUGHT_SIGNATURE, + scrubContents, + scrubHistory, } from './historyHardening.js'; import type { HistoryTurn } from '../core/agentChatHistory.js'; import { deriveStableId } from './cryptoUtils.js'; -import type { Part } from '@google/genai'; +import type { Part, Content } from '@google/genai'; describe('hardenHistory', () => { it('should return an empty array if input is empty', () => { @@ -375,4 +377,205 @@ describe('hardenHistory', () => { expect(hardened[0].content.parts![0]).not.toHaveProperty('extraProp'); expect(hardened[0].content.parts![0]).toHaveProperty('text', 'hello'); }); + + it('should completely filter out thought parts from the scrubbed history', () => { + const history: HistoryTurn[] = [ + { + id: '1', + content: { + role: 'user', + parts: [{ text: 'User prompt' }], + }, + }, + { + id: '2', + content: { + role: 'model', + parts: [ + { + text: 'Previous model thought...', + thought: true, + } as unknown as Part, + { text: 'Actual conversational text response' }, + ], + }, + }, + { + id: '3', + content: { + role: 'user', + parts: [{ text: 'User follow-up prompt' }], + }, + }, + ]; + + const hardened = hardenHistory(history); + // Model turn (Turn 2, index 1 in hardened) should only contain the actual conversational text part + const modelTurn = hardened[1]; + expect(modelTurn.content.parts).toHaveLength(1); + expect(modelTurn.content.parts![0]).toHaveProperty( + 'text', + 'Actual conversational text response', + ); + expect(modelTurn.content.parts![0]).not.toHaveProperty('thought'); + }); + + it('should remove the entire turn if it only contained thought parts and is now empty', () => { + const history: HistoryTurn[] = [ + { + id: '1', + content: { + role: 'user', + parts: [{ text: 'User prompt' }], + }, + }, + { + id: '2', + content: { + role: 'model', + parts: [ + { + text: 'Model is just thinking internally...', + thought: true, + } as unknown as Part, + ], + }, + }, + { + id: '3', + content: { + role: 'user', + parts: [{ text: 'User follow-up prompt' }], + }, + }, + ]; + + const hardened = hardenHistory(history); + // After scrubbing, Turn 2 should have 0 parts. + // The history mapping filters out empty turns, so the total turns should coalesce and reduce to 1 coalesced user turn. + // Let's inspect the hardened array: + // User prompt (Turn 1) + User follow-up prompt (Turn 3) will be coalesced into 1 User turn. + expect(hardened).toHaveLength(1); + expect(hardened[0].content.role).toBe('user'); + expect(hardened[0].content.parts).toEqual([ + { text: 'User prompt' }, + { text: 'User follow-up prompt' }, + ]); + }); +}); + +describe('scrubContents', () => { + it('should scrub non-standard fields from parts', () => { + const contents: Content[] = [ + { + role: 'user', + parts: [{ text: 'Hello', customField: 'ignored' } as unknown as Part], + }, + ]; + const scrubbed = scrubContents(contents); + expect(scrubbed).toEqual([ + { + role: 'user', + parts: [{ text: 'Hello' }], + }, + ]); + }); + + it('should filter out internal thought parts', () => { + const contents: Content[] = [ + { + role: 'model', + parts: [ + { text: 'thought', thought: true } as unknown as Part, + { text: 'response' }, + ], + }, + ]; + const scrubbed = scrubContents(contents); + expect(scrubbed).toEqual([ + { + role: 'model', + parts: [{ text: 'response' }], + }, + ]); + }); + + it('should completely filter out Content objects that have no parts left after thought scrubbing and coalesce adjacent turns of the same role', () => { + const contents: Content[] = [ + { + role: 'user', + parts: [{ text: 'Hello' }], + }, + { + role: 'model', + parts: [{ text: 'thought', thought: true } as unknown as Part], + }, + { + role: 'user', + parts: [{ text: 'How are you?' }], + }, + ]; + const scrubbed = scrubContents(contents); + expect(scrubbed).toEqual([ + { + role: 'user', + parts: [{ text: 'Hello' }, { text: 'How are you?' }], + }, + ]); + }); + + it('should coalesce adjacent turns of the same role when no filtration occurs', () => { + const contents: Content[] = [ + { + role: 'user', + parts: [{ text: 'Part 1' }], + }, + { + role: 'user', + parts: [{ text: 'Part 2' }], + }, + ]; + const scrubbed = scrubContents(contents); + expect(scrubbed).toEqual([ + { + role: 'user', + parts: [{ text: 'Part 1' }, { text: 'Part 2' }], + }, + ]); + }); +}); + +describe('scrubHistory', () => { + it('should scrub non-standard fields and filter empty turns in history', () => { + const history: HistoryTurn[] = [ + { + id: '1', + content: { + role: 'user', + parts: [{ text: 'Hello', customField: 'ignored' } as unknown as Part], + }, + }, + { + id: '2', + content: { + role: 'model', + parts: [{ text: 'thought', thought: true } as unknown as Part], + }, + }, + { + id: '3', + content: { + role: 'user', + parts: [{ text: 'World' }], + }, + }, + ]; + + const scrubbed = scrubHistory(history); + expect(scrubbed.length).toBe(1); // Since user turns are coalesced (Turn 1 + Turn 3) and Turn 2 is removed because it has 0 parts + expect(scrubbed[0].content.parts).toEqual([ + { text: 'Hello' }, + { text: 'World' }, + ]); + }); }); diff --git a/packages/core/src/utils/historyHardening.ts b/packages/core/src/utils/historyHardening.ts index e469b08e83..8a4e8a2f5c 100644 --- a/packages/core/src/utils/historyHardening.ts +++ b/packages/core/src/utils/historyHardening.ts @@ -44,8 +44,11 @@ export function hardenHistory( const sentinels = { ...DEFAULT_SENTINELS, ...options.sentinels }; + // Pass 0: Strip internal thoughts and remove empty turns + const processed = stripThoughts(history); + // Pass 1: Initial Coalesce & Empty Turn Removal - let coalesced = coalesce(history); + let coalesced = coalesce(processed); // Pass 2: Tool Pairing & Signatures (The semantic layer) coalesced = pairToolsAndEnforceSignatures(coalesced, sentinels); @@ -62,6 +65,36 @@ export function hardenHistory( return final; } +/** + * Helper to check if a Part object represents an internal thought. + */ +function isInternalThought(part: Part): boolean { + return !!part && !!(part as ThoughtPart).thought; +} + +/** + * Removes parts that represent thoughts (where part.thought === true). + * Empty turns resulting from thought removal are handled in subsequent coalescing passes. + */ +function stripThoughts(history: HistoryTurn[]): HistoryTurn[] { + return history.map((turn) => { + if (!turn.content.parts) return turn; + const hasThought = turn.content.parts.some(isInternalThought); + if (!hasThought) return turn; + + const nonThoughtParts = turn.content.parts.filter( + (p) => p && !isInternalThought(p), + ); + return { + id: turn.id, + content: { + ...turn.content, + parts: nonThoughtParts, + }, + }; + }); +} + /** * Combines adjacent turns with the same role and removes empty turns. */ @@ -70,12 +103,16 @@ function coalesce(history: HistoryTurn[]): HistoryTurn[] { for (const turn of history) { if (!turn.content.parts || turn.content.parts.length === 0) continue; - const last = result[result.length - 1]; + const lastIdx = result.length - 1; + const last = result[lastIdx]; if (last && last.content.role === turn.content.role) { - last.content.parts = [ - ...(last.content.parts || []), - ...(turn.content.parts || []), - ]; + result[lastIdx] = { + id: last.id, + content: { + ...last.content, + parts: [...(last.content.parts || []), ...(turn.content.parts || [])], + }, + }; } else { // Shallow clone the turn and content so we don't mutate the original history array structure result.push({ id: turn.id, content: { ...turn.content } }); @@ -344,23 +381,73 @@ function enforceRoleConstraints( * This ensures compatibility with strict APIs (like Vertex AI) that reject unknown fields. */ export function scrubHistory(history: HistoryTurn[]): HistoryTurn[] { - return history.map((turn) => ({ - id: turn.id, - content: scrubContents([turn.content])[0], - })); + const result: HistoryTurn[] = []; + for (const turn of history) { + const nonThoughtParts = (turn.content.parts ?? []).filter( + (p) => p && !isInternalThought(p), + ); + if (nonThoughtParts.length === 0) continue; // Skip turns that became empty + + const scrubbedParts = nonThoughtParts.map((p) => scrubPart(p)); + + const lastIdx = result.length - 1; + const last = result[lastIdx]; + if (last && last.content.role === turn.content.role) { + // Coalesce inline with strict immutability + result[lastIdx] = { + id: last.id, + content: { + ...last.content, + parts: [...(last.content.parts || []), ...scrubbedParts], + }, + }; + } else { + result.push({ + id: turn.id, + content: { + role: turn.content.role, + parts: scrubbedParts, + }, + }); + } + } + return result; } /** * Deep-scrubs an array of Content objects to remove non-standard properties. + * Coalesces adjacent turns of the same role to preserve Gemini API alternation invariants. */ export function scrubContents(contents: Content[]): Content[] { - return contents.map((content) => ({ - role: content.role, - parts: (content.parts || []).map((p) => scrubPart(p)), - })); + const result: Content[] = []; + for (const content of contents) { + const nonThoughtParts = (content.parts ?? []).filter( + (p) => p && !isInternalThought(p), + ); + if (nonThoughtParts.length === 0) continue; // Skip turns that became empty after thought stripping + + const scrubbedParts = nonThoughtParts.map((p) => scrubPart(p)); + + const lastIdx = result.length - 1; + const last = result[lastIdx]; + if (last && last.role === content.role) { + // Coalesce adjacent turns of the same role inline + result[lastIdx] = { + role: last.role, + parts: [...(last.parts || []), ...scrubbedParts], + }; + } else { + result.push({ + role: content.role, + parts: scrubbedParts, + }); + } + } + return result; } interface ThoughtPart extends Part { + thought?: boolean; thoughtSignature?: string; }