feat(core): inject memory and JIT context into subagents (#23032)

Author: abhipatel12

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('<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);
+      });
+    });
   });
 });

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<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,
+        };
 
         while (true) {
           // 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.
     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}`;