From 24311a11d48fceedc0da0d3e21a5e0d22706c8fa Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 20 Jun 2026 23:57:07 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20fix:=20Address=20Codex=20re-revi?= =?UTF-8?q?ew=20on=20memory=20capability=20(round=206)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix a P1 regression from the prior round: the execution-context agent keeps the raw 'memory' capability marker (not the expanded set_memory/delete_memory names), so the opt-in check now matches the marker. This restores memory writes/deletes AND avoids hijacking an MCP tool that merely shares the set_memory/delete_memory name (api/app/clients/tools/util/handleTools.js). - Count repeated set_memory writes to the same key as replacements, not additions, against tokenLimit — set_memory upserts, so a same-key rewrite swaps its prior token contribution instead of double-counting (packages/api/src/agents/memory.ts) + test. - Gate the memory badge, tools dropdown, and agent-builder toggle on the full memory write permissions (USE+CREATE+UPDATE) via a shared useHasMemoryAccess hook, so a read-only-memory role no longer sees an enabled Memory control the backend would refuse to wire up. --- api/app/clients/tools/util/handleTools.js | 25 +++++++++++-------- client/src/components/Chat/Input/Memory.tsx | 9 +++---- .../components/Chat/Input/ToolsDropdown.tsx | 13 ++++++---- .../SidePanel/Agents/AgentConfig.tsx | 7 ++---- client/src/hooks/Roles/index.ts | 1 + client/src/hooks/Roles/useHasMemoryAccess.ts | 25 +++++++++++++++++++ packages/api/src/agents/memory.spec.ts | 18 +++++++++++++ packages/api/src/agents/memory.ts | 21 +++++++++++----- 8 files changed, 86 insertions(+), 33 deletions(-) create mode 100644 client/src/hooks/Roles/useHasMemoryAccess.ts diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index f9f0d92f18..4383251ebf 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -86,20 +86,23 @@ async function isMemoryToolUsable(req, writePermissions = []) { } /** - * Whether the agent actually declared the requested inline memory tool. The - * event-driven executor loads tools by requested name, so a hallucinated or - * undeclared `set_memory`/`delete_memory` call (e.g. on an agent that never - * opted into the memory capability) must not be constructed. - * @param {{ toolDefinitions?: Array<{ name?: string }>, tools?: Array }} [agent] - * @param {string} toolName + * Whether the agent opted into the LibreChat inline memory capability, checked + * via the `memory` capability marker on the agent's tools. The raw agent kept + * in the tool-execution context retains that marker (not the expanded + * `set_memory`/`delete_memory` names), so this: + * - refuses a hallucinated/undeclared memory call on a non-memory agent, and + * - does not let an MCP tool that merely shares the `set_memory`/`delete_memory` + * name get replaced by the built-in memory mutator. + * @param {{ tools?: Array }} [agent] * @returns {boolean} */ -function agentDeclaredMemoryTool(agent, toolName) { +function agentOptedIntoMemory(agent) { if (!agent) { return false; } - const declared = (entry) => (typeof entry === 'string' ? entry : entry?.name) === toolName; - return (agent.toolDefinitions ?? []).some(declared) || (agent.tools ?? []).some(declared); + return (agent.tools ?? []).some( + (entry) => (typeof entry === 'string' ? entry : entry?.name) === Tools.memory, + ); } /** @@ -415,7 +418,7 @@ const loadTools = async ({ const tokenLimit = memoryConfig?.tokenLimit; requestedTools[tool] = async () => { if ( - !agentDeclaredMemoryTool(agent, SET_MEMORY_TOOL_NAME) || + !agentOptedIntoMemory(agent) || !(await isMemoryToolUsable(options.req, [Permissions.CREATE, Permissions.UPDATE])) ) { return null; @@ -439,7 +442,7 @@ const loadTools = async ({ const memoryConfig = options.req?.config?.memory; requestedTools[tool] = async () => { if ( - !agentDeclaredMemoryTool(agent, DELETE_MEMORY_TOOL_NAME) || + !agentOptedIntoMemory(agent) || !(await isMemoryToolUsable(options.req, [Permissions.UPDATE])) ) { return null; diff --git a/client/src/components/Chat/Input/Memory.tsx b/client/src/components/Chat/Input/Memory.tsx index c66d955ce7..d43610599b 100644 --- a/client/src/components/Chat/Input/Memory.tsx +++ b/client/src/components/Chat/Input/Memory.tsx @@ -1,8 +1,8 @@ import React, { memo } from 'react'; import { Brain } from 'lucide-react'; import { CheckboxButton } from '@librechat/client'; -import { Permissions, PermissionTypes, defaultAgentCapabilities } from 'librechat-data-provider'; -import { useLocalize, useHasAccess, useAgentCapabilities, useAuthContext } from '~/hooks'; +import { defaultAgentCapabilities } from 'librechat-data-provider'; +import { useLocalize, useHasMemoryAccess, useAgentCapabilities, useAuthContext } from '~/hooks'; import { useBadgeRowContext } from '~/Providers'; function Memory() { @@ -11,10 +11,7 @@ function Memory() { const context = useBadgeRowContext(); const { toggleState: memoryActive, debouncedChange, isPinned } = context?.memory ?? {}; - const canUseMemory = useHasAccess({ - permissionType: PermissionTypes.MEMORIES, - permission: Permissions.USE, - }); + const canUseMemory = useHasMemoryAccess(); const { memoryEnabled } = useAgentCapabilities( context?.agentsConfig?.capabilities ?? defaultAgentCapabilities, diff --git a/client/src/components/Chat/Input/ToolsDropdown.tsx b/client/src/components/Chat/Input/ToolsDropdown.tsx index ffe24e1d4a..63bbdd7816 100644 --- a/client/src/components/Chat/Input/ToolsDropdown.tsx +++ b/client/src/components/Chat/Input/ToolsDropdown.tsx @@ -10,7 +10,13 @@ import { defaultAgentCapabilities, } from 'librechat-data-provider'; import type { MenuItemProps } from '~/common'; -import { useLocalize, useHasAccess, useAgentCapabilities, useAuthContext } from '~/hooks'; +import { + useLocalize, + useHasAccess, + useAuthContext, + useHasMemoryAccess, + useAgentCapabilities, +} from '~/hooks'; import ArtifactsSubMenu from '~/components/Chat/Input/ArtifactsSubMenu'; import MCPSubMenu from '~/components/Chat/Input/MCPSubMenu'; import { useGetStartupConfig } from '~/data-provider'; @@ -61,10 +67,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { permission: Permissions.USE, }); - const canUseMemory = useHasAccess({ - permissionType: PermissionTypes.MEMORIES, - permission: Permissions.USE, - }); + const canUseMemory = useHasMemoryAccess(); const showMemory = canUseMemory && memoryEnabled && user?.personalization?.memories !== false; const [isPopoverActive, setIsPopoverActive] = useState(false); diff --git a/client/src/components/SidePanel/Agents/AgentConfig.tsx b/client/src/components/SidePanel/Agents/AgentConfig.tsx index 92eb010e0b..ed95912c7f 100644 --- a/client/src/components/SidePanel/Agents/AgentConfig.tsx +++ b/client/src/components/SidePanel/Agents/AgentConfig.tsx @@ -20,11 +20,11 @@ import { getIconKey, cn, } from '~/utils'; +import { useLocalize, useVisibleTools, useHasAccess, useHasMemoryAccess } from '~/hooks'; import { ToolSelectDialog, MCPToolSelectDialog } from '~/components/Tools'; import useAgentCapabilities from '~/hooks/Agents/useAgentCapabilities'; import { useListSkillsQuery, useGetAgentFiles } from '~/data-provider'; import { useFileMapContext, useAgentPanelContext } from '~/Providers'; -import { useLocalize, useVisibleTools, useHasAccess } from '~/hooks'; import { SkillSelectDialog } from '~/components/Skills/dialogs'; import AgentCategorySelector from './AgentCategorySelector'; import Action from '~/components/SidePanel/Builder/Action'; @@ -114,10 +114,7 @@ export default function AgentConfig() { permissionType: PermissionTypes.SKILLS, permission: Permissions.USE, }); - const hasMemoryAccess = useHasAccess({ - permissionType: PermissionTypes.MEMORIES, - permission: Permissions.USE, - }); + const hasMemoryAccess = useHasMemoryAccess(); const showMemory = hasMemoryAccess && memoryEnabled; const showSkills = hasSkillsAccess && skillsEnabled; const { data: skillsData } = useListSkillsQuery({ limit: 100 }, { enabled: showSkills }); diff --git a/client/src/hooks/Roles/index.ts b/client/src/hooks/Roles/index.ts index 998cfba79b..d557c95796 100644 --- a/client/src/hooks/Roles/index.ts +++ b/client/src/hooks/Roles/index.ts @@ -1 +1,2 @@ export { default as useHasAccess } from './useHasAccess'; +export { default as useHasMemoryAccess } from './useHasMemoryAccess'; diff --git a/client/src/hooks/Roles/useHasMemoryAccess.ts b/client/src/hooks/Roles/useHasMemoryAccess.ts new file mode 100644 index 0000000000..7985536814 --- /dev/null +++ b/client/src/hooks/Roles/useHasMemoryAccess.ts @@ -0,0 +1,25 @@ +import { Permissions, PermissionTypes } from 'librechat-data-provider'; +import useHasAccess from './useHasAccess'; + +/** + * The inline memory tools (`set_memory`/`delete_memory`) mutate memory, so the + * memory badge and agent-builder toggle gate on the full write permission set + * (USE + CREATE + UPDATE) — matching the backend `memoryAvailable` gate and the + * runtime tool loader. A read-only-memory role therefore never sees an enabled + * Memory control that the backend would refuse to wire up. + */ +export default function useHasMemoryAccess(): boolean { + const canUse = useHasAccess({ + permissionType: PermissionTypes.MEMORIES, + permission: Permissions.USE, + }); + const canCreate = useHasAccess({ + permissionType: PermissionTypes.MEMORIES, + permission: Permissions.CREATE, + }); + const canUpdate = useHasAccess({ + permissionType: PermissionTypes.MEMORIES, + permission: Permissions.UPDATE, + }); + return canUse && canCreate && canUpdate; +} diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 137938be9a..973a3df414 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/packages/api/src/agents/memory.spec.ts @@ -595,4 +595,22 @@ describe('createMemoryTool tokenLimit enforcement', () => { expect(setMemory).toHaveBeenCalledTimes(2); }); + + it('treats a repeat write to the same key as a replacement, not an addition', async () => { + const setMemory = jest.fn().mockResolvedValue({ ok: true }); + /** ~100 tokens; two distinct keys would exceed the 150 limit, but rewriting + * the same key only replaces its value and must stay within the cap. */ + const value = 'word '.repeat(100).trim(); + const tool = createMemoryTool({ + userId: 'user-1', + setMemory, + tokenLimit: 150, + totalTokens: 0, + }); + + await tool.invoke({ key: 'k1', value }); + await tool.invoke({ key: 'k1', value }); + + expect(setMemory).toHaveBeenCalledTimes(2); + }); }); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index 4850a62d03..64e69c39ed 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -117,6 +117,10 @@ export const createMemoryTool = ({ * check against the same stale total and collectively exceed `tokenLimit`. */ let currentTotalTokens = totalTokens; let writeChain: Promise = Promise.resolve(); + /** Tokens this instance has already committed per key. `set_memory` upserts, + * so a repeat write to the same key REPLACES its value — the running total + * must swap the prior contribution for the new one, not add both. */ + const writtenTokensByKey = new Map(); return tool( async ({ key, value }) => { @@ -132,7 +136,10 @@ export const createMemoryTool = ({ } const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base'); - const remainingTokens = tokenLimit ? tokenLimit - currentTotalTokens : Infinity; + /** Total excluding this key's prior in-instance write, so a same-key + * rewrite is measured as a replacement rather than an addition. */ + const baseTotalTokens = currentTotalTokens - (writtenTokensByKey.get(key) ?? 0); + const remainingTokens = tokenLimit ? tokenLimit - baseTotalTokens : Infinity; if (tokenLimit && remainingTokens <= 0) { const errorArtifact: MemoryArtifactRecord = { @@ -142,17 +149,18 @@ export const createMemoryTool = ({ value: JSON.stringify({ errorType: 'already_exceeded', tokenCount: Math.abs(remainingTokens), - totalTokens: currentTotalTokens, + totalTokens: baseTotalTokens, tokenLimit: tokenLimit!, }), - tokenCount: currentTotalTokens, + tokenCount: baseTotalTokens, }, }; return [`Memory storage exceeded. Cannot save new memories.`, errorArtifact]; } + const newTotalTokens = baseTotalTokens + tokenCount; + if (tokenLimit) { - const newTotalTokens = currentTotalTokens + tokenCount; const newRemainingTokens = tokenLimit - newTotalTokens; if (newRemainingTokens < 0) { @@ -166,7 +174,7 @@ export const createMemoryTool = ({ totalTokens: newTotalTokens, tokenLimit, }), - tokenCount: currentTotalTokens, + tokenCount: baseTotalTokens, }, }; return [`Memory storage would exceed limit. Cannot save this memory.`, errorArtifact]; @@ -185,7 +193,8 @@ export const createMemoryTool = ({ const result = await setMemory({ userId, key, value, tokenCount }); if (result.ok) { if (tokenLimit) { - currentTotalTokens += tokenCount; + currentTotalTokens = newTotalTokens; + writtenTokensByKey.set(key, tokenCount); } logger.debug(`Memory set for key "${key}" (${tokenCount} tokens) for user "${userId}"`); return [`Memory set for key "${key}" (${tokenCount} tokens)`, artifact];