mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-20 00:32:31 -07:00
format
This commit is contained in:
@@ -18,9 +18,10 @@ import { ContextWorkingBufferImpl } from './sidecar/contextWorkingBuffer.js';
|
||||
|
||||
export class ContextManager {
|
||||
// The master state containing the pristine graph and current active graph.
|
||||
private buffer: ContextWorkingBufferImpl = ContextWorkingBufferImpl.initialize([]);
|
||||
private buffer: ContextWorkingBufferImpl =
|
||||
ContextWorkingBufferImpl.initialize([]);
|
||||
private pristineNodes: readonly ConcreteNode[] = [];
|
||||
|
||||
|
||||
private readonly eventBus: ContextEventBus;
|
||||
|
||||
// Internal sub-components
|
||||
@@ -48,10 +49,10 @@ export class ContextManager {
|
||||
|
||||
this.eventBus.onPristineHistoryUpdated((event) => {
|
||||
this.pristineNodes = event.nodes;
|
||||
|
||||
|
||||
const existingIds = new Set(this.buffer.nodes.map((n) => n.id));
|
||||
const addedNodes = event.nodes.filter((n) => !existingIds.has(n.id));
|
||||
|
||||
|
||||
if (addedNodes.length > 0) {
|
||||
this.buffer = this.buffer.appendPristineNodes(addedNodes);
|
||||
}
|
||||
|
||||
@@ -23,8 +23,6 @@ export interface IrChunkReceivedEvent {
|
||||
targetNodeIds: Set<string>;
|
||||
}
|
||||
|
||||
|
||||
|
||||
export class ContextEventBus extends EventEmitter {
|
||||
emitPristineHistoryUpdated(event: PristineHistoryUpdatedEvent) {
|
||||
this.emit('PRISTINE_HISTORY_UPDATED', event);
|
||||
|
||||
@@ -10,29 +10,57 @@ import { createDummyNode } from '../testing/contextTestUtils.js';
|
||||
|
||||
describe('ContextWorkingBufferImpl', () => {
|
||||
it('should initialize with a pristine graph correctly', () => {
|
||||
const pristine1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
|
||||
const pristine2 = createDummyNode('ep1', 'AGENT_THOUGHT', 10, undefined, 'p2');
|
||||
|
||||
const pristine1 = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
10,
|
||||
undefined,
|
||||
'p1',
|
||||
);
|
||||
const pristine2 = createDummyNode(
|
||||
'ep1',
|
||||
'AGENT_THOUGHT',
|
||||
10,
|
||||
undefined,
|
||||
'p2',
|
||||
);
|
||||
|
||||
const buffer = ContextWorkingBufferImpl.initialize([pristine1, pristine2]);
|
||||
|
||||
|
||||
expect(buffer.nodes).toHaveLength(2);
|
||||
expect(buffer.getAuditLog()).toHaveLength(0);
|
||||
|
||||
|
||||
// Pristine nodes should point to themselves
|
||||
expect(buffer.getPristineNodes('p1')).toEqual([pristine1]);
|
||||
expect(buffer.getPristineNodes('p2')).toEqual([pristine2]);
|
||||
});
|
||||
|
||||
it('should track 1:1 replacements (e.g., masking) and append to audit log', () => {
|
||||
const pristine1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
|
||||
const pristine1 = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
10,
|
||||
undefined,
|
||||
'p1',
|
||||
);
|
||||
let buffer = ContextWorkingBufferImpl.initialize([pristine1]);
|
||||
|
||||
const maskedNode = createDummyNode('ep1', 'USER_PROMPT', 5, undefined, 'm1');
|
||||
const maskedNode = createDummyNode(
|
||||
'ep1',
|
||||
'USER_PROMPT',
|
||||
5,
|
||||
undefined,
|
||||
'm1',
|
||||
);
|
||||
// Simulate what a processor does
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(maskedNode as any).replacesId = 'p1';
|
||||
|
||||
buffer = buffer.applyProcessorResult('ToolMasking', [pristine1], [maskedNode]);
|
||||
buffer = buffer.applyProcessorResult(
|
||||
'ToolMasking',
|
||||
[pristine1],
|
||||
[maskedNode],
|
||||
);
|
||||
|
||||
expect(buffer.nodes).toHaveLength(1);
|
||||
expect(buffer.nodes[0].id).toBe('m1');
|
||||
@@ -51,17 +79,23 @@ describe('ContextWorkingBufferImpl', () => {
|
||||
const p1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
|
||||
const p2 = createDummyNode('ep1', 'AGENT_THOUGHT', 10, undefined, 'p2');
|
||||
const p3 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p3');
|
||||
|
||||
|
||||
let buffer = ContextWorkingBufferImpl.initialize([p1, p2, p3]);
|
||||
|
||||
const summaryNode = createDummyNode('ep1', 'ROLLING_SUMMARY', 15, undefined, 's1');
|
||||
const summaryNode = createDummyNode(
|
||||
'ep1',
|
||||
'ROLLING_SUMMARY',
|
||||
15,
|
||||
undefined,
|
||||
's1',
|
||||
);
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(summaryNode as any).abstractsIds = ['p1', 'p2'];
|
||||
|
||||
buffer = buffer.applyProcessorResult('Summarizer', [p1, p2], [summaryNode]);
|
||||
|
||||
// p1 and p2 are removed, p3 remains, s1 is added
|
||||
expect(buffer.nodes.map(n => n.id)).toEqual(['p3', 's1']);
|
||||
expect(buffer.nodes.map((n) => n.id)).toEqual(['p3', 's1']);
|
||||
|
||||
// Provenance lookup: The summary node should resolve to both p1 and p2!
|
||||
const roots = buffer.getPristineNodes('s1');
|
||||
@@ -81,7 +115,13 @@ describe('ContextWorkingBufferImpl', () => {
|
||||
buffer = buffer.applyProcessorResult('Masking', [p1], [gen1]);
|
||||
|
||||
// Gen 2: Summarized
|
||||
const gen2 = createDummyNode('ep1', 'ROLLING_SUMMARY', 5, undefined, 'gen2');
|
||||
const gen2 = createDummyNode(
|
||||
'ep1',
|
||||
'ROLLING_SUMMARY',
|
||||
5,
|
||||
undefined,
|
||||
'gen2',
|
||||
);
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(gen2 as any).abstractsIds = ['gen1'];
|
||||
buffer = buffer.applyProcessorResult('Summarizer', [gen1], [gen2]);
|
||||
@@ -103,14 +143,20 @@ describe('ContextWorkingBufferImpl', () => {
|
||||
const p1 = createDummyNode('ep1', 'USER_PROMPT', 10, undefined, 'p1');
|
||||
let buffer = ContextWorkingBufferImpl.initialize([p1]);
|
||||
|
||||
const injected = createDummyNode('ep1', 'SYSTEM_EVENT', 5, undefined, 'injected1');
|
||||
const injected = createDummyNode(
|
||||
'ep1',
|
||||
'SYSTEM_EVENT',
|
||||
5,
|
||||
undefined,
|
||||
'injected1',
|
||||
);
|
||||
// No replacesId or abstractsIds
|
||||
|
||||
buffer = buffer.applyProcessorResult('Injector', [], [injected]);
|
||||
|
||||
expect(buffer.nodes.map(n => n.id)).toEqual(['p1', 'injected1']);
|
||||
expect(buffer.nodes.map((n) => n.id)).toEqual(['p1', 'injected1']);
|
||||
|
||||
// It should root to itself
|
||||
expect(buffer.getPristineNodes('injected1')).toEqual([injected]);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -10,10 +10,10 @@ import type { ConcreteNode } from '../ir/types.js';
|
||||
export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
// The current active graph
|
||||
readonly nodes: readonly ConcreteNode[];
|
||||
|
||||
|
||||
// The AOT pre-calculated provenance index (Current ID -> Pristine IDs)
|
||||
private readonly provenanceMap: ReadonlyMap<string, ReadonlySet<string>>;
|
||||
|
||||
|
||||
// The original immutable pristine nodes mapping
|
||||
private readonly pristineNodesMap: ReadonlyMap<string, ConcreteNode>;
|
||||
|
||||
@@ -24,7 +24,7 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
nodes: readonly ConcreteNode[],
|
||||
pristineNodesMap: ReadonlyMap<string, ConcreteNode>,
|
||||
provenanceMap: ReadonlyMap<string, ReadonlySet<string>>,
|
||||
history: readonly GraphMutation[]
|
||||
history: readonly GraphMutation[],
|
||||
) {
|
||||
this.nodes = nodes;
|
||||
this.pristineNodesMap = pristineNodesMap;
|
||||
@@ -36,20 +36,22 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
* Initializes a brand new ContextWorkingBuffer from a pristine graph.
|
||||
* Every node's provenance points to itself.
|
||||
*/
|
||||
static initialize(pristineNodes: readonly ConcreteNode[]): ContextWorkingBufferImpl {
|
||||
static initialize(
|
||||
pristineNodes: readonly ConcreteNode[],
|
||||
): ContextWorkingBufferImpl {
|
||||
const pristineMap = new Map<string, ConcreteNode>();
|
||||
const initialProvenance = new Map<string, ReadonlySet<string>>();
|
||||
|
||||
|
||||
for (const node of pristineNodes) {
|
||||
pristineMap.set(node.id, node);
|
||||
initialProvenance.set(node.id, new Set([node.id]));
|
||||
}
|
||||
|
||||
|
||||
return new ContextWorkingBufferImpl(
|
||||
pristineNodes,
|
||||
pristineMap,
|
||||
initialProvenance,
|
||||
[] // Empty history
|
||||
[], // Empty history
|
||||
);
|
||||
}
|
||||
|
||||
@@ -57,12 +59,14 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
* Appends newly observed pristine nodes (e.g. from a user message) to the working buffer.
|
||||
* Ensures they are tracked in the pristine map and point to themselves in provenance.
|
||||
*/
|
||||
appendPristineNodes(newNodes: readonly ConcreteNode[]): ContextWorkingBufferImpl {
|
||||
appendPristineNodes(
|
||||
newNodes: readonly ConcreteNode[],
|
||||
): ContextWorkingBufferImpl {
|
||||
if (newNodes.length === 0) return this;
|
||||
|
||||
const newPristineMap = new Map<string, ConcreteNode>(this.pristineNodesMap);
|
||||
const newProvenanceMap = new Map(this.provenanceMap);
|
||||
|
||||
|
||||
for (const node of newNodes) {
|
||||
newPristineMap.set(node.id, node);
|
||||
newProvenanceMap.set(node.id, new Set([node.id]));
|
||||
@@ -72,7 +76,7 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
[...this.nodes, ...newNodes],
|
||||
newPristineMap,
|
||||
newProvenanceMap,
|
||||
[...this.history]
|
||||
[...this.history],
|
||||
);
|
||||
}
|
||||
|
||||
@@ -80,17 +84,19 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
* Generates an entirely new buffer instance by calculating the delta between the processor's input and output.
|
||||
*/
|
||||
applyProcessorResult(
|
||||
processorId: string,
|
||||
inputTargets: readonly ConcreteNode[],
|
||||
outputNodes: readonly ConcreteNode[]
|
||||
processorId: string,
|
||||
inputTargets: readonly ConcreteNode[],
|
||||
outputNodes: readonly ConcreteNode[],
|
||||
): ContextWorkingBufferImpl {
|
||||
const outputIds = new Set(outputNodes.map(n => n.id));
|
||||
const inputIds = new Set(inputTargets.map(n => n.id));
|
||||
const outputIds = new Set(outputNodes.map((n) => n.id));
|
||||
const inputIds = new Set(inputTargets.map((n) => n.id));
|
||||
|
||||
// Calculate diffs
|
||||
const removedIds = inputTargets.filter(n => !outputIds.has(n.id)).map(n => n.id);
|
||||
const addedNodes = outputNodes.filter(n => !inputIds.has(n.id));
|
||||
|
||||
const removedIds = inputTargets
|
||||
.filter((n) => !outputIds.has(n.id))
|
||||
.map((n) => n.id);
|
||||
const addedNodes = outputNodes.filter((n) => !inputIds.has(n.id));
|
||||
|
||||
// Create mutation record
|
||||
const mutation: GraphMutation = {
|
||||
processorId,
|
||||
@@ -101,17 +107,17 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
|
||||
// Calculate new node array
|
||||
const removedSet = new Set(removedIds);
|
||||
const retainedNodes = this.nodes.filter(n => !removedSet.has(n.id));
|
||||
const retainedNodes = this.nodes.filter((n) => !removedSet.has(n.id));
|
||||
const newGraph = [...retainedNodes];
|
||||
|
||||
// We append the output nodes in the same general position if possible,
|
||||
|
||||
// We append the output nodes in the same general position if possible,
|
||||
// but in a complex graph we just ensure they exist. V2 graph uses timestamps for order.
|
||||
// For simplicity, we just push added nodes to the end of the retained array
|
||||
newGraph.push(...addedNodes);
|
||||
|
||||
// Calculate new provenance map
|
||||
const newProvenanceMap = new Map(this.provenanceMap);
|
||||
|
||||
|
||||
let finalPristineMap = this.pristineNodesMap;
|
||||
|
||||
// Map the new synthetic nodes back to their pristine roots
|
||||
@@ -155,14 +161,16 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
newGraph,
|
||||
finalPristineMap,
|
||||
newProvenanceMap,
|
||||
[...this.history, mutation]
|
||||
[...this.history, mutation],
|
||||
);
|
||||
}
|
||||
|
||||
getPristineNodes(id: string): readonly ConcreteNode[] {
|
||||
const pristineIds = this.provenanceMap.get(id);
|
||||
if (!pristineIds) return [];
|
||||
return Array.from(pristineIds).map(pid => this.pristineNodesMap.get(pid)!);
|
||||
return Array.from(pristineIds).map(
|
||||
(pid) => this.pristineNodesMap.get(pid)!,
|
||||
);
|
||||
}
|
||||
|
||||
getAuditLog(): readonly GraphMutation[] {
|
||||
@@ -171,8 +179,8 @@ export class ContextWorkingBufferImpl implements ContextWorkingBuffer {
|
||||
|
||||
getLineage(id: string): readonly ConcreteNode[] {
|
||||
const lineage: ConcreteNode[] = [];
|
||||
const currentNodesMap = new Map(this.nodes.map(n => [n.id, n]));
|
||||
|
||||
const currentNodesMap = new Map(this.nodes.map((n) => [n.id, n]));
|
||||
|
||||
let current = currentNodesMap.get(id);
|
||||
while (current) {
|
||||
lineage.push(current);
|
||||
|
||||
@@ -43,7 +43,12 @@ class ModifyingProcessor implements ContextProcessor {
|
||||
text: newParts[0].text + ' [modified]',
|
||||
};
|
||||
}
|
||||
newTargets[0] = { ...prompt, id: prompt.id + '-modified', replacesId: prompt.id, semanticParts: newParts };
|
||||
newTargets[0] = {
|
||||
...prompt,
|
||||
id: prompt.id + '-modified',
|
||||
replacesId: prompt.id,
|
||||
semanticParts: newParts,
|
||||
};
|
||||
}
|
||||
return newTargets;
|
||||
}
|
||||
|
||||
@@ -5,10 +5,7 @@
|
||||
*/
|
||||
|
||||
import type { ConcreteNode } from '../ir/types.js';
|
||||
import type {
|
||||
ContextProcessor,
|
||||
ContextWorker,
|
||||
} from '../pipeline.js';
|
||||
import type { ContextProcessor, ContextWorker } from '../pipeline.js';
|
||||
import type { SidecarConfig, PipelineDef, PipelineTrigger } from './types.js';
|
||||
import type {
|
||||
ContextEnvironment,
|
||||
|
||||
@@ -132,7 +132,7 @@ export class SimulationHarness {
|
||||
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;
|
||||
|
||||
@@ -54,9 +54,7 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
],
|
||||
},
|
||||
],
|
||||
workers: [
|
||||
{ workerId: 'StateSnapshotWorker' }
|
||||
]
|
||||
workers: [{ workerId: 'StateSnapshotWorker' }],
|
||||
});
|
||||
|
||||
const mockLlmClient = createMockLlmClient([
|
||||
@@ -146,9 +144,9 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
const generousConfig: SidecarConfig = {
|
||||
budget: { maxTokens: 100000, retainedTokens: 50000 },
|
||||
pipelines: [], // No triggers
|
||||
workers: []
|
||||
workers: [],
|
||||
};
|
||||
|
||||
|
||||
const harness = await SimulationHarness.create(
|
||||
generousConfig,
|
||||
mockLlmClient,
|
||||
@@ -167,7 +165,7 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
]);
|
||||
|
||||
const goldenState = await harness.getGoldenState();
|
||||
|
||||
|
||||
// Total tokens should cleanly match character count with no synthetic nodes
|
||||
expect(goldenState).toMatchSnapshot();
|
||||
});
|
||||
@@ -177,14 +175,11 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
budget: { maxTokens: 200, retainedTokens: 100 },
|
||||
pipelines: [], // No standard pipelines
|
||||
workers: [
|
||||
{ workerId: 'StateSnapshotWorker' } // This should fire on chunk events
|
||||
]
|
||||
{ workerId: 'StateSnapshotWorker' }, // This should fire on chunk events
|
||||
],
|
||||
};
|
||||
|
||||
const harness = await SimulationHarness.create(
|
||||
gcConfig,
|
||||
mockLlmClient,
|
||||
);
|
||||
|
||||
const harness = await SimulationHarness.create(gcConfig, mockLlmClient);
|
||||
|
||||
// Turn 0
|
||||
await harness.simulateTurn([
|
||||
@@ -208,7 +203,7 @@ describe('System Lifecycle Golden Tests', () => {
|
||||
]);
|
||||
|
||||
const goldenState = await harness.getGoldenState();
|
||||
|
||||
|
||||
// We should see ROLLING_SUMMARY nodes injected into the graph, proving the worker ran in the background
|
||||
expect(goldenState).toMatchSnapshot();
|
||||
});
|
||||
|
||||
@@ -23,10 +23,7 @@ import type { Config } from '../../config/config.js';
|
||||
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
|
||||
import type { Content, GenerateContentResponse } from '@google/genai';
|
||||
import { InboxSnapshotImpl } from '../sidecar/inbox.js';
|
||||
import type {
|
||||
InboxMessage,
|
||||
ProcessArgs,
|
||||
} from '../pipeline.js';
|
||||
import type { InboxMessage, ProcessArgs } from '../pipeline.js';
|
||||
|
||||
/**
|
||||
* Creates a valid mock GenerateContentResponse with the provided text.
|
||||
@@ -184,7 +181,9 @@ export function createMockProcessArgs(
|
||||
): ProcessArgs {
|
||||
return {
|
||||
targets,
|
||||
buffer: ContextWorkingBufferImpl.initialize(bufferNodes.length ? bufferNodes : targets),
|
||||
buffer: ContextWorkingBufferImpl.initialize(
|
||||
bufferNodes.length ? bufferNodes : targets,
|
||||
),
|
||||
inbox: new InboxSnapshotImpl(inboxMessages),
|
||||
};
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user