File size: 5,260 Bytes
7a1ad33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | /**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { ContextManager } from './contextManager.js';
import {
createMockEnvironment,
createDummyNode,
} from './testing/contextTestUtils.js';
import type { ContextProfile } from './config/profiles.js';
import { NodeType, type ConcreteNode } from './graph/types.js';
import type { PipelineOrchestrator } from './pipeline/orchestrator.js';
import type { AgentChatHistory } from '../core/agentChatHistory.js';
import type { AdvancedTokenCalculator } from './utils/contextTokenCalculator.js';
import type { ContextManagementConfig } from './config/types.js';
import type { ContextEnvironment } from './pipeline/environment.js';
import type { ContextWorkingBufferImpl } from './pipeline/contextWorkingBuffer.js';
describe('ContextManager - Multi-stage and Incremental GC', () => {
let mockEnv: ReturnType<typeof createMockEnvironment>;
let mockOrchestrator: PipelineOrchestrator;
let mockChatHistory: AgentChatHistory;
let mockAdvancedTokenCalculator: AdvancedTokenCalculator;
beforeEach(() => {
mockEnv = createMockEnvironment();
mockOrchestrator = {
setNodeProvider: vi.fn(),
waitForPipelines: vi.fn().mockResolvedValue(undefined),
executeTriggerSync: vi
.fn()
.mockImplementation(async (trigger, buffer) => buffer),
executeIngestionPipeline: vi
.fn()
.mockImplementation(async (nodes) => nodes),
shutdown: vi.fn(),
} as unknown as PipelineOrchestrator;
mockChatHistory = {
all: vi.fn().mockReturnValue([]),
getHistory: vi.fn().mockReturnValue([]),
get: vi.fn().mockReturnValue([]),
subscribe: vi.fn(),
} as unknown as AgentChatHistory;
mockAdvancedTokenCalculator = {
getRawBaseUnits: vi.fn().mockReturnValue(0),
getRawBaseUnitsForContent: vi.fn().mockReturnValue(0),
calculateTokensAndBaseUnits: vi.fn(),
} as unknown as AdvancedTokenCalculator;
});
const setupManager = (config: ContextManagementConfig) => {
const sidecar: ContextProfile = {
name: 'test',
config,
buildPipelines: () => [],
buildAsyncPipelines: () => [],
};
return new ContextManager(
sidecar,
mockEnv as unknown as ContextEnvironment,
mockEnv.tracer,
mockOrchestrator,
mockChatHistory,
mockAdvancedTokenCalculator,
);
};
it('should emit NormalizeNeeded when normalizedTokens budget is exceeded', async () => {
const manager = setupManager({
budget: {
retainedTokens: 100,
normalizedTokens: 150,
maxTokens: 300,
},
} as unknown as ContextManagementConfig);
const normalizeSpy = vi.fn();
mockEnv.eventBus.onNormalizeNeeded(normalizeSpy);
const consolidationSpy = vi.fn();
mockEnv.eventBus.onConsolidationNeeded(consolidationSpy);
// Mock token calculator for evaluateTriggers
mockEnv.tokenCalculator.calculateConcreteListTokens = vi
.fn()
.mockImplementation((nodes: ConcreteNode[]) =>
nodes.reduce(
(sum: number, n: ConcreteNode) =>
// Look for the mock tokens we attached to the dummy node
sum + ((n as unknown as { _mockTokens: number })._mockTokens || 0),
0,
),
);
const createNodeWithTokens = (
id: string,
type: NodeType,
tokens: number,
) => {
const node = createDummyNode(id, type);
// @ts-expect-error - attaching mock tokens for test
node._mockTokens = tokens;
return node;
};
// Create 4 nodes, each 80 tokens. Total = 320 tokens.
// Node 1 (oldest): prior=240. 240 > 150 -> Normalization (Archiving trigger)
// Node 2: prior=160. 160 > 150 -> Normalization
// Node 3: prior=80. 80 <= 100 -> Retained
// Node 4 (newest): prior=0. 0 <= 100 -> Retained
const nodes = [
createNodeWithTokens('ep1', NodeType.USER_PROMPT, 80),
createNodeWithTokens('ep2', NodeType.AGENT_THOUGHT, 80),
createNodeWithTokens('ep3', NodeType.TOOL_EXECUTION, 80),
createNodeWithTokens('ep4', NodeType.TOOL_EXECUTION, 80),
];
// @ts-expect-error - access private method for testing
manager.buffer = { nodes } as unknown as ContextWorkingBufferImpl;
// Trigger evaluation manually with a dummy "new node" to bypass the empty check
// @ts-expect-error - access private method for testing
await manager.evaluateTriggers(nodes, new Set([nodes[3].id]), new Set());
// Nodes 3 and 4 are retained.
// Node 2 and Node 1 both fall out of normalizedTokens (160 > 150, 240 > 150).
// Therefore they should trigger NormalizeNeeded. They should NOT trigger ConsolidationNeeded
// because they exceeded normalized budget, so they skip the retained fallback.
expect(consolidationSpy).not.toHaveBeenCalled();
expect(normalizeSpy).toHaveBeenCalledOnce();
const normalizeEvent = normalizeSpy.mock.calls[0][0];
expect(normalizeEvent.targetNodeIds.has(nodes[0].id)).toBe(true);
expect(normalizeEvent.targetNodeIds.has(nodes[1].id)).toBe(true);
expect(normalizeEvent.targetNodeIds.has(nodes[2].id)).toBe(false);
});
});
|