From 9dfcd837eea0c697eb31fb97aab053711dccd991 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 20 Jun 2026 16:48:31 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20fix:=20Address=20Codex=20review?= =?UTF-8?q?=20on=20memory=20capability=20(gating,=20validKeys,=20usage=20g?= =?UTF-8?q?uard)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Strip the memory capability from the served agents capabilities when memory is not configured/enabled, so the badge, tools dropdown, agent-builder toggle, and backend capability gate stay consistent instead of exposing an inert toggle on default installs (where MEMORIES.USE defaults true). - Surface configured memory.validKeys in the inline tool definitions so the model is told the allowed keys up front, matching the runtime createMemoryTool schema. - Append a strict explicit-request usage guard to the agent instructions when inline memory tools are registered, preserving the memory-agent's privacy behavior. - Add AppService tests covering memory-capability stripping. --- packages/api/src/agents/initialize.ts | 9 ++- packages/api/src/agents/memory.ts | 73 +++++++++++++------ packages/data-schemas/src/app/service.spec.ts | 31 +++++++- packages/data-schemas/src/app/service.ts | 13 +++- 4 files changed, 98 insertions(+), 28 deletions(-) diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 305341daf2..cc7c44e26d 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -50,10 +50,10 @@ import { registerFileAuthoringTools, isFileAuthoringToolDefinition, } from './tools'; +import { registerMemoryTools, memoryToolUsageGuard } from './memory'; import { filterFilesByEndpointConfig } from '~/files'; import { generateArtifactsPrompt } from '~/prompts'; import { getProviderConfig } from '~/endpoints'; -import { registerMemoryTools } from './memory'; import { primeResources } from './resources'; /** @@ -1059,8 +1059,13 @@ export async function initializeAgent( */ const agentRequestsMemory = (agent.tools ?? []).includes(Tools.memory); if (params.memoryAvailable === true && agentRequestsMemory) { - const memoryResult = registerMemoryTools({ toolRegistry, toolDefinitions }); + const memoryResult = registerMemoryTools({ + toolRegistry, + toolDefinitions, + validKeys: req.config?.memory?.validKeys, + }); toolDefinitions = memoryResult.toolDefinitions; + appendAdditionalInstructions(agent, memoryToolUsageGuard); } else if (agentRequestsMemory) { logger.debug( `[initializeAgent] Agent "${agent.id}" requests memory but memoryAvailable=${String(params.memoryAvailable)}; skipping set_memory + delete_memory registration.`, diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index d148e45581..e33e9b3e79 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -267,42 +267,64 @@ export const createDeleteMemoryTool = ({ }, ); }; +/** + * Strict usage guard appended to the agent's instructions when the inline + * memory tools are registered, preserving the memory-agent's explicit-request + * behavior so the model never stores facts it merely observed. + */ +export const memoryToolUsageGuard = `Only use the \`set_memory\` and \`delete_memory\` tools when the user explicitly asks you to remember, update, or forget something (e.g. "remember that...", "don't forget...", "forget..."). Never store information merely because the user mentioned it in conversation.`; + /** * LLM-facing definitions for the inline memory tool pair, used by the * event-driven (definitions-only) loader. The `memory` capability string on * an agent's `tools` array expands into this pair at initialize time via * {@link registerMemoryTools}; the runtime instances created in the tool * service enforce `validKeys`/`tokenLimit` and emit memory artifacts. + * `validKeys` is surfaced in the key descriptions so the model is told the + * allowed keys up front, matching the runtime `createMemoryTool` schema. */ -export const memoryToolDefinitions: LCTool[] = [ - { - name: SET_MEMORY_TOOL_NAME, - description: SET_MEMORY_DESCRIPTION, - parameters: { - type: 'object', - properties: { - key: { type: 'string', description: 'The key identifier for this memory' }, - value: { - type: 'string', - description: - 'Value MUST be a complete sentence that fully describes relevant user information.', +export function getMemoryToolDefinitions(validKeys?: string[]): LCTool[] { + const hasValidKeys = Array.isArray(validKeys) && validKeys.length > 0; + return [ + { + name: SET_MEMORY_TOOL_NAME, + description: SET_MEMORY_DESCRIPTION, + parameters: { + type: 'object', + properties: { + key: { + type: 'string', + description: hasValidKeys + ? `The key of the memory value. Must be one of: ${validKeys!.join(', ')}` + : 'The key identifier for this memory', + }, + value: { + type: 'string', + description: + 'Value MUST be a complete sentence that fully describes relevant user information.', + }, }, + required: ['key', 'value'], }, - required: ['key', 'value'], }, - }, - { - name: DELETE_MEMORY_TOOL_NAME, - description: DELETE_MEMORY_DESCRIPTION, - parameters: { - type: 'object', - properties: { - key: { type: 'string', description: 'The key identifier of the memory to delete' }, + { + name: DELETE_MEMORY_TOOL_NAME, + description: DELETE_MEMORY_DESCRIPTION, + parameters: { + type: 'object', + properties: { + key: { + type: 'string', + description: hasValidKeys + ? `The key of the memory to delete. Must be one of: ${validKeys!.join(', ')}` + : 'The key identifier of the memory to delete', + }, + }, + required: ['key'], }, - required: ['key'], }, - }, -] as LCTool[]; + ] as LCTool[]; +} /** * Idempotently registers the inline memory tool pair (`set_memory` + @@ -314,10 +336,13 @@ export const memoryToolDefinitions: LCTool[] = [ export function registerMemoryTools({ toolRegistry, toolDefinitions, + validKeys, }: { toolRegistry?: LCToolRegistry; toolDefinitions?: LCTool[]; + validKeys?: string[]; }): { toolDefinitions: LCTool[]; registered: string[] } { + const memoryToolDefinitions = getMemoryToolDefinitions(validKeys); const inputDefinitions = toolDefinitions ?? []; const newDefs: LCTool[] = []; const registered: string[] = []; diff --git a/packages/data-schemas/src/app/service.spec.ts b/packages/data-schemas/src/app/service.spec.ts index 2a2a8ffd1d..1b8928d8ba 100644 --- a/packages/data-schemas/src/app/service.spec.ts +++ b/packages/data-schemas/src/app/service.spec.ts @@ -1,5 +1,9 @@ +import { + EModelEndpoint, + AgentCapabilities, + defaultAssistantsVersion, +} from 'librechat-data-provider'; import type { DeepPartial, TCustomConfig } from 'librechat-data-provider'; -import { EModelEndpoint, defaultAssistantsVersion } from 'librechat-data-provider'; import { AppService, loadSummarizationConfig } from './service'; import logger from '~/config/winston'; @@ -146,3 +150,28 @@ describe('AppService assistants config', () => { ); }); }); + +describe('AppService memory capability', () => { + it('strips the memory capability when no memory config is present', async () => { + const result = await AppService({ config: {} as DeepPartial }); + expect(result.endpoints?.[EModelEndpoint.agents]?.capabilities).not.toContain( + AgentCapabilities.memory, + ); + }); + + it('keeps the memory capability when memory is configured and enabled', async () => { + const config = { memory: { tokenLimit: 10000 } } as DeepPartial; + const result = await AppService({ config }); + expect(result.endpoints?.[EModelEndpoint.agents]?.capabilities).toContain( + AgentCapabilities.memory, + ); + }); + + it('strips the memory capability when memory is explicitly disabled', async () => { + const config = { memory: { disabled: true } } as DeepPartial; + const result = await AppService({ config }); + expect(result.endpoints?.[EModelEndpoint.agents]?.capabilities).not.toContain( + AgentCapabilities.memory, + ); + }); +}); diff --git a/packages/data-schemas/src/app/service.ts b/packages/data-schemas/src/app/service.ts index fe3ae1b007..e9cb4e1d03 100644 --- a/packages/data-schemas/src/app/service.ts +++ b/packages/data-schemas/src/app/service.ts @@ -1,17 +1,18 @@ import { EModelEndpoint, getConfigDefaults, + AgentCapabilities, skillSyncConfigSchema, summarizationConfigSchema, } from 'librechat-data-provider'; import type { TCustomConfig, FileSources, DeepPartial } from 'librechat-data-provider'; import type { AppConfig, FunctionTool } from '~/types/app'; +import { loadMemoryConfig, isMemoryEnabled } from './memory'; import { loadDefaultInterface } from './interface'; import { loadTurnstileConfig } from './turnstile'; import { agentsConfigSetup } from './agents'; import { loadWebSearchConfig } from './web'; import { processModelSpecs } from './specs'; -import { loadMemoryConfig } from './memory'; import { loadEndpoints } from './endpoints'; import { loadOCRConfig } from './ocr'; import logger from '~/config/winston'; @@ -158,6 +159,16 @@ export const AppService = async (params?: { const agentsDefaults = agentsConfigSetup(config); + /** The `memory` capability only functions when memory is configured and + * enabled. Drop it from the served capability set otherwise so the agent + * builder toggle, ephemeral badge, and backend capability gate stay + * consistent instead of exposing an inert memory toggle. */ + if (!isMemoryEnabled(memory) && Array.isArray(agentsDefaults.capabilities)) { + agentsDefaults.capabilities = agentsDefaults.capabilities.filter( + (capability) => capability !== AgentCapabilities.memory, + ); + } + if (!Object.keys(config).length) { const appConfig = { ...defaultConfig,