Skip to content
Merged
2 changes: 1 addition & 1 deletion packages/core/src/core/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,7 @@ describe('Gemini Client (client.ts)', () => {
.fn()
.mockReturnValue(contentGeneratorConfig),
getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
getModel: vi.fn().mockReturnValue('test-model'),
getModel: vi.fn().mockReturnValue('gemini-1.5-pro'),
getUserTier: vi.fn().mockReturnValue(undefined),
getEmbeddingModel: vi.fn().mockReturnValue('test-embedding-model'),
getApiKey: vi.fn().mockReturnValue('test-key'),
Expand Down
195 changes: 195 additions & 0 deletions packages/core/src/core/geminiChat.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
stripToolCallIdPrefixes,
type HistoryTurn,
coalesceConsecutiveRoles,
stripThoughts,
} from './geminiChat.js';
import {
type CompletedToolCall,
Expand Down Expand Up @@ -2282,6 +2283,110 @@
text: 'actual conversational response',
});
});

it('should completely filter out thought parts from getHistoryTurns when context management is disabled but model is gemini-2/modern', () => {

Check warning on line 2287 in packages/core/src/core/geminiChat.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-2". Please make sure this change is appropriate to submit.
vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(false);
vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-pro');

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',
});
});

it('should completely filter out thought parts from getHistoryTurns when model supports modern features', () => {
vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(false);
vi.mocked(mockConfig.getModel).mockReturnValue('gemini-3.1-pro-preview');

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',
});
});

it('should completely filter out model turns that end up with empty parts after stripping thoughts', () => {
vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(false);
vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-pro');

chat.setHistory([
{
role: 'user',
parts: [{ text: 'hello' }],
},
{
role: 'model',
parts: [
{ text: 'internal monologue', thought: true } as unknown as Part,
],
},
]);

const turns = chat.getHistoryTurns(true);

// Since the model turn contains only a thought part, it should be filtered out entirely.
expect(turns).toHaveLength(1);
expect(turns[0].content.role).toBe('user');
});

it('should coalesce consecutive user turns when an intermediate model turn is stripped', () => {
vi.mocked(mockConfig.isContextManagementEnabled).mockReturnValue(false);
vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-pro');

chat.setHistory([
{ role: 'user', parts: [{ text: 'Question 1' }] },
{
role: 'model',
parts: [{ text: 'thinking...', thought: true } as unknown as Part],
},
{ role: 'user', parts: [{ text: 'Question 2' }] },
]);

const turns = chat.getHistoryTurns(true);

// The model turn contains only a thought part, so it is stripped.
// The two adjacent user turns must be coalesced into one user turn.
expect(turns).toHaveLength(1);
expect(turns[0].content.role).toBe('user');
expect(turns[0].content.parts).toHaveLength(2);
expect(turns[0].content.parts![0].text).toBe('Question 1');
expect(turns[0].content.parts![1].text).toBe('Question 2');
});
});

describe('ensureActiveLoopHasThoughtSignatures', () => {
Expand Down Expand Up @@ -3196,4 +3301,94 @@
expect(coalesceConsecutiveRoles(history)).toEqual(history);
});
});

