Skip to content

Commit 01635dd

Browse files
joshualittjoshualitt
authored
fix(context): implement loose boundary policy for gc backstop. (google-gemini#26594)
1 parent 12c8469 commit 01635dd

12 files changed

Lines changed: 593 additions & 137 deletions

File tree

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -60,9 +60,7 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {
6060

6161
// Verify Episode 0 (System) was pruned, so we now start with a sentinel due to role alternation
6262
expect(projection[0].role).toBe('user');
63-
expect(projection[0].parts![0].text).toBe(
64-
'[Continuing from previous AI thoughts...]',
65-
);
63+
expect(projection[0].parts![0].text).toContain('User turn 17');
6664

6765
// Filter out synthetic Yield nodes (they are model responses without actual tool/text bodies)
6866
const contentNodes = projection.filter(

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,11 @@ export class ContextManager {
7272
event.targets,
7373
event.returnedNodes,
7474
);
75-
this.evaluateTriggers(new Set());
75+
// We explicitly DO NOT call evaluateTriggers here.
76+
// The Context Manager is a one-way assembly line. It only evaluates triggers
77+
// when fundamentally new organic context is added via PristineHistoryUpdated.
78+
// Re-evaluating after a processor finishes creates infinite feedback loops if
79+
// the processor fails to reduce the token count below the threshold.
7680
});
7781

7882
this.historyObserver.start();
@@ -126,10 +130,15 @@ export class ContextManager {
126130
// Walk backwards finding nodes that fall out of the retained budget
127131
for (let i = this.buffer.nodes.length - 1; i >= 0; i--) {
128132
const node = this.buffer.nodes[i];
133+
const priorTokens = rollingTokens;
129134
rollingTokens += this.env.tokenCalculator.calculateConcreteListTokens([
130135
node,
131136
]);
132-
if (rollingTokens > this.sidecar.config.budget.retainedTokens) {
137+
138+
// Loose Boundary Policy: If this node is the one that pushes us over the retained limit,
139+
// we KEEP it to prevent aggressive undershooting. We only age out nodes that are
140+
// strictly *older* than the boundary node.
141+
if (priorTokens > this.sidecar.config.budget.retainedTokens) {
133142
// Only age out if not protected
134143
if (!protectedIds.has(node.id)) {
135144
agedOutNodes.add(node.id);

‎packages/core/src/context/graph/render.test.ts‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,4 +61,169 @@ describe('render', () => {
6161

6262
expect(result.history).toEqual([{ text: '1' }, { text: '2' }]);
6363
});
64+
65+
it('simulates the boundary knapsack problem (loose boundary policy)', async () => {
66+
// 10k, 20k, 40k, 5k
67+
const mockNodes: ConcreteNode[] = [
68+
{
69+
id: 'D',
70+
type: NodeType.USER_PROMPT,
71+
payload: {} as Part,
72+
} as unknown as ConcreteNode,
73+
{
74+
id: 'C',
75+
type: NodeType.AGENT_THOUGHT,
76+
payload: {} as Part,
77+
} as unknown as ConcreteNode,
78+
{
79+
id: 'B',
80+
type: NodeType.USER_PROMPT,
81+
payload: {} as Part,
82+
} as unknown as ConcreteNode,
83+
{
84+
id: 'A',
85+
type: NodeType.AGENT_THOUGHT,
86+
payload: {} as Part,
87+
} as unknown as ConcreteNode,
88+
];
89+
90+
const tokenMap: Record<string, number> = {
91+
D: 5000,
92+
C: 40000,
93+
B: 20000,
94+
A: 10000,
95+
};
96+
97+
const orchestrator = {
98+
executeTriggerSync: vi.fn(async (trigger, nodes, agedOutNodes) =>
99+
nodes.filter((n: ConcreteNode) => !agedOutNodes.has(n.id)),
100+
),
101+
} as unknown as PipelineOrchestrator;
102+
103+
const sidecar = {
104+
config: {
105+
budget: { maxTokens: 150000, retainedTokens: 65000 },
106+
},
107+
} as unknown as ContextProfile;
108+
109+
const currentTokens = 160000;
110+
111+
const env = {
112+
llmClient: {
113+
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1000 }),
114+
},
115+
tokenCalculator: {
116+
calculateConcreteListTokens: vi.fn((nodes) => {
117+
if (nodes.length === 1) return tokenMap[nodes[0].id];
118+
return currentTokens;
119+
}),
120+
calculateTokenBreakdown: vi.fn(() => ({})),
121+
},
122+
graphMapper: {
123+
fromGraph: vi.fn((nodes: readonly ConcreteNode[]) =>
124+
nodes.map((n) => ({ text: n.id })),
125+
),
126+
},
127+
} as unknown as ContextEnvironment;
128+
129+
const tracer = {
130+
logEvent: vi.fn(),
131+
} as unknown as ContextTracer;
132+
133+
const result = await render(
134+
mockNodes,
135+
orchestrator,
136+
sidecar,
137+
tracer,
138+
env,
139+
new Map(),
140+
0,
141+
new Set(),
142+
);
143+
144+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
145+
const surviving = result.history.map((c: any) => c.text);
146+
// Loose Boundary: A (10k), B (20k), C (40k). Total = 70k.
147+
// Adding C pushes rolling total (70k) above retainedTokens (65k).
148+
// Under loose policy, C survives. D is strictly older and drops.
149+
expect(surviving).toEqual(['C', 'B', 'A']); // D is dropped
150+
});
151+
152+
it('drops nodes that are STRICTLY older than the boundary node', async () => {
153+
const mockNodes: ConcreteNode[] = [
154+
{
155+
id: 'A',
156+
type: NodeType.USER_PROMPT,
157+
payload: {} as Part,
158+
} as unknown as ConcreteNode,
159+
{
160+
id: 'B',
161+
type: NodeType.AGENT_THOUGHT,
162+
payload: {} as Part,
163+
} as unknown as ConcreteNode,
164+
{
165+
id: 'C',
166+
type: NodeType.USER_PROMPT,
167+
payload: {} as Part,
168+
} as unknown as ConcreteNode,
169+
];
170+
171+
const tokenMap: Record<string, number> = {
172+
C: 40000,
173+
B: 40000,
174+
A: 10000,
175+
};
176+
177+
const orchestrator = {
178+
executeTriggerSync: vi.fn(async (trigger, nodes, agedOutNodes) =>
179+
nodes.filter((n: ConcreteNode) => !agedOutNodes.has(n.id)),
180+
),
181+
} as unknown as PipelineOrchestrator;
182+
183+
const sidecar = {
184+
config: {
185+
budget: { maxTokens: 150000, retainedTokens: 65000 },
186+
},
187+
} as unknown as ContextProfile;
188+
189+
const currentTokens = 160000;
190+
191+
const env = {
192+
llmClient: {
193+
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1000 }),
194+
},
195+
tokenCalculator: {
196+
calculateConcreteListTokens: vi.fn((nodes) => {
197+
if (nodes.length === 1) return tokenMap[nodes[0].id];
198+
return currentTokens;
199+
}),
200+
calculateTokenBreakdown: vi.fn(() => ({})),
201+
},
202+
graphMapper: {
203+
fromGraph: vi.fn((nodes: readonly ConcreteNode[]) =>
204+
nodes.map((n) => ({ text: n.id })),
205+
),
206+
},
207+
} as unknown as ContextEnvironment;
208+
209+
const tracer = {
210+
logEvent: vi.fn(),
211+
} as unknown as ContextTracer;
212+
213+
const result = await render(
214+
mockNodes,
215+
orchestrator,
216+
sidecar,
217+
tracer,
218+
env,
219+
new Map(),
220+
0,
221+
new Set(),
222+
);
223+
224+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
225+
const surviving = result.history.map((c: any) => c.text);
226+
// C(40k), B(40k). Adding B pushes total to 80k. B is the boundary node and survives. A drops.
227+
expect(surviving).toEqual(['B', 'C']); // A is dropped
228+
});
64229
});

