Skip to content

Commit 50c2e67

Browse files
joshualittjoshualitt
authored
fix(context): Ensure last message is processed. (google-gemini#27232)
1 parent 906f8a3 commit 50c2e67

31 files changed

Lines changed: 1146 additions & 1128 deletions

‎integration-tests/context-fidelity.test.ts‎

Lines changed: 45 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,18 @@ import type { FakeResponse, HistoryTurn } from '@google/gemini-cli-core';
1414
describe('Context Management Fidelity E2E', () => {
1515
let rig: TestRig;
1616

17+
function generateRandomString(length: number): string {
18+
const characters =
19+
'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
20+
let result = '';
21+
for (let i = 0; i < length; i++) {
22+
result += characters.charAt(
23+
Math.floor(Math.random() * characters.length),
24+
);
25+
}
26+
return result;
27+
}
28+
1729
beforeEach(() => {
1830
rig = new TestRig();
1931
});
@@ -52,7 +64,7 @@ describe('Context Management Fidelity E2E', () => {
5264

5365
const countTokensResponse: FakeResponse = {
5466
method: 'countTokens',
55-
response: { totalTokens: 50000 },
67+
response: { totalTokens: 1000 },
5668
};
5769

5870
const streamResponse = (text: string): FakeResponse => ({
@@ -87,7 +99,6 @@ describe('Context Management Fidelity E2E', () => {
8799
},
88100
});
89101

90-
const massivePayload = 'X'.repeat(50000);
91102
const traceDir = path.join(rig.testDir!, 'traces');
92103
fs.mkdirSync(traceDir, { recursive: true });
93104
const traceLog = path.join(traceDir, 'trace.log');
@@ -105,24 +116,35 @@ describe('Context Management Fidelity E2E', () => {
105116
streamResponse('Ack 3'),
106117
streamResponse('Ack 4'),
107118
streamResponse('Ack 5'),
119+
streamResponse('Ack 6'),
120+
streamResponse('Ack 7'),
121+
streamResponse('Ack 8'),
122+
streamResponse('Ack 9'),
123+
streamResponse('Ack 10'),
124+
streamResponse('Ack 11'),
125+
streamResponse('Ack 12'),
108126
];
109127
for (let i = 0; i < 50; i++) {
110128
runMocks.push(snapshotResponse);
111129
runMocks.push(countTokensResponse);
112130
}
113131

114-
// Turn 1: Initial massive payload to put pressure
115-
await rig.run({
116-
args: [
117-
'--debug',
118-
'--fake-responses-non-strict',
119-
setupResponses('resp1.json', runMocks),
120-
],
121-
stdin: 'Turn 1: ' + massivePayload,
122-
env: commonEnv,
123-
});
132+
// Turns 1-10: Build up history
133+
for (let i = 1; i <= 10; i++) {
134+
await rig.run({
135+
args: [
136+
'--debug',
137+
i === 1 ? '' : '--resume',
138+
i === 1 ? '' : 'latest',
139+
'--fake-responses-non-strict',
140+
setupResponses(`resp_init_${i}.json`, runMocks),
141+
].filter(Boolean),
142+
stdin: `Turn ${i}: ` + generateRandomString(900),
143+
env: commonEnv,
144+
});
145+
}
124146

125-
// Turn 2: Another turn, resuming Turn 1
147+
// Turn 11: Penultimate turn
126148
await rig.run({
127149
args: [
128150
'--debug',
@@ -131,11 +153,11 @@ describe('Context Management Fidelity E2E', () => {
131153
'--fake-responses-non-strict',
132154
setupResponses('resp2.json', runMocks),
133155
],
134-
stdin: 'Turn 2: ' + massivePayload,
156+
stdin: 'Turn 11: ' + generateRandomString(900),
135157
env: commonEnv,
136158
});
137159

138-
// Turn 3: Third turn to force GC, resuming Turn 2
160+
// Turn 12: Breach threshold and force GC
139161
await rig.run({
140162
args: [
141163
'--debug',
@@ -144,7 +166,7 @@ describe('Context Management Fidelity E2E', () => {
144166
'--fake-responses-non-strict',
145167
setupResponses('resp3.json', runMocks),
146168
],
147-
stdin: 'Turn 3: ' + massivePayload,
169+
stdin: 'Turn 12: ' + generateRandomString(900),
148170
env: commonEnv,
149171
});
150172

@@ -214,12 +236,16 @@ describe('Context Management Fidelity E2E', () => {
214236

215237
// Most importantly, synthetic IDs (like summaries) must be stable.
216238
const syntheticTurns = contextBeforeExit!.filter(
217-
(t: HistoryTurn) => t.id && t.id.length === 32,
218-
); // deriveStableId produces 32-char hex
239+
(t: HistoryTurn) =>
240+
t.content.parts?.some((p) => p.text?.includes('active_tasks')) ||
241+
(t.id && t.id.length === 32),
242+
);
219243
expect(syntheticTurns.length).toBeGreaterThan(0);
220244

221245
const syntheticTurnsAfter = contextAfterResume!.filter(
222-
(t: HistoryTurn) => t.id && t.id.length === 32,
246+
(t: HistoryTurn) =>
247+
t.content.parts?.some((p) => p.text?.includes('active_tasks')) ||
248+
(t.id && t.id.length === 32),
223249
);
224250
expect(syntheticTurnsAfter.length).toBeGreaterThanOrEqual(
225251
syntheticTurns.length,

‎packages/core/src/context/config/profiles.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -177,8 +177,8 @@ export const stressTestProfile: ContextProfile = {
177177
name: 'Stress Test',
178178
config: {
179179
budget: {
180-
retainedTokens: 4000,
181-
maxTokens: 10000,
180+
retainedTokens: 1500,
181+
maxTokens: 5000,
182182
},
183183
processorOptions: {
184184
ToolMasking: {

‎packages/core/src/context/contextManager.barrier.test.ts‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
createSyntheticHistory,
1212
createMockContextConfig,
1313
setupContextComponentTest,
14+
deriveStableId,
1415
} from './testing/contextTestUtils.js';
1516

1617
describe('ContextManager Sync Pressure Barrier Tests', () => {
@@ -32,10 +33,14 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {
3233
);
3334

3435
// 2. Add System Prompt (Episode 0 - Protected)
36+
const envId = deriveStableId(['environment-context']);
3537
chatHistory.set([
3638
{
37-
id: 'h1',
38-
content: { role: 'user', parts: [{ text: 'System prompt' }] },
39+
id: envId,
40+
content: {
41+
role: 'user',
42+
parts: [{ text: '<session_context>\nSystem prompt' }],
43+
},
3944
},
4045
{
4146
id: 'h2',
@@ -74,8 +79,10 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {
7479

7580
expect(projection.length).toBeLessThan(rawHistoryLength);
7681

77-
// Verify Episode 0 (System) was pruned, so we now start with a sentinel due to role alternation
82+
// Verify Episode 0 (System) was PRESERVED because it is pinned Turn 0.
83+
expect(projection[0].id).toBe(envId);
7884
expect(projection[0].content.role).toBe('user');
85+
7986
const projectionString = JSON.stringify(projection);
8087
expect(projectionString).toContain('User turn 17');
8188
// Filter out synthetic Yield nodes (they are model responses without actual tool/text bodies)
@@ -86,19 +93,13 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {
8693
);
8794

8895
// Verify the latest turn is perfectly preserved at the back
89-
// Note: The HistoryHardener appends a "Please continue." user turn if we end on model,
90-
// so we look at the turns before the sentinel.
91-
const lastSentinel = contentNodes[contentNodes.length - 1].content;
92-
const lastModel = contentNodes[contentNodes.length - 2].content;
93-
const lastUser = contentNodes[contentNodes.length - 3].content;
96+
const lastModel = contentNodes[contentNodes.length - 1].content;
97+
const lastUser = contentNodes[contentNodes.length - 2].content;
9498

95-
expect(lastSentinel.role).toBe('user');
96-
expect(lastSentinel.parts![0].text).toBe('Please continue.');
99+
expect(lastModel.role).toBe('model');
100+
expect(lastModel.parts![0].text).toBe('Final answer.');
97101

98102
expect(lastUser.role).toBe('user');
99103
expect(lastUser.parts![0].text).toBe('Final question.');
100-
101-
expect(lastModel.role).toBe('model');
102-
expect(lastModel.parts![0].text).toBe('Final answer.');
103104
});
104105
});

‎packages/core/src/context/contextManager.test.ts‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
/**
2+
* @license
3+
* Copyright 2026 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { vi, describe, it, expect, beforeEach, type Mock } from 'vitest';
8+
import { ContextManager } from './contextManager.js';
9+
import type { ContextProfile } from './config/profiles.js';
10+
import type { ContextEnvironment } from './pipeline/environment.js';
11+
import type { ContextTracer } from './tracer.js';
12+
import type { PipelineOrchestrator } from './pipeline/orchestrator.js';
13+
import type {
14+
AgentChatHistory,
15+
HistoryTurn,
16+
} from '../core/agentChatHistory.js';
17+
import type { AdvancedTokenCalculator } from './utils/contextTokenCalculator.js';
18+
import { createMockEnvironment } from './testing/contextTestUtils.js';
19+
20+
describe('ContextManager', () => {
21+
let mockSidecar: ContextProfile;
22+
let mockEnv: ContextEnvironment;
23+
let mockTracer: ContextTracer;
24+
let mockOrchestrator: PipelineOrchestrator;
25+
let mockChatHistory: AgentChatHistory;
26+
let mockAdvancedTokenCalculator: AdvancedTokenCalculator;
27+
28+
beforeEach(() => {
29+
vi.resetAllMocks();
30+
31+
mockSidecar = {
32+
name: 'test-profile',
33+
config: { budget: { retainedTokens: 1000, maxTokens: 2000 } },
34+
buildPipelines: vi.fn().mockReturnValue([]),
35+
buildAsyncPipelines: vi.fn().mockReturnValue([]),
36+
} as unknown as ContextProfile;
37+
38+
mockEnv = createMockEnvironment();
39+
mockTracer = mockEnv.tracer;
40+
41+
mockOrchestrator = {
42+
setNodeProvider: vi.fn(),
43+
waitForPipelines: vi.fn().mockResolvedValue(undefined),
44+
executeTriggerSync: vi
45+
.fn()
46+
.mockImplementation(async (trigger, nodes) => nodes),
47+
shutdown: vi.fn(),
48+
} as unknown as PipelineOrchestrator;
49+
50+
mockChatHistory = {
51+
all: vi.fn().mockReturnValue([]),
52+
last: vi.fn(),
53+
getById: vi.fn(),
54+
getTurnById: vi.fn(),
55+
getTurnsByIds: vi.fn(),
56+
getNeighboringTurns: vi.fn(),
57+
getHistory: vi.fn().mockReturnValue([]),
58+
get: vi.fn().mockReturnValue([]),
59+
setHistory: vi.fn(),
60+
getHistoryTurns: vi.fn().mockReturnValue([]),
61+
getRawHistory: vi.fn().mockReturnValue([]),
62+
addTurn: vi.fn(),
63+
updateTurn: vi.fn(),
64+
addListener: vi.fn(),
65+
removeListener: vi.fn(),
66+
clear: vi.fn(),
67+
subscribe: vi.fn(),
68+
} as unknown as AgentChatHistory;
69+
70+
mockAdvancedTokenCalculator = {
71+
getRawBaseUnits: vi.fn().mockReturnValue(0),
72+
getRawBaseUnitsForContent: vi.fn().mockReturnValue(0),
73+
calculateTokensAndBaseUnits: vi
74+
.fn()
75+
.mockReturnValue({ tokens: 0, baseUnits: 0 }),
76+
} as unknown as AdvancedTokenCalculator;
77+
});
78+
79+
it('renderHistory should process pendingRequest via the new_message pipeline', async () => {
80+
const contextManager = new ContextManager(
81+
mockSidecar,
82+
mockEnv,
83+
mockTracer,
84+
mockOrchestrator,
85+
mockChatHistory,
86+
mockAdvancedTokenCalculator,
87+
);
88+
89+
const largeToolOutput = 'a'.repeat(10000);
90+
const pendingRequest: HistoryTurn = {
91+
id: 'pending-turn-1',
92+
content: {
93+
role: 'user',
94+
parts: [
95+
{
96+
functionResponse: {
97+
name: 'run_shell_command',
98+
response: {
99+
output: largeToolOutput,
100+
},
101+
},
102+
},
103+
],
104+
},
105+
};
106+
107+
await contextManager.renderHistory(pendingRequest);
108+
109+
expect(mockOrchestrator.executeTriggerSync).toHaveBeenCalledExactlyOnceWith(
110+
'new_message',
111+
expect.any(Array),
112+
expect.any(Set),
113+
);
114+
115+
// Check that the node passed to the orchestrator corresponds to our pendingRequest
116+
const call = (mockOrchestrator.executeTriggerSync as unknown as Mock).mock
117+
.calls[0];
118+
const passedNodes = call[1];
119+
const passedNodeIds = call[2];
120+
121+
expect(passedNodes).toHaveLength(1);
122+
expect(passedNodes[0].type).toBe('TOOL_EXECUTION');
123+
expect(passedNodes[0].payload.functionResponse.response.output).toBe(
124+
largeToolOutput,
125+
);
126+
expect(passedNodeIds.has(passedNodes[0].id)).toBe(true);
127+
});
128+
129+
it('renderHistory should exclude pendingRequest from the result (late binding)', async () => {
130+
const contextManager = new ContextManager(
131+
mockSidecar,
132+
mockEnv,
133+
mockTracer,
134+
mockOrchestrator,
135+
mockChatHistory,
136+
mockAdvancedTokenCalculator,
137+
);
138+
139+
const pendingRequest: HistoryTurn = {
140+
id: 'pending-turn-1',
141+
content: { role: 'user', parts: [{ text: 'Active prompt' }] },
142+
};
143+
144+
const { history, apiHistory } =
145+
await contextManager.renderHistory(pendingRequest);
146+
147+
// Should be empty because mockChatHistory has no historical turns
148+
expect(history).toHaveLength(0);
149+
expect(apiHistory).toHaveLength(0);
150+
});
151+
});

0 commit comments

Comments
 (0)