Skip to content
Merged
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
4 changes: 1 addition & 3 deletions packages/core/src/context/contextManager.barrier.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,7 @@ describe('ContextManager Sync Pressure Barrier Tests', () => {

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

// Filter out synthetic Yield nodes (they are model responses without actual tool/text bodies)
const contentNodes = projection.filter(
Expand Down
13 changes: 11 additions & 2 deletions packages/core/src/context/contextManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,11 @@ export class ContextManager {
event.targets,
event.returnedNodes,
);
this.evaluateTriggers(new Set());
// We explicitly DO NOT call evaluateTriggers here.
// The Context Manager is a one-way assembly line. It only evaluates triggers
// when fundamentally new organic context is added via PristineHistoryUpdated.
// Re-evaluating after a processor finishes creates infinite feedback loops if
// the processor fails to reduce the token count below the threshold.
});

this.historyObserver.start();
Expand Down Expand Up @@ -126,10 +130,15 @@ export class ContextManager {
// Walk backwards finding nodes that fall out of the retained budget
for (let i = this.buffer.nodes.length - 1; i >= 0; i--) {
const node = this.buffer.nodes[i];
const priorTokens = rollingTokens;
rollingTokens += this.env.tokenCalculator.calculateConcreteListTokens([
node,
]);
if (rollingTokens > this.sidecar.config.budget.retainedTokens) {

// Loose Boundary Policy: If this node is the one that pushes us over the retained limit,
// we KEEP it to prevent aggressive undershooting. We only age out nodes that are
// strictly *older* than the boundary node.
if (priorTokens > this.sidecar.config.budget.retainedTokens) {
// Only age out if not protected
if (!protectedIds.has(node.id)) {
agedOutNodes.add(node.id);
Expand Down
165 changes: 165 additions & 0 deletions packages/core/src/context/graph/render.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,4 +61,169 @@ describe('render', () => {

expect(result.history).toEqual([{ text: '1' }, { text: '2' }]);
});

it('simulates the boundary knapsack problem (loose boundary policy)', async () => {
// 10k, 20k, 40k, 5k
const mockNodes: ConcreteNode[] = [
{
id: 'D',
type: NodeType.USER_PROMPT,
payload: {} as Part,
} as unknown as ConcreteNode,
{
id: 'C',
type: NodeType.AGENT_THOUGHT,
payload: {} as Part,
} as unknown as ConcreteNode,
{
id: 'B',
type: NodeType.USER_PROMPT,
payload: {} as Part,
} as unknown as ConcreteNode,
{
id: 'A',
type: NodeType.AGENT_THOUGHT,
payload: {} as Part,
} as unknown as ConcreteNode,
];

const tokenMap: Record<string, number> = {
D: 5000,
C: 40000,
B: 20000,
A: 10000,
};

const orchestrator = {
executeTriggerSync: vi.fn(async (trigger, nodes, agedOutNodes) =>
nodes.filter((n: ConcreteNode) => !agedOutNodes.has(n.id)),
),
} as unknown as PipelineOrchestrator;

const sidecar = {
config: {
budget: { maxTokens: 150000, retainedTokens: 65000 },
},
} as unknown as ContextProfile;

const currentTokens = 160000;

const env = {
llmClient: {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1000 }),
},
tokenCalculator: {
calculateConcreteListTokens: vi.fn((nodes) => {
if (nodes.length === 1) return tokenMap[nodes[0].id];
return currentTokens;
}),
calculateTokenBreakdown: vi.fn(() => ({})),
},
graphMapper: {
fromGraph: vi.fn((nodes: readonly ConcreteNode[]) =>
nodes.map((n) => ({ text: n.id })),
),
},
} as unknown as ContextEnvironment;

const tracer = {
logEvent: vi.fn(),
} as unknown as ContextTracer;

const result = await render(
mockNodes,
orchestrator,
sidecar,
tracer,
env,
new Map(),
0,
new Set(),
);

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const surviving = result.history.map((c: any) => c.text);
// Loose Boundary: A (10k), B (20k), C (40k). Total = 70k.
// Adding C pushes rolling total (70k) above retainedTokens (65k).
// Under loose policy, C survives. D is strictly older and drops.
expect(surviving).toEqual(['C', 'B', 'A']); // D is dropped
});

it('drops nodes that are STRICTLY older than the boundary node', async () => {
const mockNodes: ConcreteNode[] = [
{
id: 'A',
type: NodeType.USER_PROMPT,
payload: {} as Part,
} as unknown as ConcreteNode,
{
id: 'B',
type: NodeType.AGENT_THOUGHT,
payload: {} as Part,
} as unknown as ConcreteNode,
{
id: 'C',
type: NodeType.USER_PROMPT,
payload: {} as Part,
} as unknown as ConcreteNode,
];

const tokenMap: Record<string, number> = {
C: 40000,
B: 40000,
A: 10000,
};

const orchestrator = {
executeTriggerSync: vi.fn(async (trigger, nodes, agedOutNodes) =>
nodes.filter((n: ConcreteNode) => !agedOutNodes.has(n.id)),
),
} as unknown as PipelineOrchestrator;

const sidecar = {
config: {
budget: { maxTokens: 150000, retainedTokens: 65000 },
},
} as unknown as ContextProfile;

const currentTokens = 160000;

const env = {
llmClient: {
countTokens: vi.fn().mockResolvedValue({ totalTokens: 1000 }),
},
tokenCalculator: {
calculateConcreteListTokens: vi.fn((nodes) => {
if (nodes.length === 1) return tokenMap[nodes[0].id];
return currentTokens;
}),
calculateTokenBreakdown: vi.fn(() => ({})),
},
graphMapper: {
fromGraph: vi.fn((nodes: readonly ConcreteNode[]) =>
nodes.map((n) => ({ text: n.id })),
),
},
} as unknown as ContextEnvironment;

const tracer = {
logEvent: vi.fn(),
} as unknown as ContextTracer;

const result = await render(
mockNodes,
orchestrator,
sidecar,
tracer,
env,
new Map(),
0,
new Set(),
);

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const surviving = result.history.map((c: any) => c.text);
// C(40k), B(40k). Adding B pushes total to 80k. B is the boundary node and survives. A drops.
expect(surviving).toEqual(['B', 'C']); // A is dropped
});
});
8 changes: 7 additions & 1 deletion packages/core/src/context/graph/render.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import type { ContextTracer } from '../tracer.js';
import type { ContextProfile } from '../config/profiles.js';
import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import type { ContextEnvironment } from '../pipeline/environment.js';
import { performCalibration } from '../utils/tokenCalibration.js';

/**
* Maps the Episodic Context Graph back into a raw Gemini Content[] array for transmission.
Expand Down Expand Up @@ -68,6 +69,7 @@ export async function render(
tracer.logEvent('Render', 'Render Context for LLM', {
renderedContext: contents,
});
performCalibration(env, visibleNodes, contents);
return { history: contents, didApplyManagement: false };
}
const targetDelta = currentTokens - sidecar.config.budget.retainedTokens;
Expand All @@ -83,9 +85,12 @@ export async function render(
// Start from newest and count backwards
for (let i = nodes.length - 1; i >= 0; i--) {
const node = nodes[i];
const priorTokens = rollingTokens;
const nodeTokens = env.tokenCalculator.calculateConcreteListTokens([node]);
rollingTokens += nodeTokens;
if (rollingTokens > sidecar.config.budget.retainedTokens) {

// Loose Boundary Policy: Keep the node that crosses the boundary
if (priorTokens > sidecar.config.budget.retainedTokens) {
Comment thread
joshualitt marked this conversation as resolved.
agedOutNodes.add(node.id);
}
}
Expand Down Expand Up @@ -113,5 +118,6 @@ export async function render(
tracer.logEvent('Render', 'Render Sanitized Context for LLM', {
renderedContextSanitized: contents,
});
performCalibration(env, visibleNodes, contents);
return { history: contents, didApplyManagement: true };
}
4 changes: 4 additions & 0 deletions packages/core/src/context/initializer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,10 @@ export async function initializeContextManager(
tracer,
4,
eventBus,
{
calibrateTokenCalculation:
!!process.env['GEMINI_CONTEXT_CALIBRATE_TOKEN_CALCULATIONS'],
},
);

const orchestrator = new PipelineOrchestrator(
Expand Down
5 changes: 5 additions & 0 deletions packages/core/src/context/pipeline/environment.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,10 @@ import type { ContextGraphMapper } from '../graph/mapper.js';

export type { ContextTracer, ContextEventBus };

export interface RenderOptions {
calibrateTokenCalculation?: boolean;
}

export interface ContextEnvironment {
readonly llmClient: BaseLlmClient;
readonly promptId: string;
Expand All @@ -26,4 +30,5 @@ export interface ContextEnvironment {
readonly inbox: LiveInbox;
readonly behaviorRegistry: NodeBehaviorRegistry;
readonly graphMapper: ContextGraphMapper;
readonly renderOptions?: RenderOptions;
}
3 changes: 2 additions & 1 deletion packages/core/src/context/pipeline/environmentImpl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import type { ContextTracer } from '../tracer.js';
import type { ContextEnvironment } from './environment.js';
import type { ContextEnvironment, RenderOptions } from './environment.js';
import type { ContextEventBus } from '../eventBus.js';
import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
import { LiveInbox } from './inbox.js';
Expand All @@ -29,6 +29,7 @@ export class ContextEnvironmentImpl implements ContextEnvironment {
readonly tracer: ContextTracer,
readonly charsPerToken: number,
readonly eventBus: ContextEventBus,
readonly renderOptions?: RenderOptions,
) {
this.behaviorRegistry = new NodeBehaviorRegistry();
registerBuiltInBehaviors(this.behaviorRegistry);
Expand Down
Loading
Loading