diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index e0a21e01e3..65f3b76877 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -3373,5 +3373,104 @@ describe('LocalAgentExecutor', () => { const uniqueNames = new Set(names); 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(''); + }); + + 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 = + '\nExtension memory rule\n'; + 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); + }); + }); }); }); diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 9dc92d1321..a860e1e597 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -32,6 +32,7 @@ import { CompressionStatus } from '../core/turn.js'; import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { ChatCompressionService } from '../services/chatCompressionService.js'; import { getDirectoryContextString } from '../utils/environmentContext.js'; +import { renderUserMemory } from '../prompts/snippets.js'; import { promptIdContext } from '../utils/promptIdContext.js'; import { logAgentStart, @@ -585,12 +586,24 @@ export class LocalAgentExecutor { ); const formattedInitialHints = formatUserHintsForModel(initialHints); - let currentMessage: Content = formattedInitialHints - ? { - role: 'user', - parts: [{ text: formattedInitialHints }, { text: query }], - } - : { role: 'user', parts: [{ text: query }] }; + // Inject loaded memory files (JIT + extension/project memory) + const environmentMemory = this.context.config.isJitContextEnabled?.() + ? this.context.config.getSessionMemory() + : this.context.config.getEnvironmentMemory(); + + const initialParts: Part[] = []; + if (environmentMemory) { + initialParts.push({ text: environmentMemory }); + } + if (formattedInitialHints) { + initialParts.push({ text: formattedInitialHints }); + } + initialParts.push({ text: query }); + + let currentMessage: Content = { + role: 'user', + parts: initialParts, + }; while (true) { // Check for termination conditions like max turns. @@ -1375,6 +1388,12 @@ export class LocalAgentExecutor { // Inject user inputs into the prompt template. 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). const dirContext = await getDirectoryContextString(this.context.config); finalPrompt += `\n\n# Environment Context\n${dirContext}`;