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