describe('stripThoughts', () => {
it('should return empty history if empty array is passed', () => {
expect(stripThoughts([])).toEqual([]);
});

it('should strip thought parts and keep the turn if other parts remain', () => {
const history: HistoryTurn[] = [
{
id: '1',
content: {
role: 'model',
parts: [
{ text: 'internal monologue', thought: true } as unknown as Part,
{ text: 'visible response' },
],
},
},
];
expect(stripThoughts(history)).toEqual([
{
id: '1',
content: {
role: 'model',
parts: [{ text: 'visible response' }],
},
},
]);
});

it('should completely remove a turn if all its parts are thought parts', () => {
const history: HistoryTurn[] = [
{
id: '1',
content: {
role: 'user',
parts: [{ text: 'hello' }],
},
},
{
id: '2',
content: {
role: 'model',
parts: [
{ text: 'internal monologue', thought: true } as unknown as Part,
],
},
},
];
expect(stripThoughts(history)).toEqual([
{
id: '1',
content: {
role: 'user',
parts: [{ text: 'hello' }],
},
},
]);
});

it('should preserve turns that do not have parts arrays', () => {
const history: HistoryTurn[] = [{ id: '1', content: { role: 'user' } }];
expect(stripThoughts(history)).toEqual(history);
});

it('should preserve top-level metadata when stripping thoughts', () => {
const history: HistoryTurn[] = [
{
id: '1',
content: {
role: 'model',
parts: [
{ text: 'internal monologue', thought: true } as unknown as Part,
{ text: 'visible response' },
],
},
// top-level turn metadata
timestamp: '2026-07-23T00:00:00.000Z',
metadata: { some: 'value' },
} as unknown as HistoryTurn,
];
const stripped = stripThoughts(history);
expect(stripped).toHaveLength(1);
expect(stripped[0]).toHaveProperty(
'timestamp',
'2026-07-23T00:00:00.000Z',
);
expect(stripped[0]).toHaveProperty('metadata', { some: 'value' });
});
});
});
52 changes: 42 additions & 10 deletions packages/core/src/core/geminiChat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,11 @@ import {
getRetryErrorType,
} from '../utils/retry.js';
import type { ValidationRequiredError } from '../utils/googleQuotaErrors.js';
import { resolveModel, supportsModernFeatures } from '../config/models.js';
import {
resolveModel,
supportsModernFeatures,
isGemini2Model,
} from '../config/models.js';
import { hasCycleInSchema } from '../tools/tools.js';
import type { StructuredError } from './turn.js';
import type { CompletedToolCall } from '../scheduler/types.js';
Expand Down Expand Up @@ -766,9 +770,10 @@ export class GeminiChat {
abortSignal,
};

let contentsToUse: Content[] = supportsModernFeatures(modelToUse)
? [...contentsForPreviewModel]
: [...requestContents];
let contentsToUse: Content[] =
supportsModernFeatures(modelToUse) || isGemini2Model(modelToUse)
? [...contentsForPreviewModel]
: [...requestContents];

const hookSystem = this.context.config.getHookSystem();
if (hookSystem) {
Expand Down Expand Up @@ -810,9 +815,10 @@ export class GeminiChat {
);
lastModelToUse = modelToUse;
// Re-evaluate contentsToUse based on the new model's feature support
contentsToUse = supportsModernFeatures(modelToUse)
? [...contentsForPreviewModel]
: [...requestContents];
contentsToUse =
supportsModernFeatures(modelToUse) || isGemini2Model(modelToUse)
? [...contentsForPreviewModel]
: [...requestContents];
}
if (beforeModelResult.modifiedConfig) {
Object.assign(config, beforeModelResult.modifiedConfig);
Expand Down Expand Up @@ -956,9 +962,16 @@ export class GeminiChat {
? extractCuratedHistory(this.agentHistory.get())
: [...this.agentHistory.get()];

return this.context.config.isContextManagementEnabled()
? scrubHistory(history)
: history;
if (this.context.config.isContextManagementEnabled()) {
return scrubHistory(history);
}

const model = this.context.config.getModel();
if (isGemini2Model(model) || supportsModernFeatures(model)) {
return coalesceConsecutiveRoles(stripThoughts(history));
}

return history;
}

/**
Expand Down Expand Up @@ -1503,3 +1516,22 @@ export function coalesceConsecutiveRoles(
}
return result;
}

export function stripThoughts(history: HistoryTurn[]): HistoryTurn[] {
return history
.map((turn) => {
if (!turn.content.parts) return turn;
const hasThought = turn.content.parts.some((p) => p && p.thought);
if (!hasThought) return turn;

const nonThoughtParts = turn.content.parts.filter((p) => p && !p.thought);
return {
...turn,
content: {
...turn.content,
parts: nonThoughtParts,
},
};
})
.filter((turn) => !turn.content.parts || turn.content.parts.length > 0);
}
Comment thread
DavidAPierce marked this conversation as resolved.
Loading