mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-28 05:55:17 -07:00
feat(core): introduce decoupled ContextManager and Sidecar architecture (#24752)
This commit is contained in:
joshualitt
File diff suppressed because one or more lines are too long
@@ -0,0 +1,246 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest';
|
||||
import { SimulationHarness } from './simulationHarness.js';
|
||||
import { createMockLlmClient } from '../testing/contextTestUtils.js';
|
||||
import type { ContextProfile } from '../config/profiles.js';
|
||||
import { createToolMaskingProcessor } from '../processors/toolMaskingProcessor.js';
|
||||
import { createBlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
|
||||
import { createStateSnapshotProcessor } from '../processors/stateSnapshotProcessor.js';
|
||||
import { createHistoryTruncationProcessor } from '../processors/historyTruncationProcessor.js';
|
||||
import { createStateSnapshotAsyncProcessor } from '../processors/stateSnapshotAsyncProcessor.js';
|
||||
|
||||
expect.addSnapshotSerializer({
|
||||
test: (val) =>
|
||||
typeof val === 'string' &&
|
||||
(/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(
|
||||
val,
|
||||
) ||
|
||||
/^\/tmp\/sim/.test(val)), // Mask temp directories and UUIDs
|
||||
print: (val) =>
|
||||
typeof val === 'string' && /^\/tmp\/sim/.test(val)
|
||||
? '"<MOCKED_DIR>"'
|
||||
: '"<UUID>"',
|
||||
});
|
||||
|
||||
describe('System Lifecycle Golden Tests', () => {
|
||||
beforeAll(() => {
|
||||
vi.spyOn(Math, 'random').mockReturnValue(0.5);
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
const getAggressiveConfig = (): ContextProfile => ({
|
||||
config: {
|
||||
budget: { maxTokens: 1000, retainedTokens: 500 }, // Extremely tight limits
|
||||
},
|
||||
buildPipelines: (env) => [
|
||||
{
|
||||
name: 'Pressure Relief', // Emits from eventBus 'retained_exceeded'
|
||||
triggers: ['retained_exceeded'],
|
||||
processors: [
|
||||
createBlobDegradationProcessor('BlobDegradationProcessor', env),
|
||||
createToolMaskingProcessor('ToolMaskingProcessor', env, {
|
||||
stringLengthThresholdTokens: 50,
|
||||
}),
|
||||
createStateSnapshotProcessor('StateSnapshotProcessor', env, {}),
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Immediate Sanitization', // The magic string the projector is hardcoded to use
|
||||
triggers: ['retained_exceeded'],
|
||||
processors: [
|
||||
createHistoryTruncationProcessor(
|
||||
'HistoryTruncationProcessor',
|
||||
env,
|
||||
{},
|
||||
),
|
||||
],
|
||||
},
|
||||
],
|
||||
buildAsyncPipelines: (env) => [
|
||||
{
|
||||
name: 'Async',
|
||||
triggers: ['nodes_aged_out'],
|
||||
processors: [
|
||||
createStateSnapshotAsyncProcessor(
|
||||
'StateSnapshotAsyncProcessor',
|
||||
env,
|
||||
{},
|
||||
),
|
||||
],
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const mockLlmClient = createMockLlmClient([
|
||||
'<MOCKED_STATE_SNAPSHOT_SUMMARY>',
|
||||
]);
|
||||
|
||||
it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => {
|
||||
const harness = await SimulationHarness.create(
|
||||
getAggressiveConfig(),
|
||||
mockLlmClient,
|
||||
);
|
||||
|
||||
// Turn 0: System Prompt
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'System Instructions' }] },
|
||||
{ role: 'model', parts: [{ text: 'Ack.' }] },
|
||||
]);
|
||||
|
||||
// Turn 1: Normal conversation
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Hello!' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi, how can I help?' }] },
|
||||
]);
|
||||
|
||||
// Turn 2: Massive Tool Output (Should trigger ToolMaskingProcessor in background)
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Read the logs.' }] },
|
||||
{
|
||||
role: 'model',
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: 'run_shell_command',
|
||||
args: { cmd: 'cat server.log' },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: 'run_shell_command',
|
||||
response: { output: 'LOG '.repeat(5000) },
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'The logs are very long.' }] },
|
||||
]);
|
||||
|
||||
// Turn 3: Multi-modal blob (Should trigger BlobDegradationProcessor)
|
||||
await harness.simulateTurn([
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{ text: 'Look at this architecture diagram:' },
|
||||
{
|
||||
inlineData: {
|
||||
mimeType: 'image/png',
|
||||
data: 'fake_base64_data_'.repeat(1000),
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
{ role: 'model', parts: [{ text: 'Nice diagram.' }] },
|
||||
]);
|
||||
|
||||
// Turn 4: More conversation to trigger StateSnapshot
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Can we refactor?' }] },
|
||||
{ role: 'model', parts: [{ text: 'Yes we can.' }] },
|
||||
]);
|
||||
|
||||
// Get final state
|
||||
const goldenState = await harness.getGoldenState();
|
||||
|
||||
// In a perfectly functioning opportunistic system, the token trajectory should show
|
||||
// the massive spikes in Turn 2 and 3 being immediately resolved by the background tasks.
|
||||
// The final projection should fit neatly under the Max Tokens limit.
|
||||
|
||||
expect(goldenState).toMatchSnapshot();
|
||||
});
|
||||
|
||||
it('Scenario 2: Under Budget (No Modifications)', async () => {
|
||||
const generousConfig: ContextProfile = {
|
||||
config: {
|
||||
budget: { maxTokens: 100000, retainedTokens: 50000 },
|
||||
},
|
||||
buildPipelines: () => [],
|
||||
buildAsyncPipelines: () => [],
|
||||
};
|
||||
|
||||
const harness = await SimulationHarness.create(
|
||||
generousConfig,
|
||||
mockLlmClient,
|
||||
);
|
||||
|
||||
// Turn 0: System Prompt
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'System Instructions' }] },
|
||||
{ role: 'model', parts: [{ text: 'Ack.' }] },
|
||||
]);
|
||||
|
||||
// Turn 1: Normal conversation
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'Hello!' }] },
|
||||
{ role: 'model', parts: [{ text: 'Hi, how can I help?' }] },
|
||||
]);
|
||||
|
||||
const goldenState = await harness.getGoldenState();
|
||||
|
||||
// Total tokens should cleanly match character count with no synthetic nodes
|
||||
expect(goldenState).toMatchSnapshot();
|
||||
});
|
||||
|
||||
it('Scenario 3: Async-Driven Background GC', async () => {
|
||||
const gcConfig: ContextProfile = {
|
||||
config: {
|
||||
budget: { maxTokens: 200, retainedTokens: 100 },
|
||||
},
|
||||
buildPipelines: () => [],
|
||||
buildAsyncPipelines: (env) => [
|
||||
{
|
||||
name: 'Async',
|
||||
triggers: ['nodes_aged_out'],
|
||||
processors: [
|
||||
createStateSnapshotAsyncProcessor(
|
||||
'StateSnapshotAsyncProcessor',
|
||||
env,
|
||||
{},
|
||||
),
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const harness = await SimulationHarness.create(gcConfig, mockLlmClient);
|
||||
|
||||
// Turn 0
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'A'.repeat(50) }] },
|
||||
{ role: 'model', parts: [{ text: 'B'.repeat(50) }] },
|
||||
]);
|
||||
|
||||
// Turn 1 (Should trigger StateSnapshotasync pipeline because we exceed 100 retainedTokens)
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'C'.repeat(50) }] },
|
||||
{ role: 'model', parts: [{ text: 'D'.repeat(50) }] },
|
||||
]);
|
||||
|
||||
// Give the async background pipeline an extra beat to complete its async execution and emit variants
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
// Turn 2
|
||||
await harness.simulateTurn([
|
||||
{ role: 'user', parts: [{ text: 'E'.repeat(50) }] },
|
||||
{ role: 'model', parts: [{ text: 'F'.repeat(50) }] },
|
||||
]);
|
||||
|
||||
const goldenState = await harness.getGoldenState();
|
||||
|
||||
// We should see ROLLING_SUMMARY nodes injected into the graph, proving the async pipeline ran in the background
|
||||
expect(goldenState).toMatchSnapshot();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,157 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { ContextManager } from '../contextManager.js';
|
||||
import { AgentChatHistory } from '../../core/agentChatHistory.js';
|
||||
import type { Content } from '@google/genai';
|
||||
import type { ContextProfile } from '../config/profiles.js';
|
||||
import { ContextEnvironmentImpl } from '../pipeline/environmentImpl.js';
|
||||
import { ContextTracer } from '../tracer.js';
|
||||
import { ContextEventBus } from '../eventBus.js';
|
||||
import { PipelineOrchestrator } from '../pipeline/orchestrator.js';
|
||||
import { debugLogger } from '../../utils/debugLogger.js';
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
|
||||
export interface TurnSummary {
|
||||
turnIndex: number;
|
||||
tokensBeforeBackground: number;
|
||||
tokensAfterBackground: number;
|
||||
}
|
||||
|
||||
export class SimulationHarness {
|
||||
readonly chatHistory: AgentChatHistory;
|
||||
contextManager!: ContextManager;
|
||||
env!: ContextEnvironmentImpl;
|
||||
orchestrator!: PipelineOrchestrator;
|
||||
readonly eventBus: ContextEventBus;
|
||||
config!: ContextProfile;
|
||||
private tracer!: ContextTracer;
|
||||
private currentTurnIndex = 0;
|
||||
private tokenTrajectory: TurnSummary[] = [];
|
||||
|
||||
static async create(
|
||||
config: ContextProfile,
|
||||
mockLlmClient: BaseLlmClient,
|
||||
mockTempDir = '/tmp/sim',
|
||||
): Promise<SimulationHarness> {
|
||||
const harness = new SimulationHarness();
|
||||
await harness.init(config, mockLlmClient, mockTempDir);
|
||||
return harness;
|
||||
}
|
||||
|
||||
private constructor() {
|
||||
this.chatHistory = new AgentChatHistory();
|
||||
this.eventBus = new ContextEventBus();
|
||||
}
|
||||
|
||||
private async init(
|
||||
config: ContextProfile,
|
||||
mockLlmClient: BaseLlmClient,
|
||||
mockTempDir: string,
|
||||
) {
|
||||
this.config = config;
|
||||
|
||||
this.tracer = new ContextTracer({
|
||||
targetDir: mockTempDir,
|
||||
sessionId: 'sim-session',
|
||||
});
|
||||
this.env = new ContextEnvironmentImpl(
|
||||
mockLlmClient,
|
||||
'sim-prompt',
|
||||
'sim-session',
|
||||
mockTempDir,
|
||||
mockTempDir,
|
||||
this.tracer,
|
||||
1, // 1 char per token average
|
||||
this.eventBus,
|
||||
);
|
||||
|
||||
this.orchestrator = new PipelineOrchestrator(
|
||||
config.buildPipelines(this.env),
|
||||
config.buildAsyncPipelines(this.env),
|
||||
this.env,
|
||||
this.eventBus,
|
||||
this.tracer,
|
||||
);
|
||||
this.contextManager = new ContextManager(
|
||||
config,
|
||||
this.env,
|
||||
this.tracer,
|
||||
this.orchestrator,
|
||||
this.chatHistory,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Simulates a single "Turn" (User input + Model/Tool outputs)
|
||||
* A turn might consist of multiple Content messages (e.g. user prompt -> model call -> user response -> model answer)
|
||||
*/
|
||||
async simulateTurn(messages: Content[]) {
|
||||
// 1. Append the new messages
|
||||
const currentHistory = this.chatHistory.get();
|
||||
this.chatHistory.set([...currentHistory, ...messages]);
|
||||
|
||||
// 2. Measure tokens immediately after append (Before background processing)
|
||||
const tokensBefore = this.env.tokenCalculator.calculateConcreteListTokens(
|
||||
this.contextManager.getNodes(),
|
||||
);
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Tokens BEFORE: ${tokensBefore}`,
|
||||
);
|
||||
|
||||
// 3. Yield to event loop to allow internal async subscribers and orchestrator to finish
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
// 3.1 Simulate what projectCompressedHistory does with the sync handlers
|
||||
let currentView = this.contextManager.getNodes();
|
||||
const currentTokens =
|
||||
this.env.tokenCalculator.calculateConcreteListTokens(currentView);
|
||||
if (
|
||||
this.config.config.budget &&
|
||||
currentTokens > this.config.config.budget.maxTokens
|
||||
) {
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.config.budget.maxTokens}`,
|
||||
);
|
||||
const orchestrator = this.orchestrator;
|
||||
// In the V2 simulation, we trigger the 'gc_backstop' to simulate emergency pressure.
|
||||
// Since contextManager owns its buffer natively, the simulation now properly matches reality
|
||||
// where the manager runs the orchestrator and keeps the resulting modified view.
|
||||
const modifiedView = await orchestrator.executeTriggerSync(
|
||||
'gc_backstop',
|
||||
currentView,
|
||||
new Set(currentView.map((e) => e.id)),
|
||||
new Set<string>(),
|
||||
);
|
||||
|
||||
// In the real system, ContextManager triggers this and retains it.
|
||||
// We will emulate that behavior internally in the test loop for token counting.
|
||||
currentView = modifiedView;
|
||||
}
|
||||
|
||||
// 4. Measure tokens after background processors have processed inboxes
|
||||
const tokensAfter = this.env.tokenCalculator.calculateConcreteListTokens(
|
||||
this.contextManager.getNodes(),
|
||||
);
|
||||
debugLogger.log(
|
||||
`[Turn ${this.currentTurnIndex}] Tokens AFTER: ${tokensAfter}`,
|
||||
);
|
||||
|
||||
this.tokenTrajectory.push({
|
||||
turnIndex: this.currentTurnIndex++,
|
||||
tokensBeforeBackground: tokensBefore,
|
||||
tokensAfterBackground: tokensAfter,
|
||||
});
|
||||
}
|
||||
|
||||
async getGoldenState() {
|
||||
const finalProjection = await this.contextManager.renderHistory();
|
||||
return {
|
||||
tokenTrajectory: this.tokenTrajectory,
|
||||
finalProjection,
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user