diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 7cb265464e..299c2e6ac6 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -46,6 +46,7 @@ const { attachAskUserQuestionArgs, createContentIndexOffsetHandlers, getRequestMemories, + getMemoryAgentId, createMemoryProcessor, agentHasInlineMemoryTools, loadAgent: loadAgentFn, @@ -536,10 +537,34 @@ class AgentClient extends BaseClient { * keys + token metadata) is reserved for agents that can call * `delete_memory`; everyone else gets the unkeyed values only. */ const memories = await this.useMemory(); + /** Partition the loaded memories belong to (the primary agent's). */ + const loadedMemoryAgentId = getMemoryAgentId(this.options.agent); const buildMemoryContext = (text) => text ? `${memoryInstructions}\n\n# Existing memory about the user:\n${text}` : undefined; - const memoryContext = buildMemoryContext(memories?.withoutKeys); - const keyedMemoryContext = buildMemoryContext(memories?.withKeys); + /** Resolves formatted memories for an agent's own partition. A defined + * `memories` means the run-level gates (permission, opt-out, config) + * passed; agents on other partitions fetch through the request-scoped + * cache so repeated partitions share one query. */ + const getAgentPartitionMemories = async (agent) => { + if (!memories) { + return undefined; + } + const agentPartition = getMemoryAgentId(agent); + if (agentPartition === loadedMemoryAgentId) { + return memories; + } + try { + return await getRequestMemories({ + req: this.options.req, + userId: this.options.req.user.id + '', + agentId: agentPartition, + getFormattedMemories: db.getFormattedMemories, + }); + } catch (error) { + logger.error('[AgentClient] Error loading partition memories', error); + return undefined; + } + }; const sharedRunContext = sharedRunContextParts.join('\n\n'); const memoryAgentEnabled = isMemoryAgentEnabled(this.options.req.config?.memory); @@ -588,15 +613,17 @@ class AgentClient extends BaseClient { const configServers = await resolveConfigServers(this.options.req); await Promise.all( - allAgents.map(({ agent, agentId }) => { + allAgents.map(async ({ agent, agentId }) => { const agentRunContextParts = [sharedRunContext]; const agentHasMemory = agentHasInlineMemoryTools(agent); - const agentMemoryContext = agentHasMemory ? keyedMemoryContext : memoryContext; - if ( - agentMemoryContext && - (agentId === this.options.agent.id || memoryAgentEnabled || agentHasMemory) - ) { - agentRunContextParts.push(agentMemoryContext); + if (agentId === this.options.agent.id || memoryAgentEnabled || agentHasMemory) { + const partitionMemories = await getAgentPartitionMemories(agent); + const agentMemoryContext = buildMemoryContext( + agentHasMemory ? partitionMemories?.withKeys : partitionMemories?.withoutKeys, + ); + if (agentMemoryContext) { + agentRunContextParts.push(agentMemoryContext); + } } const scopedContext = agentScopedContext.get(agentId); if (scopedContext) { @@ -674,6 +701,8 @@ class AgentClient extends BaseClient { } const userId = this.options.req.user.id + ''; + /** Memory partition of the primary agent; undefined = shared personal pool */ + const memoryAgentId = getMemoryAgentId(this.options.agent); this.processMemory = undefined; if (!isMemoryAgentEnabled(memoryConfig)) { @@ -681,6 +710,7 @@ class AgentClient extends BaseClient { const { withKeys, withoutKeys } = await getRequestMemories({ req: this.options.req, userId, + agentId: memoryAgentId, getFormattedMemories: db.getFormattedMemories, }); return { withKeys, withoutKeys }; @@ -786,6 +816,7 @@ class AgentClient extends BaseClient { const streamId = this.options.req?._resumableStreamId || null; const [withoutKeys, processMemory] = await createMemoryProcessor({ userId, + agentId: memoryAgentId, config, messageId, streamId, @@ -805,6 +836,7 @@ class AgentClient extends BaseClient { ({ withKeys } = await getRequestMemories({ req: this.options.req, userId, + agentId: memoryAgentId, getFormattedMemories: db.getFormattedMemories, })); } catch (error) { diff --git a/api/server/routes/memories.js b/api/server/routes/memories.js index e71e94f457..bfd15640a0 100644 --- a/api/server/routes/memories.js +++ b/api/server/routes/memories.js @@ -1,13 +1,21 @@ const express = require('express'); const { Tokenizer, generateCheckAccess } = require('@librechat/api'); -const { PermissionTypes, Permissions } = require('librechat-data-provider'); +const { + PermissionTypes, + PermissionBits, + ResourceType, + Permissions, +} = require('librechat-data-provider'); +const { findAccessibleResources } = require('~/server/services/PermissionService'); const { getAllUserMemories, + getUserMemories, toggleUserMemories, getRoleByName, createMemory, deleteMemory, setMemory, + getAgents, } = require('~/models'); const { requireJwtAuth, configMiddleware } = require('~/server/middleware'); @@ -43,6 +51,37 @@ const checkMemoryOptOut = generateCheckAccess({ router.use(requireJwtAuth); +/** Normalizes the optional agent partition param; undefined = shared personal pool */ +const getAgentIdParam = (value) => + typeof value === 'string' && value.trim() !== '' ? value.trim() : undefined; + +/** Resolves agent display names for agent-partitioned memories, restricted + * to agents the requester can VIEW — `agentId` is caller-supplied on write, + * so an unrestricted lookup would leak private agents' names. */ +const withAgentNames = async (memories, user) => { + const agentIds = [...new Set(memories.map((m) => m.agentId).filter(Boolean))]; + if (agentIds.length === 0) { + return memories; + } + try { + const accessibleIds = await findAccessibleResources({ + userId: user.id, + role: user.role, + resourceType: ResourceType.AGENT, + requiredPermissions: PermissionBits.VIEW, + }); + const agents = await getAgents({ id: { $in: agentIds }, _id: { $in: accessibleIds } }); + const namesById = new Map(agents.map((agent) => [agent.id, agent.name])); + return memories.map((memory) => + memory.agentId + ? { ...memory, agentName: namesById.get(memory.agentId) ?? undefined } + : memory, + ); + } catch (_error) { + return memories; + } +}; + /** * GET /memories * Returns all memories for the authenticated user, sorted by updated_at (newest first). @@ -52,12 +91,14 @@ router.get('/', checkMemoryRead, configMiddleware, async (req, res) => { try { const memories = await getAllUserMemories(req.user.id); - const sortedMemories = memories.sort( + const sortedMemories = (await withAgentNames(memories, req.user)).sort( (a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime(), ); + /** Usage totals reflect the shared personal pool only — `tokenLimit` + * applies per partition, matching the inline tools' enforcement. */ const totalTokens = memories.reduce((sum, memory) => { - return sum + (memory.tokenCount || 0); + return sum + (memory.agentId ? 0 : memory.tokenCount || 0); }, 0); const appConfig = req.config; @@ -90,6 +131,7 @@ router.get('/', checkMemoryRead, configMiddleware, async (req, res) => { */ router.post('/', memoryPayloadLimit, checkMemoryCreate, configMiddleware, async (req, res) => { const { key, value } = req.body; + const agentId = getAgentIdParam(req.body.agentId); if (typeof key !== 'string' || key.trim() === '') { return res.status(400).json({ error: 'Key is required and must be a non-empty string.' }); @@ -118,7 +160,7 @@ router.post('/', memoryPayloadLimit, checkMemoryCreate, configMiddleware, async try { const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base'); - const memories = await getAllUserMemories(req.user.id); + const memories = await getUserMemories({ userId: req.user.id, agentId }); const appConfig = req.config; const memoryConfig = appConfig?.memory; @@ -141,13 +183,14 @@ router.post('/', memoryPayloadLimit, checkMemoryCreate, configMiddleware, async key: key.trim(), value: value.trim(), tokenCount, + agentId, }); if (!result.ok) { return res.status(500).json({ error: 'Failed to create memory.' }); } - const updatedMemories = await getAllUserMemories(req.user.id); + const updatedMemories = await getUserMemories({ userId: req.user.id, agentId }); const newMemory = updatedMemories.find((m) => m.key === key.trim()); res.status(201).json({ created: true, memory: newMemory }); @@ -199,6 +242,7 @@ router.patch('/preferences', checkMemoryOptOut, async (req, res) => { router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, async (req, res) => { const { key: urlKey } = req.params; const { key: bodyKey, value } = req.body || {}; + const agentId = getAgentIdParam(req.query.agentId); if (typeof value !== 'string' || value.trim() === '') { return res.status(400).json({ error: 'Value is required and must be a non-empty string.' }); @@ -224,7 +268,7 @@ router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, a try { const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base'); - const memories = await getAllUserMemories(req.user.id); + const memories = await getUserMemories({ userId: req.user.id, agentId }); const existingMemory = memories.find((m) => m.key === urlKey); if (!existingMemory) { @@ -242,13 +286,14 @@ router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, a key: newKey, value, tokenCount, + agentId, }); if (!createResult.ok) { return res.status(500).json({ error: 'Failed to create new memory.' }); } - const deleteResult = await deleteMemory({ userId: req.user.id, key: urlKey }); + const deleteResult = await deleteMemory({ userId: req.user.id, key: urlKey, agentId }); if (!deleteResult.ok) { return res.status(500).json({ error: 'Failed to delete old memory.' }); } @@ -258,6 +303,7 @@ router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, a key: newKey, value, tokenCount, + agentId, }); if (!result.ok) { @@ -265,7 +311,7 @@ router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, a } } - const updatedMemories = await getAllUserMemories(req.user.id); + const updatedMemories = await getUserMemories({ userId: req.user.id, agentId }); const updatedMemory = updatedMemories.find((m) => m.key === newKey); res.json({ updated: true, memory: updatedMemory }); @@ -281,9 +327,10 @@ router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, a */ router.delete('/:key', checkMemoryDelete, async (req, res) => { const { key } = req.params; + const agentId = getAgentIdParam(req.query.agentId); try { - const result = await deleteMemory({ userId: req.user.id, key }); + const result = await deleteMemory({ userId: req.user.id, key, agentId }); if (!result.ok) { return res.status(404).json({ error: 'Memory not found.' }); diff --git a/client/src/common/agents-types.ts b/client/src/common/agents-types.ts index b5e6aefb37..dd6bd79017 100644 --- a/client/src/common/agents-types.ts +++ b/client/src/common/agents-types.ts @@ -5,6 +5,7 @@ import type { AgentToolOptions, SupportContact, AgentProvider, + MemoryScope, GraphEdge, Agent, } from 'librechat-data-provider'; @@ -40,6 +41,8 @@ export type AgentForm = { tool_options?: AgentToolOptions; skills?: string[]; skills_enabled?: boolean; + /** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */ + memory_scope?: MemoryScope; provider?: AgentProvider | OptionWithIcon; /** @deprecated Use edges instead */ agent_ids?: string[]; diff --git a/client/src/components/SidePanel/Agents/AgentPanel.tsx b/client/src/components/SidePanel/Agents/AgentPanel.tsx index c575939794..4e6787bcc4 100644 --- a/client/src/components/SidePanel/Agents/AgentPanel.tsx +++ b/client/src/components/SidePanel/Agents/AgentPanel.tsx @@ -5,6 +5,7 @@ import { useWatch, useForm, FormProvider } from 'react-hook-form'; import { useGetModelsQuery } from 'librechat-data-provider/react-query'; import { Tools, + MemoryScope, SystemRoles, ResourceType, EModelEndpoint, @@ -78,6 +79,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n tool_options, skills, skills_enabled, + memory_scope, avatar_action: avatarActionState, } = data; @@ -107,6 +109,9 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n tool_options, skills, skills_enabled, + /** A hidden stale 'agent' scope must not survive disabling memory — + * runtime partitioning keys off memory_scope alone. */ + memory_scope: data.memory === true ? memory_scope : MemoryScope.user, ...(shouldResetAvatar ? { avatar: null } : {}), }, provider, diff --git a/client/src/components/SidePanel/Agents/Memory.tsx b/client/src/components/SidePanel/Agents/Memory.tsx index 0ce1b489af..69c409f8f6 100644 --- a/client/src/components/SidePanel/Agents/Memory.tsx +++ b/client/src/components/SidePanel/Agents/Memory.tsx @@ -1,6 +1,6 @@ import { memo } from 'react'; -import { AgentCapabilities } from 'librechat-data-provider'; -import { useFormContext, Controller } from 'react-hook-form'; +import { useFormContext, Controller, useWatch } from 'react-hook-form'; +import { MemoryScope, AgentCapabilities } from 'librechat-data-provider'; import { Checkbox, HoverCard, @@ -17,50 +17,101 @@ function Memory() { const localize = useLocalize(); const methods = useFormContext(); const { control } = methods; + const memoryEnabled = useWatch({ control, name: AgentCapabilities.memory }); return ( - -
- ( - - )} - /> - - - - - - -
-

{localize('com_agents_memory_info')}

-
-
-
-
-
+ {localize('com_agents_enable_memory')} + + + + + + +
+

{localize('com_agents_memory_info')}

+
+
+
+ + + {memoryEnabled === true && ( + +
+ ( + + field.onChange(checked === true ? MemoryScope.agent : MemoryScope.user) + } + className="relative float-left mr-2 inline-flex h-4 w-4 cursor-pointer" + value={(field.value === MemoryScope.agent).toString()} + aria-labelledby="memory-scope-label" + /> + )} + /> + + + + + + +
+

+ {localize('com_agents_memory_scope_info')} +

+
+
+
+
+
+ )} + ); } diff --git a/client/src/components/SidePanel/Memories/MemoryCard.tsx b/client/src/components/SidePanel/Memories/MemoryCard.tsx index b160aaf455..6234f84c80 100644 --- a/client/src/components/SidePanel/Memories/MemoryCard.tsx +++ b/client/src/components/SidePanel/Memories/MemoryCard.tsx @@ -27,9 +27,17 @@ export default function MemoryCard({ memory, hasUpdateAccess }: MemoryCardProps) 'hover:bg-surface-secondary', )} > - {/* Row 1: Key + Token count + Actions */} + {/* Row 1: Key + Agent badge + Token count + Actions */}
{memory.key} + {memory.agentId != null && ( + + {memory.agentName ?? memory.agentId} + + )} {memory.tokenCount !== undefined && ( {memory.tokenCount}{' '} diff --git a/client/src/components/SidePanel/Memories/MemoryCardActions.tsx b/client/src/components/SidePanel/Memories/MemoryCardActions.tsx index c1dcb87eed..564d485cd1 100644 --- a/client/src/components/SidePanel/Memories/MemoryCardActions.tsx +++ b/client/src/components/SidePanel/Memories/MemoryCardActions.tsx @@ -39,15 +39,18 @@ export default function MemoryCardActions({ memory }: MemoryCardActionsProps) { ); const confirmDelete = () => { - deleteMemory(memory.key, { - onSuccess: () => { - showToast({ message: localize('com_ui_deleted'), status: 'success' }); - setDeleteOpen(false); + deleteMemory( + { key: memory.key, agentId: memory.agentId }, + { + onSuccess: () => { + showToast({ message: localize('com_ui_deleted'), status: 'success' }); + setDeleteOpen(false); + }, + onError: () => { + showToast({ message: localize('com_ui_error'), status: 'error' }); + }, }, - onError: () => { - showToast({ message: localize('com_ui_error'), status: 'error' }); - }, - }); + ); }; return ( diff --git a/client/src/components/SidePanel/Memories/MemoryEditDialog.tsx b/client/src/components/SidePanel/Memories/MemoryEditDialog.tsx index b6a18e086d..5ec5d4c985 100644 --- a/client/src/components/SidePanel/Memories/MemoryEditDialog.tsx +++ b/client/src/components/SidePanel/Memories/MemoryEditDialog.tsx @@ -117,6 +117,7 @@ export default function MemoryEditDialog({ updateMemory({ key: key.trim(), value: value.trim(), + agentId: memory.agentId, ...(originalKey !== key.trim() && { originalKey }), }); }; diff --git a/client/src/components/SidePanel/Memories/MemoryList.tsx b/client/src/components/SidePanel/Memories/MemoryList.tsx index 30b4236468..6076896044 100644 --- a/client/src/components/SidePanel/Memories/MemoryList.tsx +++ b/client/src/components/SidePanel/Memories/MemoryList.tsx @@ -23,7 +23,7 @@ export default function MemoryList({ return (
{memories.map((memory) => ( -
+
))} diff --git a/client/src/components/SidePanel/Memories/MemoryPanel.tsx b/client/src/components/SidePanel/Memories/MemoryPanel.tsx index 01ec7cb562..89f3e39658 100644 --- a/client/src/components/SidePanel/Memories/MemoryPanel.tsx +++ b/client/src/components/SidePanel/Memories/MemoryPanel.tsx @@ -6,6 +6,7 @@ import { Button, Checkbox, Spinner, + Dropdown, FilterInput, TooltipAnchor, OGDialogTrigger, @@ -25,6 +26,10 @@ import MemoryList from './MemoryList'; const pageSize = 10; +/** Partition filter sentinels; any other value is an agent id */ +const PARTITION_ALL = 'all'; +const PARTITION_PERSONAL = 'personal'; + export default function MemoryPanel() { const localize = useLocalize(); const { user } = useAuthContext(); @@ -33,6 +38,7 @@ export default function MemoryPanel() { const { showToast } = useToastContext(); const [pageIndex, setPageIndex] = useState(0); const [searchQuery, setSearchQuery] = useState(''); + const [partitionFilter, setPartitionFilter] = useState(PARTITION_ALL); const [createDialogOpen, setCreateDialogOpen] = useState(false); const [referenceSavedMemories, setReferenceSavedMemories] = useState(true); @@ -85,20 +91,57 @@ export default function MemoryPanel() { const memories: TUserMemory[] = useMemo(() => memData?.memories ?? [], [memData]); + const partitionOptions = useMemo(() => { + const agentsById = new Map(); + for (const memory of memories) { + if (memory.agentId != null && !agentsById.has(memory.agentId)) { + agentsById.set(memory.agentId, memory.agentName ?? memory.agentId); + } + } + if (agentsById.size === 0) { + return null; + } + return [ + { value: PARTITION_ALL, label: localize('com_ui_memories_all') }, + { value: PARTITION_PERSONAL, label: localize('com_ui_memories_personal') }, + ...[...agentsById.entries()].map(([value, label]) => ({ value, label })), + ]; + }, [memories, localize]); + + /** Falls back to "all" when the selected partition no longer exists + * (e.g. the last memory of that agent was deleted), so the panel never + * gets stuck filtering on a removed partition. */ + const activePartition = useMemo(() => { + if (partitionFilter === PARTITION_ALL || partitionFilter === PARTITION_PERSONAL) { + return partitionFilter; + } + return partitionOptions?.some((option) => option.value === partitionFilter) + ? partitionFilter + : PARTITION_ALL; + }, [partitionOptions, partitionFilter]); + const filteredMemories = useMemo(() => { - return matchSorter(memories, searchQuery, { + const partitionMemories = + activePartition === PARTITION_ALL + ? memories + : memories.filter((memory) => + activePartition === PARTITION_PERSONAL + ? memory.agentId == null + : memory.agentId === activePartition, + ); + return matchSorter(partitionMemories, searchQuery, { keys: ['key', 'value'], }); - }, [memories, searchQuery]); + }, [memories, searchQuery, activePartition]); const currentRows = useMemo(() => { return filteredMemories.slice(pageIndex * pageSize, (pageIndex + 1) * pageSize); }, [filteredMemories, pageIndex]); - // Reset page when search changes + // Reset page when search or partition changes useEffect(() => { setPageIndex(0); - }, [searchQuery]); + }, [searchQuery, activePartition]); if (isLoading) { return ( @@ -155,6 +198,18 @@ export default function MemoryPanel() { )}
+ {/* Partition filter (only when agent-scoped memories exist) */} + {partitionOptions && ( + + )} + {/* Controls: Usage Badge + Memory Toggle */} {(memData?.tokenLimit != null || hasOptOutAccess) && (
diff --git a/client/src/data-provider/Memories/queries.ts b/client/src/data-provider/Memories/queries.ts index 53c2025228..357d2799a9 100644 --- a/client/src/data-provider/Memories/queries.ts +++ b/client/src/data-provider/Memories/queries.ts @@ -1,6 +1,6 @@ /* Memories */ -import { QueryKeys, MutationKeys, dataService } from 'librechat-data-provider'; import { useQuery, useQueryClient, useMutation } from '@tanstack/react-query'; +import { QueryKeys, MutationKeys, dataService } from 'librechat-data-provider'; import type { UseQueryOptions, UseMutationOptions, @@ -19,23 +19,32 @@ export const useMemoriesQuery = ( }); }; +export type DeleteMemoryParams = { key: string; agentId?: string }; export const useDeleteMemoryMutation = () => { const queryClient = useQueryClient(); - return useMutation((key: string) => dataService.deleteMemory(key), { - onSuccess: () => { - queryClient.invalidateQueries([QueryKeys.memories]); + return useMutation( + ({ key, agentId }: DeleteMemoryParams) => dataService.deleteMemory(key, agentId), + { + onSuccess: () => { + queryClient.invalidateQueries([QueryKeys.memories]); + }, }, - }); + ); }; -export type UpdateMemoryParams = { key: string; value: string; originalKey?: string }; +export type UpdateMemoryParams = { + key: string; + value: string; + originalKey?: string; + agentId?: string; +}; export const useUpdateMemoryMutation = ( options?: UseMutationOptions, ) => { const queryClient = useQueryClient(); return useMutation( - ({ key, value, originalKey }: UpdateMemoryParams) => - dataService.updateMemory(key, value, originalKey), + ({ key, value, originalKey, agentId }: UpdateMemoryParams) => + dataService.updateMemory(key, value, originalKey, agentId), { ...options, onSuccess: (...params) => { @@ -74,7 +83,7 @@ export const useUpdateMemoryPreferencesMutation = ( ); }; -export type CreateMemoryParams = { key: string; value: string }; +export type CreateMemoryParams = { key: string; value: string; agentId?: string }; export type CreateMemoryResponse = { created: boolean; memory: TUserMemory }; export const useCreateMemoryMutation = ( @@ -82,7 +91,8 @@ export const useCreateMemoryMutation = ( ) => { const queryClient = useQueryClient(); return useMutation( - ({ key, value }: CreateMemoryParams) => dataService.createMemory({ key, value }), + ({ key, value, agentId }: CreateMemoryParams) => + dataService.createMemory({ key, value, agentId }), { ...options, onSuccess: (data, variables, context) => { @@ -90,8 +100,9 @@ export const useCreateMemoryMutation = ( if (!oldData) return oldData; const newMemories = [...oldData.memories, data.memory]; + /** Usage totals track the shared personal pool only */ const totalTokens = newMemories.reduce( - (sum, memory) => sum + (memory.tokenCount || 0), + (sum, memory) => sum + (memory.agentId ? 0 : memory.tokenCount || 0), 0, ); const tokenLimit = oldData.tokenLimit; diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 1f19dd928d..a5976e8873 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -83,6 +83,8 @@ "com_agents_mcp_name_placeholder": "Custom Tool", "com_agents_mcp_trust_subtext": "Custom connectors are not verified by LibreChat", "com_agents_memory_info": "When enabled, the agent can save, update, and delete memories about the user during the conversation using the set_memory and delete_memory tools.", + "com_agents_memory_scope": "Keep memories separate for this agent", + "com_agents_memory_scope_info": "When enabled, this agent reads and writes its own memories, isolated from your shared memories and from other agents. Existing shared memories are not shown to this agent.", "com_agents_missing_name": "Please type in a name before creating an agent.", "com_agents_missing_provider_model": "Please select a provider and model before creating an agent.", "com_agents_name_placeholder": "Agent name", @@ -1368,8 +1370,12 @@ "com_ui_memories_allow_read": "Allow reading Memories", "com_ui_memories_allow_update": "Allow updating Memories", "com_ui_memories_allow_use": "Allow using Memories", + "com_ui_memories_all": "All memories", "com_ui_memories_filter": "Filter memories...", + "com_ui_memories_partition_filter": "Filter memories by agent", + "com_ui_memories_personal": "Personal", "com_ui_memory": "Memory", + "com_ui_memory_agent_badge": "Agent-scoped memory", "com_ui_memory_already_exceeded": "Memory storage already full - exceeded by {{tokens}} tokens. Delete existing memories before adding new ones.", "com_ui_memory_created": "Memory created successfully", "com_ui_memory_deleted": "Memory deleted", diff --git a/client/src/utils/memory.ts b/client/src/utils/memory.ts index 49b383e9e4..ac26fb2e33 100644 --- a/client/src/utils/memory.ts +++ b/client/src/utils/memory.ts @@ -5,6 +5,13 @@ type HandleMemoryArtifactParams = { currentData: MemoriesResponse; }; +/** Usage totals track the shared personal pool only; agent-partition + * writes never affect the personal usage badge. */ +const isPersonal = (agentId?: string) => agentId == null; + +const samePartition = (memory: TUserMemory, agentId?: string) => + (memory.agentId ?? undefined) === (agentId ?? undefined); + /** * Pure function to handle memory artifact updates * @param params - Object containing memoryArtifact and currentData @@ -14,14 +21,14 @@ export function handleMemoryArtifact({ memoryArtifact, currentData, }: HandleMemoryArtifactParams): MemoriesResponse | undefined { - const { type, key, value, tokenCount = 0 } = memoryArtifact; + const { type, key, value, tokenCount = 0, agentId } = memoryArtifact; if (type === 'update' && !value) { return undefined; } const memories = currentData.memories; - const existingIndex = memories.findIndex((m) => m.key === key); + const existingIndex = memories.findIndex((m) => m.key === key && samePartition(m, agentId)); if (type === 'delete') { if (existingIndex === -1) { @@ -32,7 +39,9 @@ export function handleMemoryArtifact({ const newMemories = [...memories]; newMemories.splice(existingIndex, 1); - const totalTokens = currentData.totalTokens - (deletedMemory.tokenCount || 0); + const totalTokens = isPersonal(agentId) + ? currentData.totalTokens - (deletedMemory.tokenCount || 0) + : currentData.totalTokens; const usagePercentage = currentData.tokenLimit ? Math.min(100, Math.round((totalTokens / currentData.tokenLimit) * 100)) : null; @@ -49,29 +58,30 @@ export function handleMemoryArtifact({ const timestamp = new Date().toISOString(); let totalTokens = currentData.totalTokens; let newMemories: TUserMemory[]; + const updatedMemory: TUserMemory = { + key, + value: value!, + tokenCount, + updated_at: timestamp, + ...(agentId != null ? { agentId } : {}), + }; if (existingIndex >= 0) { const oldTokenCount = memories[existingIndex].tokenCount || 0; - totalTokens = totalTokens - oldTokenCount + tokenCount; + if (isPersonal(agentId)) { + totalTokens = totalTokens - oldTokenCount + tokenCount; + } newMemories = [...memories]; newMemories[existingIndex] = { - key, - value: value!, - tokenCount, - updated_at: timestamp, + ...updatedMemory, + agentName: memories[existingIndex].agentName, }; } else { - totalTokens = totalTokens + tokenCount; - newMemories = [ - ...memories, - { - key, - value: value!, - tokenCount, - updated_at: timestamp, - }, - ]; + if (isPersonal(agentId)) { + totalTokens = totalTokens + tokenCount; + } + newMemories = [...memories, updatedMemory]; } const usagePercentage = currentData.tokenLimit diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 124cf346ff..463a115d00 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/packages/api/src/agents/memory.spec.ts @@ -1,12 +1,14 @@ import { Types } from 'mongoose'; import { Run, Providers } from '@librechat/agents'; +import { MemoryScope } from 'librechat-data-provider'; import type { IUser } from '@librechat/data-schemas'; import type { Response } from 'express'; import { processMemory, createMemoryTool, - createDeleteMemoryTool, + getMemoryAgentId, getRequestMemories, + createDeleteMemoryTool, invalidateRequestMemories, agentHasInlineMemoryTools, } from './memory'; @@ -706,4 +708,44 @@ describe('getRequestMemories caching', () => { await getRequestMemories({ req, userId: 'user-1', getFormattedMemories }); expect(getFormattedMemories).toHaveBeenCalledTimes(2); }); + + it('caches and invalidates per partition', async () => { + const getFormattedMemories = jest + .fn() + .mockResolvedValue({ withKeys: '', withoutKeys: '', totalTokens: 10 }); + const req = {}; + + await getRequestMemories({ req, userId: 'user-1', getFormattedMemories }); + await getRequestMemories({ req, userId: 'user-1', agentId: 'agent_a', getFormattedMemories }); + await getRequestMemories({ req, userId: 'user-1', agentId: 'agent_a', getFormattedMemories }); + /** Personal pool and agent partition are distinct cache entries. */ + expect(getFormattedMemories).toHaveBeenCalledTimes(2); + expect(getFormattedMemories).toHaveBeenLastCalledWith({ + userId: 'user-1', + agentId: 'agent_a', + }); + + /** Invalidating one partition leaves the other cached. */ + invalidateRequestMemories(req, 'agent_a'); + await getRequestMemories({ req, userId: 'user-1', getFormattedMemories }); + expect(getFormattedMemories).toHaveBeenCalledTimes(2); + await getRequestMemories({ req, userId: 'user-1', agentId: 'agent_a', getFormattedMemories }); + expect(getFormattedMemories).toHaveBeenCalledTimes(3); + }); +}); + +describe('getMemoryAgentId', () => { + it('resolves the agent partition only for memory_scope "agent"', () => { + expect(getMemoryAgentId({ id: 'agent_a', memory_scope: MemoryScope.agent })).toBe('agent_a'); + expect(getMemoryAgentId({ id: 'agent_a', memory_scope: MemoryScope.user })).toBeUndefined(); + expect(getMemoryAgentId({ id: 'agent_a' })).toBeUndefined(); + expect(getMemoryAgentId({ memory_scope: MemoryScope.agent })).toBeUndefined(); + expect(getMemoryAgentId(null)).toBeUndefined(); + }); + + it('strips runtime id suffixes so added-conversation runs share the persisted partition', () => { + expect(getMemoryAgentId({ id: 'agent_a____1', memory_scope: MemoryScope.agent })).toBe( + 'agent_a', + ); + }); }); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index e6b3dbebb3..1f4c362be8 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -6,10 +6,12 @@ import { Run, Providers, GraphEvents } from '@librechat/agents'; import { HumanMessage } from '@librechat/agents/langchain/messages'; import { Tools, + MemoryScope, Permissions, EModelEndpoint, PermissionTypes, AgentCapabilities, + stripAgentIdSuffix, } from 'librechat-data-provider'; import type { OpenAIClientOptions, @@ -117,6 +119,7 @@ type MemoryArtifactRecord = Record; */ export const createMemoryTool = ({ userId, + agentId, setMemory, validKeys, charLimit, @@ -125,6 +128,8 @@ export const createMemoryTool = ({ onWrite, }: { userId: string | ObjectId; + /** Agent partition to write to; omit for the shared personal pool */ + agentId?: string; setMemory: MemoryMethods['setMemory']; validKeys?: string[]; charLimit?: number; @@ -220,10 +225,11 @@ export const createMemoryTool = ({ value, tokenCount, type: 'update', + ...(agentId ? { agentId } : {}), }, }; - const result = await setMemory({ userId, key, value, tokenCount }); + const result = await setMemory({ userId, key, value, tokenCount, agentId }); if (result.ok) { if (tokenLimit) { currentTotalTokens = newTotalTokens; @@ -274,11 +280,14 @@ export const createMemoryTool = ({ */ export const createDeleteMemoryTool = ({ userId, + agentId, deleteMemory, validKeys, onWrite, }: { userId: string | ObjectId; + /** Agent partition to delete from; omit for the shared personal pool */ + agentId?: string; deleteMemory: MemoryMethods['deleteMemory']; validKeys?: string[]; onWrite?: () => void; @@ -299,10 +308,11 @@ export const createDeleteMemoryTool = ({ [Tools.memory]: { key, type: 'delete', + ...(agentId ? { agentId } : {}), }, }; - const result = await deleteMemory({ userId, key }); + const result = await deleteMemory({ userId, key, agentId }); if (result.ok) { onWrite?.(); logger.debug(`Memory deleted for key "${key}" for user "${userId}"`); @@ -433,7 +443,25 @@ type GetRoleByName = ( fieldsToSelect?: string | string[], ) => Promise; -type InlineMemoryAgent = { tools?: unknown[]; memoryToolsRegistered?: boolean } | null | undefined; +type InlineMemoryAgent = + | { id?: string; tools?: unknown[]; memoryToolsRegistered?: boolean; memory_scope?: MemoryScope } + | null + | undefined; + +/** + * Resolves the memory partition an agent reads/writes: its persisted id when + * the agent opted into isolated memory (`memory_scope: 'agent'`), otherwise + * `undefined` for the shared personal pool. Runtime ids carry a `____N` + * suffix in added-conversation paths, so the suffix is stripped to keep the + * partition stable across single- and multi-agent runs. Ephemeral agents + * never carry `memory_scope`, so they always resolve to the shared pool. + */ +export function getMemoryAgentId(agent: InlineMemoryAgent): string | undefined { + if (agent?.memory_scope === MemoryScope.agent && typeof agent.id === 'string' && agent.id) { + return stripAgentIdSuffix(agent.id); + } + return undefined; +} /** * Whether an agent carries the inline memory tools. Prefers the LibreChat-only @@ -463,35 +491,44 @@ export function agentHasInlineMemoryTools(agent: InlineMemoryAgent): boolean { /** * Request-scoped cache so that multiple memory-enabled agents in one run (and * the run's memory context load) share a single `getFormattedMemories` call - * instead of each re-fetching the same user's memories. + * per partition instead of each re-fetching the same memories. */ -const requestMemoriesCache = new WeakMap>(); +const requestMemoriesCache = new WeakMap>>(); export function getRequestMemories({ req, userId, + agentId, getFormattedMemories, }: { req: object; userId: string | ObjectId; + /** Agent partition; omit for the shared personal pool */ + agentId?: string; getFormattedMemories: MemoryMethods['getFormattedMemories']; }): Promise { - let cached = requestMemoriesCache.get(req); + let partitions = requestMemoriesCache.get(req); + if (!partitions) { + partitions = new Map(); + requestMemoriesCache.set(req, partitions); + } + const partitionKey = agentId ?? ''; + let cached = partitions.get(partitionKey); if (!cached) { - cached = getFormattedMemories({ userId }); - requestMemoriesCache.set(req, cached); + cached = getFormattedMemories({ userId, agentId }); + partitions.set(partitionKey, cached); } return cached; } /** - * Drops the cached memories for a request so the next {@link getRequestMemories} - * re-fetches. Inline `set_memory`/`delete_memory` writes call this on success so - * a later tool round in the same response is seeded with the post-write usage - * total instead of a stale pre-write one. + * Drops the cached memories for a request partition so the next + * {@link getRequestMemories} re-fetches. Inline `set_memory`/`delete_memory` + * writes call this on success so a later tool round in the same response is + * seeded with the post-write usage total instead of a stale pre-write one. */ -export function invalidateRequestMemories(req: object): void { - requestMemoriesCache.delete(req); +export function invalidateRequestMemories(req: object, agentId?: string): void { + requestMemoriesCache.get(req)?.delete(agentId ?? ''); } /** @@ -563,6 +600,7 @@ export async function buildInlineMemoryTool({ const memoryConfig = req?.config?.memory; const validKeys = memoryConfig?.validKeys as string[] | undefined; + const memoryAgentId = getMemoryAgentId(agent); if (toolName === DELETE_MEMORY_TOOL_NAME) { const allowed = await isMemoryToolAllowed({ @@ -575,9 +613,10 @@ export async function buildInlineMemoryTool({ } return createDeleteMemoryTool({ userId, + agentId: memoryAgentId, deleteMemory: memoryMethods.deleteMemory, validKeys, - onWrite: () => invalidateRequestMemories(req), + onWrite: () => invalidateRequestMemories(req, memoryAgentId), }); } @@ -598,6 +637,7 @@ export async function buildInlineMemoryTool({ const formatted = await getRequestMemories({ req, userId, + agentId: memoryAgentId, getFormattedMemories: memoryMethods.getFormattedMemories, }); totalTokens = formatted?.totalTokens ?? 0; @@ -611,12 +651,13 @@ export async function buildInlineMemoryTool({ return createMemoryTool({ userId, + agentId: memoryAgentId, setMemory: memoryMethods.setMemory, validKeys, charLimit, tokenLimit, totalTokens, - onWrite: () => invalidateRequestMemories(req), + onWrite: () => invalidateRequestMemories(req, memoryAgentId), }); } @@ -647,6 +688,7 @@ export class BasicToolEndHandler implements EventHandler { export async function processMemory({ res, userId, + agentId, setMemory, deleteMemory, messages, @@ -665,6 +707,8 @@ export async function processMemory({ setMemory: MemoryMethods['setMemory']; deleteMemory: MemoryMethods['deleteMemory']; userId: string | ObjectId; + /** Agent partition; omit for the shared personal pool */ + agentId?: string; memory: string; messageId: string; conversationId: string; @@ -680,6 +724,7 @@ export async function processMemory({ try { const memoryTool = createMemoryTool({ userId, + agentId, tokenLimit, setMemory, validKeys, @@ -687,6 +732,7 @@ export async function processMemory({ }); const deleteMemoryTool = createDeleteMemoryTool({ userId, + agentId, validKeys, deleteMemory, }); @@ -874,6 +920,7 @@ ${memory ?? 'No existing memories'}`; export async function createMemoryProcessor({ res, userId, + agentId, messageId, memoryMethods, conversationId, @@ -885,6 +932,8 @@ export async function createMemoryProcessor({ messageId: string; conversationId: string; userId: string | ObjectId; + /** Agent partition; omit for the shared personal pool */ + agentId?: string; memoryMethods: RequiredMemoryMethods; config?: MemoryConfig; streamId?: string | null; @@ -895,6 +944,7 @@ export async function createMemoryProcessor({ const { withKeys, withoutKeys, totalTokens } = await memoryMethods.getFormattedMemories({ userId, + agentId, }); return [ @@ -904,6 +954,7 @@ export async function createMemoryProcessor({ return await processMemory({ res, userId, + agentId, messages, validKeys, llmConfig, diff --git a/packages/api/src/agents/validation.ts b/packages/api/src/agents/validation.ts index 8fc63414e5..666dea3d9f 100644 --- a/packages/api/src/agents/validation.ts +++ b/packages/api/src/agents/validation.ts @@ -1,5 +1,5 @@ import { z } from 'zod'; -import { MAX_SUBAGENTS, ViolationTypes, ErrorTypes } from 'librechat-data-provider'; +import { MemoryScope, MAX_SUBAGENTS, ViolationTypes, ErrorTypes } from 'librechat-data-provider'; import type { Agent, TModelsConfig } from 'librechat-data-provider'; import type { Request, Response } from 'express'; @@ -388,6 +388,7 @@ export const agentBaseSchema: z.ZodObject< tools: z.ZodOptional>; skills: z.ZodOptional>; skills_enabled: z.ZodOptional; + memory_scope: z.ZodOptional>; /** @deprecated Use edges instead */ agent_ids: z.ZodOptional>; edges: z.ZodOptional< @@ -688,6 +689,7 @@ export const agentBaseSchema: z.ZodObject< tools: z.array(z.string()).optional(), skills: z.array(z.string()).optional(), skills_enabled: z.boolean().optional(), + memory_scope: z.nativeEnum(MemoryScope).optional(), /** @deprecated Use edges instead */ agent_ids: z.array(z.string()).optional(), edges: z.array(graphEdgeSchema).optional(), @@ -732,6 +734,7 @@ export const agentCreateSchema: z.ZodObject< model_parameters: z.ZodOptional>; skills: z.ZodOptional>; skills_enabled: z.ZodOptional; + memory_scope: z.ZodOptional>; agent_ids: z.ZodOptional>; edges: z.ZodOptional< z.ZodArray< @@ -1042,6 +1045,7 @@ export const agentUpdateSchema: z.ZodObject< tools: z.ZodOptional>; skills: z.ZodOptional>; skills_enabled: z.ZodOptional; + memory_scope: z.ZodOptional>; agent_ids: z.ZodOptional>; edges: z.ZodOptional< z.ZodArray< diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index b08e9c5106..374a377f19 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -490,7 +490,8 @@ export const verifyTwoFactorTemp = () => `${BASE_URL}/api/auth/2fa/verify-temp`; /* Memories */ export const memories = () => `${BASE_URL}/api/memories`; -export const memory = (key: string) => `${memories()}/${encodeURIComponent(key)}`; +export const memory = (key: string, agentId?: string) => + `${memories()}/${encodeURIComponent(key)}${agentId ? `?agentId=${encodeURIComponent(agentId)}` : ''}`; export const memoryPreferences = () => `${memories()}/preferences`; export const searchPrincipals = (params: q.PrincipalSearchParams) => { diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index d06c138e0d..4d5fa37fa2 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -1307,16 +1307,17 @@ export const getMemories = (): Promise => { return request.get(endpoints.memories()); }; -export const deleteMemory = (key: string): Promise => { - return request.delete(endpoints.memory(key)); +export const deleteMemory = (key: string, agentId?: string): Promise => { + return request.delete(endpoints.memory(key, agentId)); }; export const updateMemory = ( key: string, value: string, originalKey?: string, + agentId?: string, ): Promise => { - return request.patch(endpoints.memory(originalKey || key), { key, value }); + return request.patch(endpoints.memory(originalKey || key, agentId), { key, value }); }; export const updateMemoryPreferences = (preferences: { @@ -1328,6 +1329,7 @@ export const updateMemoryPreferences = (preferences: { export const createMemory = (data: { key: string; value: string; + agentId?: string; }): Promise<{ created: boolean; memory: q.TUserMemory }> => { return request.post(endpoints.memories(), data); }; diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 47e40bf589..f36dee6258 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -344,6 +344,8 @@ export const defaultAgentFormValues = { subagents: undefined as | { enabled?: boolean; allowSelf?: boolean; agent_ids?: string[] } | undefined, + /** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */ + memory_scope: undefined as MemoryScope | undefined, }; export const ImageVisionTool: FunctionTool = { @@ -837,11 +839,23 @@ export const tMessageSchema = z.object({ quotes: z.array(z.string()).optional(), }); +/** + * Which memory partition an agent reads/writes. + * `user` = the shared personal pool (default); `agent` = a partition + * isolated per (user, agent) so the agent only sees its own memories. + */ +export enum MemoryScope { + user = 'user', + agent = 'agent', +} + export type MemoryArtifact = { key: string; value?: string; tokenCount?: number; type: 'update' | 'delete' | 'error'; + /** Agent partition the write targeted; absent = shared personal pool */ + agentId?: string; }; export type UIResource = { diff --git a/packages/data-provider/src/types/assistants.ts b/packages/data-provider/src/types/assistants.ts index bae2a7a72f..57e20f5dbd 100644 --- a/packages/data-provider/src/types/assistants.ts +++ b/packages/data-provider/src/types/assistants.ts @@ -1,5 +1,5 @@ import type { OpenAPIV3 } from 'openapi-types'; -import type { AssistantsEndpoint, AgentProvider } from 'src/schemas'; +import type { AssistantsEndpoint, AgentProvider, MemoryScope } from 'src/schemas'; import type { Agents, GraphEdge } from './agents'; import type { ContentTypes } from './runs'; import type { TFile } from './files'; @@ -301,6 +301,8 @@ export type Agent = { skills_enabled?: boolean; /** Subagent spawning configuration — isolated-context child agents. */ subagents?: AgentSubagentsConfig; + /** Memory partition: `agent` isolates memories per (user, agent); default shared pool */ + memory_scope?: MemoryScope; }; export type TAgentsMap = Record; @@ -329,6 +331,7 @@ export type AgentCreateParams = { | 'skills' | 'skills_enabled' | 'subagents' + | 'memory_scope' >; export type AgentUpdateParams = { @@ -356,6 +359,7 @@ export type AgentUpdateParams = { | 'skills' | 'skills_enabled' | 'subagents' + | 'memory_scope' >; export type AgentListParams = { diff --git a/packages/data-provider/src/types/queries.ts b/packages/data-provider/src/types/queries.ts index f763632d81..b7c968c4e4 100644 --- a/packages/data-provider/src/types/queries.ts +++ b/packages/data-provider/src/types/queries.ts @@ -150,6 +150,10 @@ export type TUserMemory = { value: string; updated_at: string; tokenCount?: number; + /** Agent partition this memory belongs to; absent = shared personal pool */ + agentId?: string; + /** Display name of the partition's agent, resolved server-side when available */ + agentName?: string; }; export type MemoriesResponse = { diff --git a/packages/data-schemas/src/methods/memory.spec.ts b/packages/data-schemas/src/methods/memory.spec.ts new file mode 100644 index 0000000000..48c8c3b5ab --- /dev/null +++ b/packages/data-schemas/src/methods/memory.spec.ts @@ -0,0 +1,156 @@ +import mongoose from 'mongoose'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import type { IMemoryEntry } from '~/types'; +import { createMemoryMethods, type MemoryMethods } from './memory'; +import { createModels } from '~/models'; + +jest.mock('~/config/winston', () => ({ + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), +})); + +let mongoServer: InstanceType; +let MemoryEntry: mongoose.Model; +let methods: MemoryMethods; + +const userId = new mongoose.Types.ObjectId(); +const otherUserId = new mongoose.Types.ObjectId(); +const agentId = 'agent_partition_a'; +const otherAgentId = 'agent_partition_b'; + +beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + createModels(mongoose); + MemoryEntry = mongoose.models.MemoryEntry as mongoose.Model; + methods = createMemoryMethods(mongoose); + await mongoose.connect(mongoServer.getUri()); +}); + +afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); +}); + +beforeEach(async () => { + await MemoryEntry.deleteMany({}); +}); + +describe('memory partitions', () => { + it('writes to the shared personal pool when agentId is omitted', async () => { + await methods.setMemory({ userId, key: 'preference', value: 'likes tea', tokenCount: 3 }); + + const personal = await methods.getUserMemories({ userId }); + expect(personal).toHaveLength(1); + expect(personal[0].agentId == null).toBe(true); + + const agentScoped = await methods.getUserMemories({ userId, agentId }); + expect(agentScoped).toHaveLength(0); + }); + + it('isolates the same key across partitions', async () => { + await methods.setMemory({ userId, key: 'context', value: 'personal', tokenCount: 1 }); + await methods.setMemory({ userId, key: 'context', value: 'agent a', tokenCount: 2, agentId }); + await methods.setMemory({ + userId, + key: 'context', + value: 'agent b', + tokenCount: 3, + agentId: otherAgentId, + }); + + const personal = await methods.getUserMemories({ userId }); + const partitionA = await methods.getUserMemories({ userId, agentId }); + const partitionB = await methods.getUserMemories({ userId, agentId: otherAgentId }); + + expect(personal.map((m) => m.value)).toEqual(['personal']); + expect(partitionA.map((m) => m.value)).toEqual(['agent a']); + expect(partitionB.map((m) => m.value)).toEqual(['agent b']); + expect(await MemoryEntry.countDocuments({ userId })).toBe(3); + }); + + it('upserts within a partition instead of duplicating', async () => { + await methods.setMemory({ userId, key: 'context', value: 'first', tokenCount: 1, agentId }); + await methods.setMemory({ userId, key: 'context', value: 'second', tokenCount: 2, agentId }); + + const partitionA = await methods.getUserMemories({ userId, agentId }); + expect(partitionA).toHaveLength(1); + expect(partitionA[0].value).toBe('second'); + }); + + it('treats legacy entries without an agentId field as the personal pool', async () => { + await MemoryEntry.collection.insertOne({ + userId, + key: 'legacy', + value: 'pre-partition entry', + tokenCount: 4, + updated_at: new Date(), + }); + + const personal = await methods.getUserMemories({ userId }); + expect(personal.map((m) => m.key)).toEqual(['legacy']); + + await methods.setMemory({ userId, key: 'legacy', value: 'updated', tokenCount: 4 }); + expect(await MemoryEntry.countDocuments({ userId })).toBe(1); + expect((await methods.getUserMemories({ userId }))[0].value).toBe('updated'); + }); + + it('deletes only within the targeted partition', async () => { + await methods.setMemory({ userId, key: 'context', value: 'personal', tokenCount: 1 }); + await methods.setMemory({ userId, key: 'context', value: 'agent a', tokenCount: 2, agentId }); + + const agentDelete = await methods.deleteMemory({ userId, key: 'context', agentId }); + expect(agentDelete.ok).toBe(true); + expect(await methods.getUserMemories({ userId })).toHaveLength(1); + expect(await methods.getUserMemories({ userId, agentId })).toHaveLength(0); + + const missingDelete = await methods.deleteMemory({ userId, key: 'context', agentId }); + expect(missingDelete.ok).toBe(false); + }); + + it('scopes createMemory duplicate detection to the partition', async () => { + await methods.createMemory({ userId, key: 'context', value: 'personal', tokenCount: 1 }); + + await expect( + methods.createMemory({ userId, key: 'context', value: 'dupe', tokenCount: 1 }), + ).rejects.toThrow('already exists'); + + const created = await methods.createMemory({ + userId, + key: 'context', + value: 'agent a', + tokenCount: 2, + agentId, + }); + expect(created.ok).toBe(true); + }); + + it('formats memories per partition with partition-scoped totals', async () => { + await methods.setMemory({ userId, key: 'one', value: 'personal one', tokenCount: 5 }); + await methods.setMemory({ userId, key: 'two', value: 'agent two', tokenCount: 7, agentId }); + + const personal = await methods.getFormattedMemories({ userId }); + expect(personal.totalTokens).toBe(5); + expect(personal.withoutKeys).toContain('personal one'); + expect(personal.withoutKeys).not.toContain('agent two'); + + const partitionA = await methods.getFormattedMemories({ userId, agentId }); + expect(partitionA.totalTokens).toBe(7); + expect(partitionA.withKeys).toContain('"key": "two"'); + expect(partitionA.withKeys).not.toContain('personal one'); + }); + + it('returns every partition from getAllUserMemories and wipes them all on deleteAllUserMemories', async () => { + await methods.setMemory({ userId, key: 'one', value: 'personal', tokenCount: 1 }); + await methods.setMemory({ userId, key: 'two', value: 'agent a', tokenCount: 1, agentId }); + await methods.setMemory({ userId: otherUserId, key: 'three', value: 'other', tokenCount: 1 }); + + const all = await methods.getAllUserMemories(userId); + expect(all).toHaveLength(2); + + const deleted = await methods.deleteAllUserMemories(userId); + expect(deleted).toBe(2); + expect(await MemoryEntry.countDocuments({ userId: otherUserId })).toBe(1); + }); +}); diff --git a/packages/data-schemas/src/methods/memory.ts b/packages/data-schemas/src/methods/memory.ts index 7c89535198..6e5cc82263 100644 --- a/packages/data-schemas/src/methods/memory.ts +++ b/packages/data-schemas/src/methods/memory.ts @@ -9,14 +9,32 @@ const formatDate = (date: Date): string => { return date.toISOString().split('T')[0]; }; +/** Partition filter: `agentId: null` matches both null and missing fields, + * so pre-partition entries remain part of the shared personal pool. */ +const partitionFilter = (agentId?: string) => ({ agentId: agentId ?? null }); + // Factory function that takes mongoose instance and returns the methods export function createMemoryMethods(mongoose: typeof import('mongoose')): { - setMemory: ({ userId, key, value, tokenCount }: t.SetMemoryParams) => Promise; - createMemory: ({ userId, key, value, tokenCount }: t.SetMemoryParams) => Promise; - deleteMemory: ({ userId, key }: t.DeleteMemoryParams) => Promise; + setMemory: ({ + userId, + key, + value, + tokenCount, + agentId, + }: t.SetMemoryParams) => Promise; + createMemory: ({ + userId, + key, + value, + tokenCount, + agentId, + }: t.SetMemoryParams) => Promise; + deleteMemory: ({ userId, key, agentId }: t.DeleteMemoryParams) => Promise; getAllUserMemories: (userId: string | Types.ObjectId) => Promise; + getUserMemories: ({ userId, agentId }: t.GetUserMemoriesParams) => Promise; getFormattedMemories: ({ userId, + agentId, }: t.GetFormattedMemoriesParams) => Promise; deleteAllUserMemories: (userId: string | Types.ObjectId) => Promise; } { @@ -29,6 +47,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { key, value, tokenCount = 0, + agentId, }: t.SetMemoryParams): Promise { try { if (key?.toLowerCase() === 'nothing') { @@ -36,7 +55,11 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { } const MemoryEntry = mongoose.models.MemoryEntry; - const existingMemory = await MemoryEntry.findOne({ userId, key }); + const existingMemory = await MemoryEntry.findOne({ + userId, + key, + ...partitionFilter(agentId), + }); if (existingMemory) { throw new Error('Memory with this key already exists'); } @@ -46,6 +69,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { key, value, tokenCount, + ...(agentId ? { agentId } : {}), updated_at: new Date(), }); @@ -65,6 +89,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { key, value, tokenCount = 0, + agentId, }: t.SetMemoryParams): Promise { try { if (key?.toLowerCase() === 'nothing') { @@ -73,10 +98,11 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { const MemoryEntry = mongoose.models.MemoryEntry; await MemoryEntry.findOneAndUpdate( - { userId, key }, + { userId, key, ...partitionFilter(agentId) }, { value, tokenCount, + ...(agentId ? { agentId } : {}), updated_at: new Date(), }, { @@ -96,10 +122,18 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { /** * Deletes a specific memory entry for a user */ - async function deleteMemory({ userId, key }: t.DeleteMemoryParams): Promise { + async function deleteMemory({ + userId, + key, + agentId, + }: t.DeleteMemoryParams): Promise { try { const MemoryEntry = mongoose.models.MemoryEntry; - const result = await MemoryEntry.findOneAndDelete({ userId, key }); + const result = await MemoryEntry.findOneAndDelete({ + userId, + key, + ...partitionFilter(agentId), + }); return { ok: !!result }; } catch (error) { throw new Error( @@ -109,7 +143,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { } /** - * Gets all memory entries for a user + * Gets all memory entries for a user across every partition */ async function getAllUserMemories( userId: string | Types.ObjectId, @@ -125,13 +159,35 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { } /** - * Gets and formats all memories for a user in two different formats + * Gets a user's memory entries for a single partition + * (the shared personal pool when `agentId` is omitted) + */ + async function getUserMemories({ + userId, + agentId, + }: t.GetUserMemoriesParams): Promise { + try { + const MemoryEntry = mongoose.models.MemoryEntry; + return (await MemoryEntry.find({ + userId, + ...partitionFilter(agentId), + }).lean()) as t.IMemoryEntryLean[]; + } catch (error) { + throw new Error( + `Failed to get memories: ${error instanceof Error ? error.message : 'Unknown error'}`, + ); + } + } + + /** + * Gets and formats a partition's memories for a user in two different formats */ async function getFormattedMemories({ userId, + agentId, }: t.GetFormattedMemoriesParams): Promise { try { - const memories = await getAllUserMemories(userId); + const memories = await getUserMemories({ userId, agentId }); if (!memories || memories.length === 0) { return { withKeys: '', withoutKeys: '', totalTokens: 0 }; @@ -187,6 +243,7 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): { createMemory, deleteMemory, getAllUserMemories, + getUserMemories, getFormattedMemories, deleteAllUserMemories, }; diff --git a/packages/data-schemas/src/schema/agent.ts b/packages/data-schemas/src/schema/agent.ts index 5ef6a9cf0d..7000d28688 100644 --- a/packages/data-schemas/src/schema/agent.ts +++ b/packages/data-schemas/src/schema/agent.ts @@ -125,6 +125,12 @@ const agentSchema: Schema = new Schema( type: Schema.Types.Mixed, default: undefined, }, + /** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */ + memory_scope: { + type: String, + enum: ['user', 'agent'], + default: undefined, + }, tenantId: { type: String, index: true, diff --git a/packages/data-schemas/src/schema/memory.ts b/packages/data-schemas/src/schema/memory.ts index 773fa87115..fc58833458 100644 --- a/packages/data-schemas/src/schema/memory.ts +++ b/packages/data-schemas/src/schema/memory.ts @@ -20,6 +20,11 @@ const MemoryEntrySchema: Schema = new Schema({ type: String, required: true, }, + /** Agent partition; null/absent = shared personal pool */ + agentId: { + type: String, + default: undefined, + }, tokenCount: { type: Number, default: 0, @@ -34,4 +39,6 @@ const MemoryEntrySchema: Schema = new Schema({ }, }); +MemoryEntrySchema.index({ userId: 1, agentId: 1, key: 1 }); + export default MemoryEntrySchema; diff --git a/packages/data-schemas/src/types/agent.ts b/packages/data-schemas/src/types/agent.ts index 416112141d..5c833c3484 100644 --- a/packages/data-schemas/src/types/agent.ts +++ b/packages/data-schemas/src/types/agent.ts @@ -1,6 +1,7 @@ import { Document, Types } from 'mongoose'; import type { GraphEdge, + MemoryScope, AgentToolOptions, AgentToolResources, AgentSubagentsConfig, @@ -50,5 +51,7 @@ export interface IAgent extends Omit { tool_options?: AgentToolOptions; /** Subagent spawning configuration — isolated-context child agents. */ subagents?: AgentSubagentsConfig; + /** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */ + memory_scope?: MemoryScope; tenantId?: string; } diff --git a/packages/data-schemas/src/types/memory.ts b/packages/data-schemas/src/types/memory.ts index 4d9b4fbefd..06f5905496 100644 --- a/packages/data-schemas/src/types/memory.ts +++ b/packages/data-schemas/src/types/memory.ts @@ -5,6 +5,8 @@ export interface IMemoryEntry extends Document { userId: Types.ObjectId; key: string; value: string; + /** Agent partition; null/absent = shared personal pool */ + agentId?: string; tokenCount?: number; updated_at?: Date; tenantId?: string; @@ -15,6 +17,7 @@ export interface IMemoryEntryLean { userId: Types.ObjectId; key: string; value: string; + agentId?: string; tokenCount?: number; updated_at?: Date; __v?: number; @@ -26,15 +29,24 @@ export interface SetMemoryParams { key: string; value: string; tokenCount?: number; + /** Agent partition; omit for the shared personal pool */ + agentId?: string; } export interface DeleteMemoryParams { userId: string | Types.ObjectId; key: string; + agentId?: string; +} + +export interface GetUserMemoriesParams { + userId: string | Types.ObjectId; + agentId?: string; } export interface GetFormattedMemoriesParams { userId: string | Types.ObjectId; + agentId?: string; } // Result interfaces