mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
🔏 fix: Address Codex re-review on memory capability (round 8)
- Enforce memory size limits on inline writes: createMemoryTool now rejects keys over 1000 chars and values over memory.charLimit, matching the REST memory routes, so an inline-memory agent can't persist blobs the memory UI/API would reject (packages/api/src/agents/memory.ts, api/app/clients/tools/util/handleTools.js) + test. - Recheck the agents 'memory' endpoint capability at execution time, so a stale/hallucinated set_memory/delete_memory call can't mutate memory after an admin removes the capability while the agent document still carries the marker (api/app/clients/tools/util/handleTools.js).
This commit is contained in:
parent
e5d56a920c
commit
6a62841bab
3 changed files with 43 additions and 1 deletions
|
|
@ -20,6 +20,7 @@ const {
|
|||
Tools,
|
||||
Constants,
|
||||
Permissions,
|
||||
EModelEndpoint,
|
||||
EToolResources,
|
||||
PermissionTypes,
|
||||
} = require('librechat-data-provider');
|
||||
|
|
@ -66,6 +67,13 @@ const { getRoleByName, setMemory, deleteMemory, getFormattedMemories } = require
|
|||
* @returns {Promise<boolean>}
|
||||
*/
|
||||
async function isMemoryToolUsable(req, writePermissions = []) {
|
||||
/** Re-check the agents `memory` capability at execution time: an admin can
|
||||
* disable it after agents were created with the marker, and the raw agent in
|
||||
* the execution context still carries that marker. */
|
||||
const agentsCapabilities = req?.config?.endpoints?.[EModelEndpoint.agents]?.capabilities;
|
||||
if (!Array.isArray(agentsCapabilities) || !agentsCapabilities.includes(Tools.memory)) {
|
||||
return false;
|
||||
}
|
||||
if (!isMemoryEnabled(req?.config?.memory)) {
|
||||
return false;
|
||||
}
|
||||
|
|
@ -419,6 +427,7 @@ const loadTools = async ({
|
|||
} else if (tool === SET_MEMORY_TOOL_NAME) {
|
||||
const memoryConfig = options.req?.config?.memory;
|
||||
const validKeys = memoryConfig?.validKeys;
|
||||
const charLimit = memoryConfig?.charLimit;
|
||||
const tokenLimit = memoryConfig?.tokenLimit;
|
||||
requestedTools[tool] = async () => {
|
||||
if (
|
||||
|
|
@ -439,7 +448,14 @@ const loadTools = async ({
|
|||
return null;
|
||||
}
|
||||
}
|
||||
return createMemoryTool({ userId: user, setMemory, validKeys, tokenLimit, totalTokens });
|
||||
return createMemoryTool({
|
||||
userId: user,
|
||||
setMemory,
|
||||
validKeys,
|
||||
charLimit,
|
||||
tokenLimit,
|
||||
totalTokens,
|
||||
});
|
||||
};
|
||||
continue;
|
||||
} else if (tool === DELETE_MEMORY_TOOL_NAME) {
|
||||
|
|
|
|||
|
|
@ -596,6 +596,15 @@ describe('createMemoryTool tokenLimit enforcement', () => {
|
|||
expect(setMemory).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('rejects values longer than charLimit without writing', async () => {
|
||||
const setMemory = jest.fn().mockResolvedValue({ ok: true });
|
||||
const tool = createMemoryTool({ userId: 'user-1', setMemory, charLimit: 10 });
|
||||
|
||||
await tool.invoke({ key: 'k1', value: 'this value is far longer than ten characters' });
|
||||
|
||||
expect(setMemory).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -58,6 +58,9 @@ export const memoryInstructions =
|
|||
export const SET_MEMORY_TOOL_NAME = 'set_memory';
|
||||
export const DELETE_MEMORY_TOOL_NAME = 'delete_memory';
|
||||
|
||||
/** Maximum memory key length, matching the REST memory routes. */
|
||||
const MEMORY_KEY_CHAR_LIMIT = 1000;
|
||||
|
||||
const SET_MEMORY_DESCRIPTION = 'Saves important information about the user into memory.';
|
||||
const DELETE_MEMORY_DESCRIPTION =
|
||||
'Deletes specific memory data about the user using the provided key. For updating existing memories, use the `set_memory` tool instead';
|
||||
|
|
@ -102,12 +105,14 @@ export const createMemoryTool = ({
|
|||
userId,
|
||||
setMemory,
|
||||
validKeys,
|
||||
charLimit,
|
||||
tokenLimit,
|
||||
totalTokens = 0,
|
||||
}: {
|
||||
userId: string | ObjectId;
|
||||
setMemory: MemoryMethods['setMemory'];
|
||||
validKeys?: string[];
|
||||
charLimit?: number;
|
||||
tokenLimit?: number;
|
||||
totalTokens?: number;
|
||||
}): DynamicStructuredTool => {
|
||||
|
|
@ -135,6 +140,18 @@ export const createMemoryTool = ({
|
|||
return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
|
||||
}
|
||||
|
||||
/** Mirror the REST memory routes' size guards so inline writes can't
|
||||
* persist values the normal memory UI/API would reject. */
|
||||
if (key.length > MEMORY_KEY_CHAR_LIMIT) {
|
||||
return [
|
||||
`Key exceeds maximum length of ${MEMORY_KEY_CHAR_LIMIT} characters.`,
|
||||
undefined,
|
||||
];
|
||||
}
|
||||
if (charLimit && value.length > charLimit) {
|
||||
return [`Value exceeds maximum length of ${charLimit} characters.`, undefined];
|
||||
}
|
||||
|
||||
const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base');
|
||||
/** Total excluding this key's prior in-instance write, so a same-key
|
||||
* rewrite is measured as a replacement rather than an addition. */
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue