Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 99 additions & 0 deletions packages/core/src/agents/local-executor.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3203,5 +3203,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('<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);
});
});
});
});
31 changes: 25 additions & 6 deletions packages/core/src/agents/local-executor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,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,
Expand Down Expand Up @@ -576,12 +577,24 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
);
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,
};
Comment thread
abhipatel12 marked this conversation as resolved.

while (true) {
// Check for termination conditions like max turns.
Expand Down Expand Up @@ -1326,6 +1339,12 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
// 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)}`;
}
Comment thread
abhipatel12 marked this conversation as resolved.
Comment thread
abhipatel12 marked this conversation as resolved.

// Append environment context (CWD and folder structure).
const dirContext = await getDirectoryContextString(this.context.config);
finalPrompt += `\n\n# Environment Context\n${dirContext}`;
Expand Down
Loading