‎packages/core/src/context/graph/render.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import type { ContextTracer } from '../tracer.js';
1010
import type { ContextProfile } from '../config/profiles.js';
1111
import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
1212
import type { ContextEnvironment } from '../pipeline/environment.js';
13+
import { performCalibration } from '../utils/tokenCalibration.js';
1314

1415
/**
1516
* Maps the Episodic Context Graph back into a raw Gemini Content[] array for transmission.
@@ -68,6 +69,7 @@ export async function render(
6869
tracer.logEvent('Render', 'Render Context for LLM', {
6970
renderedContext: contents,
7071
});
72+
performCalibration(env, visibleNodes, contents);
7173
return { history: contents, didApplyManagement: false };
7274
}
7375
const targetDelta = currentTokens - sidecar.config.budget.retainedTokens;
@@ -83,9 +85,12 @@ export async function render(
8385
// Start from newest and count backwards
8486
for (let i = nodes.length - 1; i >= 0; i--) {
8587
const node = nodes[i];
88+
const priorTokens = rollingTokens;
8689
const nodeTokens = env.tokenCalculator.calculateConcreteListTokens([node]);
8790
rollingTokens += nodeTokens;
88-
if (rollingTokens > sidecar.config.budget.retainedTokens) {
91+
92+
// Loose Boundary Policy: Keep the node that crosses the boundary
93+
if (priorTokens > sidecar.config.budget.retainedTokens) {
8994
agedOutNodes.add(node.id);
9095
}
9196
}
@@ -113,5 +118,6 @@ export async function render(
113118
tracer.logEvent('Render', 'Render Sanitized Context for LLM', {
114119
renderedContextSanitized: contents,
115120
});
121+
performCalibration(env, visibleNodes, contents);
116122
return { history: contents, didApplyManagement: true };
117123
}

‎packages/core/src/context/initializer.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,10 @@ export async function initializeContextManager(
9494
tracer,
9595
4,
9696
eventBus,
97+
{
98+
calibrateTokenCalculation:
99+
!!process.env['GEMINI_CONTEXT_CALIBRATE_TOKEN_CALCULATIONS'],
100+
},
97101
);
98102

99103
const orchestrator = new PipelineOrchestrator(

‎packages/core/src/context/pipeline/environment.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ import type { ContextGraphMapper } from '../graph/mapper.js';
1313

1414
export type { ContextTracer, ContextEventBus };
1515

16+
export interface RenderOptions {
17+
calibrateTokenCalculation?: boolean;
18+
}
19+
1620
export interface ContextEnvironment {
1721
readonly llmClient: BaseLlmClient;
1822
readonly promptId: string;
@@ -26,4 +30,5 @@ export interface ContextEnvironment {
2630
readonly inbox: LiveInbox;
2731
readonly behaviorRegistry: NodeBehaviorRegistry;
2832
readonly graphMapper: ContextGraphMapper;
33+
readonly renderOptions?: RenderOptions;
2934
}

‎packages/core/src/context/pipeline/environmentImpl.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
88
import type { ContextTracer } from '../tracer.js';
9-
import type { ContextEnvironment } from './environment.js';
9+
import type { ContextEnvironment, RenderOptions } from './environment.js';
1010
import type { ContextEventBus } from '../eventBus.js';
1111
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
1212
import { LiveInbox } from './inbox.js';
@@ -29,6 +29,7 @@ export class ContextEnvironmentImpl implements ContextEnvironment {
2929
readonly tracer: ContextTracer,
3030
readonly charsPerToken: number,
3131
readonly eventBus: ContextEventBus,
32+
readonly renderOptions?: RenderOptions,
3233
) {
3334
this.behaviorRegistry = new NodeBehaviorRegistry();
3435
registerBuiltInBehaviors(this.behaviorRegistry);

0 commit comments

Comments
 (0)