Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions packages/core/src/core/geminiChat.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import {
ThinkingLevel,
type Content,
type GenerateContentResponse,
type Part,
} from '@google/genai';
import type { ContentGenerator } from '../core/contentGenerator.js';
import {
Expand Down Expand Up @@ -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',
Comment thread
amelidev marked this conversation as resolved.
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, '', [], []);
Expand Down
189 changes: 188 additions & 1 deletion packages/core/src/utils/historyHardening.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down Expand Up @@ -375,4 +377,189 @@ 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', () => {
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' }],
},
{
role: 'user',
parts: [{ text: 'How are you?' }],
},
]);
});
});

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' },
]);
});
});
100 changes: 86 additions & 14 deletions packages/core/src/utils/historyHardening.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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 as ThoughtPart).thought;
}
Comment thread
amelidev marked this conversation as resolved.

/**
* 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) => !isInternalThought(p),
);
return {
id: turn.id,
content: {
...turn.content,
parts: nonThoughtParts,
},
};
});
}
Comment thread
amelidev marked this conversation as resolved.

/**
* Combines adjacent turns with the same role and removes empty turns.
*/
Expand All @@ -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 } });
Expand Down Expand Up @@ -344,23 +381,58 @@ 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) => !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.
*/
export function scrubContents(contents: Content[]): Content[] {
return contents.map((content) => ({
role: content.role,
parts: (content.parts || []).map((p) => scrubPart(p)),
}));
return contents
.map((content) => {
const nonThoughtParts = (content.parts ?? []).filter(
(p) => !isInternalThought(p),
);
return {
role: content.role,
parts: nonThoughtParts.map((p) => scrubPart(p)),
};
})
.filter((content) => content.parts.length > 0);
}

interface ThoughtPart extends Part {
thought?: boolean;
thoughtSignature?: string;
}

Expand Down
Loading