mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-27 05:24:34 -07:00
feat(core): introduce decoupled ContextManager and Sidecar architecture (#24752)
This commit is contained in:
joshualitt
@@ -0,0 +1,94 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
||||
import { loadContextManagementConfig } from './configLoader.js';
|
||||
import { defaultContextProfile } from './profiles.js';
|
||||
import { ContextProcessorRegistry } from './registry.js';
|
||||
import * as fs from 'node:fs/promises';
|
||||
import * as path from 'node:path';
|
||||
import * as os from 'node:os';
|
||||
import type { Config } from '../../config/config.js';
|
||||
import type { JSONSchemaType } from 'ajv';
|
||||
|
||||
describe('SidecarLoader (Real FS)', () => {
|
||||
let tmpDir: string;
|
||||
let registry: ContextProcessorRegistry;
|
||||
let sidecarPath: string;
|
||||
let mockConfig: Config;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'gemini-sidecar-test-'));
|
||||
sidecarPath = path.join(tmpDir, 'sidecar.json');
|
||||
registry = new ContextProcessorRegistry();
|
||||
registry.registerProcessor({
|
||||
id: 'NodeTruncation',
|
||||
schema: {
|
||||
type: 'object',
|
||||
properties: { maxTokens: { type: 'number' } },
|
||||
required: ['maxTokens'],
|
||||
} as unknown as JSONSchemaType<{ maxTokens: number }>,
|
||||
});
|
||||
|
||||
mockConfig = {
|
||||
getExperimentalContextManagementConfig: () => sidecarPath,
|
||||
} as unknown as Config;
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await fs.rm(tmpDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
it('returns default profile if file does not exist', async () => {
|
||||
const result = await loadContextManagementConfig(mockConfig, registry);
|
||||
expect(result).toBe(defaultContextProfile);
|
||||
});
|
||||
|
||||
it('returns default profile if file exists but is 0 bytes', async () => {
|
||||
await fs.writeFile(sidecarPath, '');
|
||||
const result = await loadContextManagementConfig(mockConfig, registry);
|
||||
expect(result).toBe(defaultContextProfile);
|
||||
});
|
||||
|
||||
it('returns parsed config if file is valid', async () => {
|
||||
const validConfig = {
|
||||
budget: { retainedTokens: 1000, maxTokens: 2000 },
|
||||
processorOptions: {
|
||||
myTruncation: {
|
||||
type: 'NodeTruncation',
|
||||
options: { maxTokens: 500 },
|
||||
},
|
||||
},
|
||||
};
|
||||
await fs.writeFile(sidecarPath, JSON.stringify(validConfig));
|
||||
const result = await loadContextManagementConfig(mockConfig, registry);
|
||||
expect(result.config.budget?.maxTokens).toBe(2000);
|
||||
expect(result.config.processorOptions?.['myTruncation']).toBeDefined();
|
||||
});
|
||||
|
||||
it('throws validation error if processorOptions contains invalid data for the schema', async () => {
|
||||
const invalidConfig = {
|
||||
budget: { retainedTokens: 1000, maxTokens: 2000 },
|
||||
processorOptions: {
|
||||
myTruncation: {
|
||||
type: 'NodeTruncation',
|
||||
options: { maxTokens: 'this should be a number' },
|
||||
},
|
||||
},
|
||||
};
|
||||
await fs.writeFile(sidecarPath, JSON.stringify(invalidConfig));
|
||||
await expect(
|
||||
loadContextManagementConfig(mockConfig, registry),
|
||||
).rejects.toThrow('Validation error');
|
||||
});
|
||||
|
||||
it('throws validation error if file is empty whitespace', async () => {
|
||||
await fs.writeFile(sidecarPath, ' \n ');
|
||||
await expect(
|
||||
loadContextManagementConfig(mockConfig, registry),
|
||||
).rejects.toThrow('Unexpected end of JSON input');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,90 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../../config/config.js';
|
||||
import * as fsSync from 'node:fs';
|
||||
import * as fs from 'node:fs/promises';
|
||||
import type { ContextManagementConfig } from './types.js';
|
||||
import { defaultContextProfile, type ContextProfile } from './profiles.js';
|
||||
import { SchemaValidator } from '../../utils/schemaValidator.js';
|
||||
import { getContextManagementConfigSchema } from './schema.js';
|
||||
import type { ContextProcessorRegistry } from './registry.js';
|
||||
import { getErrorMessage } from '../../utils/errors.js';
|
||||
|
||||
/**
|
||||
* Loads and validates a sidecar config from a specific file path.
|
||||
* Throws an error if the file cannot be read, parsed, or fails schema validation.
|
||||
*/
|
||||
async function loadConfigFromFile(
|
||||
sidecarPath: string,
|
||||
registry: ContextProcessorRegistry,
|
||||
): Promise<ContextProfile> {
|
||||
const fileContent = await fs.readFile(sidecarPath, 'utf8');
|
||||
let parsed: unknown;
|
||||
try {
|
||||
parsed = JSON.parse(fileContent);
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Failed to parse Sidecar configuration file at ${sidecarPath}: ${getErrorMessage(
|
||||
error,
|
||||
)}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Validate the complete structure, including deep options
|
||||
const validationError = SchemaValidator.validate(
|
||||
getContextManagementConfigSchema(registry),
|
||||
parsed,
|
||||
);
|
||||
|
||||
if (validationError) {
|
||||
throw new Error(
|
||||
`Invalid sidecar configuration in ${sidecarPath}. Validation error: ${validationError}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Extract strictly what we need.
|
||||
// Why this unsafe cast is acceptable:
|
||||
// SchemaValidator just ran \`getSidecarConfigSchema(registry)\` against \`parsed\`.
|
||||
// That function dynamically maps the \`processorOptions\` to strict JSON schema definitions,
|
||||
// so we know with absolute certainty at runtime that \`parsed\` conforms to this shape.
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
const validConfig = parsed as ContextManagementConfig;
|
||||
return {
|
||||
...defaultContextProfile,
|
||||
config: {
|
||||
...defaultContextProfile.config,
|
||||
...(validConfig.budget ? { budget: validConfig.budget } : {}),
|
||||
...(validConfig.processorOptions
|
||||
? { processorOptions: validConfig.processorOptions }
|
||||
: {}),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a Sidecar JSON graph from the experimental config file path or defaults.
|
||||
* If a config file is present but invalid, this will THROW to prevent silent misconfiguration.
|
||||
*/
|
||||
export async function loadContextManagementConfig(
|
||||
config: Config,
|
||||
registry: ContextProcessorRegistry,
|
||||
): Promise<ContextProfile> {
|
||||
const sidecarPath = config.getExperimentalContextManagementConfig();
|
||||
|
||||
if (sidecarPath && fsSync.existsSync(sidecarPath)) {
|
||||
const size = fsSync.statSync(sidecarPath).size;
|
||||
// If the file exists but is completely empty (0 bytes), it's safe to fallback.
|
||||
if (size === 0) {
|
||||
return defaultContextProfile;
|
||||
}
|
||||
|
||||
// If the file has content, enforce strict validation and throw on failure.
|
||||
return loadConfigFromFile(sidecarPath, registry);
|
||||
}
|
||||
|
||||
return defaultContextProfile;
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type {
|
||||
AsyncPipelineDef,
|
||||
ContextManagementConfig,
|
||||
PipelineDef,
|
||||
} from './types.js';
|
||||
import type { ContextEnvironment } from '../pipeline/environment.js';
|
||||
|
||||
// Import factories
|
||||
import { createToolMaskingProcessor } from '../processors/toolMaskingProcessor.js';
|
||||
import { createBlobDegradationProcessor } from '../processors/blobDegradationProcessor.js';
|
||||
import { createNodeTruncationProcessor } from '../processors/nodeTruncationProcessor.js';
|
||||
import { createNodeDistillationProcessor } from '../processors/nodeDistillationProcessor.js';
|
||||
import { createStateSnapshotProcessor } from '../processors/stateSnapshotProcessor.js';
|
||||
import { createStateSnapshotAsyncProcessor } from '../processors/stateSnapshotAsyncProcessor.js';
|
||||
|
||||
/**
|
||||
* Helper to safely merge static default options with dynamically loaded
|
||||
* JSON overrides from the SidecarConfig.
|
||||
*
|
||||
* Why the unsafe cast is acceptable here:
|
||||
* Before the \`config\` object ever reaches this function, \`SidecarLoader.ts\`
|
||||
* passes the raw JSON through \`SchemaValidator\`. The schema dynamically generates
|
||||
* a \`oneOf\` map linking every \`type\` discriminator to its corresponding processor
|
||||
* schema definition. By the time we access \`options\` here, its shape has been
|
||||
* strictly validated against the corresponding Zod/JSONSchema definition at runtime,
|
||||
* making the generic cast to \`<T>\` structurally safe.
|
||||
*/
|
||||
function resolveProcessorOptions<T>(
|
||||
config: ContextManagementConfig | undefined,
|
||||
id: string,
|
||||
defaultOptions: T,
|
||||
): T {
|
||||
if (config?.processorOptions && config.processorOptions[id]) {
|
||||
return {
|
||||
...defaultOptions,
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
...(config.processorOptions[id].options as T),
|
||||
};
|
||||
}
|
||||
return defaultOptions;
|
||||
}
|
||||
|
||||
export interface ContextProfile {
|
||||
config: ContextManagementConfig;
|
||||
buildPipelines: (
|
||||
env: ContextEnvironment,
|
||||
config?: ContextManagementConfig,
|
||||
) => PipelineDef[];
|
||||
buildAsyncPipelines: (
|
||||
env: ContextEnvironment,
|
||||
config?: ContextManagementConfig,
|
||||
) => AsyncPipelineDef[];
|
||||
}
|
||||
|
||||
/**
|
||||
* The standard default context management profile.
|
||||
* Optimized for safety, precision, and reliable summarization.
|
||||
*/
|
||||
export const defaultContextProfile: ContextProfile = {
|
||||
config: {
|
||||
budget: {
|
||||
retainedTokens: 65000,
|
||||
maxTokens: 150000,
|
||||
},
|
||||
},
|
||||
|
||||
buildPipelines: (
|
||||
env: ContextEnvironment,
|
||||
config?: ContextManagementConfig,
|
||||
): PipelineDef[] =>
|
||||
// Helper to merge default options with dynamically loaded processorOptions by ID
|
||||
[
|
||||
{
|
||||
name: 'Immediate Sanitization',
|
||||
triggers: ['new_message'],
|
||||
processors: [
|
||||
createToolMaskingProcessor(
|
||||
'ToolMasking',
|
||||
env,
|
||||
resolveProcessorOptions(config, 'ToolMasking', {
|
||||
stringLengthThresholdTokens: 8000,
|
||||
}),
|
||||
),
|
||||
createBlobDegradationProcessor('BlobDegradation', env), // No options
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Normalization',
|
||||
triggers: ['retained_exceeded'],
|
||||
processors: [
|
||||
createNodeTruncationProcessor(
|
||||
'NodeTruncation',
|
||||
env,
|
||||
resolveProcessorOptions(config, 'NodeTruncation', {
|
||||
maxTokensPerNode: 3000,
|
||||
}),
|
||||
),
|
||||
createNodeDistillationProcessor(
|
||||
'NodeDistillation',
|
||||
env,
|
||||
resolveProcessorOptions(config, 'NodeDistillation', {
|
||||
nodeThresholdTokens: 5000,
|
||||
}),
|
||||
),
|
||||
],
|
||||
},
|
||||
{
|
||||
name: 'Emergency Backstop',
|
||||
triggers: ['gc_backstop'],
|
||||
processors: [
|
||||
createStateSnapshotProcessor(
|
||||
'StateSnapshotSync',
|
||||
env,
|
||||
resolveProcessorOptions(config, 'StateSnapshotSync', {
|
||||
target: 'max',
|
||||
}),
|
||||
),
|
||||
],
|
||||
},
|
||||
],
|
||||
buildAsyncPipelines: (
|
||||
env: ContextEnvironment,
|
||||
config?: ContextManagementConfig,
|
||||
): AsyncPipelineDef[] => [
|
||||
{
|
||||
name: 'Async Background GC',
|
||||
triggers: ['nodes_aged_out'],
|
||||
processors: [
|
||||
createStateSnapshotAsyncProcessor(
|
||||
'StateSnapshotAsync',
|
||||
env,
|
||||
resolveProcessorOptions(config, 'StateSnapshotAsync', {
|
||||
type: 'accumulate',
|
||||
}),
|
||||
),
|
||||
],
|
||||
},
|
||||
],
|
||||
};
|
||||
@@ -0,0 +1,42 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { JSONSchemaType } from 'ajv';
|
||||
|
||||
export interface ContextProcessorDef<T = unknown> {
|
||||
readonly id: string;
|
||||
readonly schema: JSONSchemaType<T>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Registry for validating declarative sidecar configuration schemas.
|
||||
* (Dynamic instantiation has been replaced by static ContextProfiles)
|
||||
*/
|
||||
export class ContextProcessorRegistry {
|
||||
private readonly processors = new Map<string, ContextProcessorDef>();
|
||||
|
||||
registerProcessor<T>(def: ContextProcessorDef<T>) {
|
||||
// Erasing the type.
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
this.processors.set(def.id, def as unknown as ContextProcessorDef<unknown>);
|
||||
}
|
||||
|
||||
getSchema(id: string): object | undefined {
|
||||
return this.processors.get(id)?.schema;
|
||||
}
|
||||
|
||||
getSchemaDefs(): ContextProcessorDef[] {
|
||||
const defs = [];
|
||||
for (const def of this.processors.values()) {
|
||||
if (def.schema) defs.push({ id: def.id, schema: def.schema });
|
||||
}
|
||||
return defs;
|
||||
}
|
||||
|
||||
clear() {
|
||||
this.processors.clear();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessorRegistry } from './registry.js';
|
||||
|
||||
export function getContextManagementConfigSchema(
|
||||
registry: ContextProcessorRegistry,
|
||||
) {
|
||||
// We use a registry to deeply validate processor overrides.
|
||||
// We do this by generating a `oneOf` list that matches the `type` discriminator
|
||||
// to the specific processor `options` schema.
|
||||
const processorOptionSchemas = registry.getSchemaDefs().map((def) => ({
|
||||
type: 'object',
|
||||
required: ['type', 'options'],
|
||||
properties: {
|
||||
type: { const: def.id },
|
||||
options: def.schema,
|
||||
},
|
||||
}));
|
||||
|
||||
return {
|
||||
$schema: 'http://json-schema.org/draft-07/schema#',
|
||||
title: 'ContextManagementConfig',
|
||||
description: 'The Hyperparameter schema for a Context Profile.',
|
||||
type: 'object',
|
||||
properties: {
|
||||
budget: {
|
||||
type: 'object',
|
||||
description: 'Defines the token ceilings and limits for the pipeline.',
|
||||
required: ['retainedTokens', 'maxTokens'],
|
||||
properties: {
|
||||
retainedTokens: {
|
||||
type: 'number',
|
||||
description:
|
||||
'The ideal token count the pipeline tries to shrink down to.',
|
||||
},
|
||||
maxTokens: {
|
||||
type: 'number',
|
||||
description:
|
||||
'The absolute maximum token count allowed before synchronous truncation kicks in.',
|
||||
},
|
||||
},
|
||||
},
|
||||
processorOptions: {
|
||||
type: 'object',
|
||||
description:
|
||||
'Named hyperparameter configurations for ContextProcessors and AsyncProcessors.',
|
||||
additionalProperties: { oneOf: processorOptionSchemas },
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { ContextProcessor, AsyncContextProcessor } from '../pipeline.js';
|
||||
|
||||
export type PipelineTrigger =
|
||||
| 'new_message'
|
||||
| 'retained_exceeded'
|
||||
| 'gc_backstop'
|
||||
| 'nodes_added'
|
||||
| 'nodes_aged_out'
|
||||
| { type: 'timer'; intervalMs: number };
|
||||
|
||||
export interface PipelineDef {
|
||||
name: string;
|
||||
triggers: PipelineTrigger[];
|
||||
processors: ContextProcessor[];
|
||||
}
|
||||
|
||||
export interface AsyncPipelineDef {
|
||||
name: string;
|
||||
triggers: PipelineTrigger[];
|
||||
processors: AsyncContextProcessor[];
|
||||
}
|
||||
|
||||
export interface ContextBudget {
|
||||
retainedTokens: number;
|
||||
maxTokens: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* The Data-Driven Schema for the Context Manager.
|
||||
*/
|
||||
export interface ContextManagementConfig {
|
||||
/** Defines the token ceilings and limits for the pipeline. */
|
||||
budget: ContextBudget;
|
||||
|
||||
/**
|
||||
* Dynamic hyperparameter overrides for individual ContextProcessors and AsyncProcessors.
|
||||
* Keys are named identifiers (e.g. "gentleTruncation").
|
||||
*/
|
||||
processorOptions?: Record<string, { type: string; options: unknown }>;
|
||||
}
|
||||
Reference in New Issue
Block a user