Exclude extension context from skill extraction agent (#26879)

This commit is contained in:
Sandy Tao
2026-05-12 10:45:19 -07:00
committed by GitHub
parent bc730b2c0f
commit ebe15553a9
7 changed files with 183 additions and 10 deletions
@@ -4199,6 +4199,49 @@ describe('LocalAgentExecutor', () => {
expect(memoryPart).toBeDefined();
expect(memoryPart?.text).toContain(mockMemory);
});
it('should omit extension context from session memory when disabled by the agent', async () => {
const definition = createTestDefinition();
definition.includeExtensionContext = false;
const executor = await LocalAgentExecutor.create(
definition,
mockConfig,
onActivity,
);
const getSessionMemorySpy = vi
.spyOn(mockConfig, 'getSessionMemory')
.mockImplementation(
(options?: { includeExtensionContext?: boolean }) =>
options?.includeExtensionContext === false
? '<loaded_context>\n<project_context>\nProject memory rule\n</project_context>\n</loaded_context>'
: '<loaded_context>\n<extension_context>\nExtension memory rule\n</extension_context>\n<project_context>\nProject memory rule\n</project_context>\n</loaded_context>',
);
vi.spyOn(mockConfig, 'isJitContextEnabled').mockReturnValue(true);
mockModelResponse([
{
name: COMPLETE_TASK_TOOL_NAME,
args: { finalResult: 'done' },
id: 'call1',
},
]);
await executor.run({ goal: 'test' }, signal);
expect(getSessionMemorySpy).toHaveBeenCalledWith({
includeExtensionContext: false,
});
const { message } = getMockMessageParams(0);
const parts = message as Part[];
const memoryPart = parts.find((p) =>
p.text?.includes('<loaded_context>'),
);
expect(memoryPart?.text).toContain('Project memory rule');
expect(memoryPart?.text).not.toContain('<extension_context>');
expect(memoryPart?.text).not.toContain('Extension memory rule');
});
});
});
});
+13 -4
View File
@@ -640,10 +640,19 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
);
const formattedInitialHints = formatUserHintsForModel(initialHints);
// Inject loaded memory files (JIT + extension/project memory)
const environmentMemory = this.context.config.isJitContextEnabled?.()
? this.context.config.getSessionMemory()
: this.context.config.getEnvironmentMemory();
// Inject loaded memory files. Some background agents opt out of
// extension memory while still retaining project session context.
let environmentMemory: string;
if (this.context.config.isJitContextEnabled?.()) {
environmentMemory =
this.definition.includeExtensionContext === false
? this.context.config.getSessionMemory({
includeExtensionContext: false,
})
: this.context.config.getSessionMemory();
} else {
environmentMemory = this.context.config.getEnvironmentMemory();
}
const initialParts: Part[] = [];
if (environmentMemory) {
@@ -37,6 +37,7 @@ describe('SkillExtractionAgent', () => {
expect(agent.modelConfig.model).toBe(PREVIEW_GEMINI_FLASH_MODEL);
expect(agent.memoryInboxAccess).toBe(true);
expect(agent.autoMemoryExtractionWriteAccess).toBe(true);
expect(agent.includeExtensionContext).toBe(false);
expect(agent.toolConfig?.tools).toEqual(
expect.arrayContaining([
READ_FILE_TOOL_NAME,
@@ -415,6 +415,7 @@ export const SkillExtractionAgent = (
},
memoryInboxAccess: true,
autoMemoryExtractionWriteAccess: true,
includeExtensionContext: false,
toolConfig: {
tools: [
ACTIVATE_SKILL_TOOL_NAME,
+6
View File
@@ -244,6 +244,12 @@ export interface LocalAgentDefinition<
*/
autoMemoryExtractionWriteAccess?: boolean;
/**
* Controls whether extension memory is injected into this agent's initial
* session context when JIT context is enabled. Defaults to true.
*/
includeExtensionContext?: boolean;
/**
* Optional inline MCP servers for this agent.
*/