mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-13 05:12:55 -07:00
feat(core): inject memory and JIT context into subagents (#23032)
This commit is contained in:
@@ -3373,5 +3373,104 @@ describe('LocalAgentExecutor', () => {
|
|||||||
const uniqueNames = new Set(names);
|
const uniqueNames = new Set(names);
|
||||||
expect(uniqueNames.size).toBe(names.length);
|
expect(uniqueNames.size).toBe(names.length);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('Memory Injection', () => {
|
||||||
|
it('should inject system instruction memory into system prompt', async () => {
|
||||||
|
const definition = createTestDefinition();
|
||||||
|
const executor = await LocalAgentExecutor.create(
|
||||||
|
definition,
|
||||||
|
mockConfig,
|
||||||
|
onActivity,
|
||||||
|
);
|
||||||
|
|
||||||
|
const mockMemory = 'Global memory constraint';
|
||||||
|
vi.spyOn(mockConfig, 'getSystemInstructionMemory').mockReturnValue(
|
||||||
|
mockMemory,
|
||||||
|
);
|
||||||
|
|
||||||
|
mockModelResponse([
|
||||||
|
{
|
||||||
|
name: TASK_COMPLETE_TOOL_NAME,
|
||||||
|
args: { finalResult: 'done' },
|
||||||
|
id: 'call1',
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
await executor.run({ goal: 'test' }, signal);
|
||||||
|
|
||||||
|
const chatConstructorArgs = MockedGeminiChat.mock.calls[0];
|
||||||
|
const systemInstruction = chatConstructorArgs[1] as string;
|
||||||
|
|
||||||
|
expect(systemInstruction).toContain(mockMemory);
|
||||||
|
expect(systemInstruction).toContain('<loaded_context>');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should inject environment memory into the first message when JIT is disabled', async () => {
|
||||||
|
const definition = createTestDefinition();
|
||||||
|
const executor = await LocalAgentExecutor.create(
|
||||||
|
definition,
|
||||||
|
mockConfig,
|
||||||
|
onActivity,
|
||||||
|
);
|
||||||
|
|
||||||
|
const mockMemory = 'Project memory rule';
|
||||||
|
vi.spyOn(mockConfig, 'getEnvironmentMemory').mockReturnValue(
|
||||||
|
mockMemory,
|
||||||
|
);
|
||||||
|
vi.spyOn(mockConfig, 'isJitContextEnabled').mockReturnValue(false);
|
||||||
|
|
||||||
|
mockModelResponse([
|
||||||
|
{
|
||||||
|
name: TASK_COMPLETE_TOOL_NAME,
|
||||||
|
args: { finalResult: 'done' },
|
||||||
|
id: 'call1',
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
await executor.run({ goal: 'test' }, signal);
|
||||||
|
|
||||||
|
const { message } = getMockMessageParams(0);
|
||||||
|
const parts = message as Part[];
|
||||||
|
|
||||||
|
expect(parts).toBeDefined();
|
||||||
|
const memoryPart = parts.find((p) => p.text?.includes(mockMemory));
|
||||||
|
expect(memoryPart).toBeDefined();
|
||||||
|
expect(memoryPart?.text).toBe(mockMemory);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should inject session memory into the first message when JIT is enabled', async () => {
|
||||||
|
const definition = createTestDefinition();
|
||||||
|
const executor = await LocalAgentExecutor.create(
|
||||||
|
definition,
|
||||||
|
mockConfig,
|
||||||
|
onActivity,
|
||||||
|
);
|
||||||
|
|
||||||
|
const mockMemory =
|
||||||
|
'<loaded_context>\nExtension memory rule\n</loaded_context>';
|
||||||
|
vi.spyOn(mockConfig, 'getSessionMemory').mockReturnValue(mockMemory);
|
||||||
|
vi.spyOn(mockConfig, 'isJitContextEnabled').mockReturnValue(true);
|
||||||
|
|
||||||
|
mockModelResponse([
|
||||||
|
{
|
||||||
|
name: TASK_COMPLETE_TOOL_NAME,
|
||||||
|
args: { finalResult: 'done' },
|
||||||
|
id: 'call1',
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
await executor.run({ goal: 'test' }, signal);
|
||||||
|
|
||||||
|
const { message } = getMockMessageParams(0);
|
||||||
|
const parts = message as Part[];
|
||||||
|
|
||||||
|
expect(parts).toBeDefined();
|
||||||
|
const memoryPart = parts.find((p) =>
|
||||||
|
p.text?.includes('Extension memory rule'),
|
||||||
|
);
|
||||||
|
expect(memoryPart).toBeDefined();
|
||||||
|
expect(memoryPart?.text).toContain(mockMemory);
|
||||||
|
});
|
||||||
|
});
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ import { CompressionStatus } from '../core/turn.js';
|
|||||||
import { type ToolCallRequestInfo } from '../scheduler/types.js';
|
import { type ToolCallRequestInfo } from '../scheduler/types.js';
|
||||||
import { ChatCompressionService } from '../services/chatCompressionService.js';
|
import { ChatCompressionService } from '../services/chatCompressionService.js';
|
||||||
import { getDirectoryContextString } from '../utils/environmentContext.js';
|
import { getDirectoryContextString } from '../utils/environmentContext.js';
|
||||||
|
import { renderUserMemory } from '../prompts/snippets.js';
|
||||||
import { promptIdContext } from '../utils/promptIdContext.js';
|
import { promptIdContext } from '../utils/promptIdContext.js';
|
||||||
import {
|
import {
|
||||||
logAgentStart,
|
logAgentStart,
|
||||||
@@ -585,12 +586,24 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
|||||||
);
|
);
|
||||||
const formattedInitialHints = formatUserHintsForModel(initialHints);
|
const formattedInitialHints = formatUserHintsForModel(initialHints);
|
||||||
|
|
||||||
let currentMessage: Content = formattedInitialHints
|
// Inject loaded memory files (JIT + extension/project memory)
|
||||||
? {
|
const environmentMemory = this.context.config.isJitContextEnabled?.()
|
||||||
role: 'user',
|
? this.context.config.getSessionMemory()
|
||||||
parts: [{ text: formattedInitialHints }, { text: query }],
|
: this.context.config.getEnvironmentMemory();
|
||||||
|
|
||||||
|
const initialParts: Part[] = [];
|
||||||
|
if (environmentMemory) {
|
||||||
|
initialParts.push({ text: environmentMemory });
|
||||||
}
|
}
|
||||||
: { role: 'user', parts: [{ text: query }] };
|
if (formattedInitialHints) {
|
||||||
|
initialParts.push({ text: formattedInitialHints });
|
||||||
|
}
|
||||||
|
initialParts.push({ text: query });
|
||||||
|
|
||||||
|
let currentMessage: Content = {
|
||||||
|
role: 'user',
|
||||||
|
parts: initialParts,
|
||||||
|
};
|
||||||
|
|
||||||
while (true) {
|
while (true) {
|
||||||
// Check for termination conditions like max turns.
|
// Check for termination conditions like max turns.
|
||||||
@@ -1375,6 +1388,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
|||||||
// Inject user inputs into the prompt template.
|
// Inject user inputs into the prompt template.
|
||||||
let finalPrompt = templateString(promptConfig.systemPrompt, inputs);
|
let finalPrompt = templateString(promptConfig.systemPrompt, inputs);
|
||||||
|
|
||||||
|
// Append memory context if available.
|
||||||
|
const systemMemory = this.context.config.getSystemInstructionMemory();
|
||||||
|
if (systemMemory) {
|
||||||
|
finalPrompt += `\n\n${renderUserMemory(systemMemory)}`;
|
||||||
|
}
|
||||||
|
|
||||||
// Append environment context (CWD and folder structure).
|
// Append environment context (CWD and folder structure).
|
||||||
const dirContext = await getDirectoryContextString(this.context.config);
|
const dirContext = await getDirectoryContextString(this.context.config);
|
||||||
finalPrompt += `\n\n# Environment Context\n${dirContext}`;
|
finalPrompt += `\n\n# Environment Context\n${dirContext}`;
|
||||||
|
|||||||
Reference in New Issue
Block a user