mirror of
https://github.com/twentyhq/twenty
synced 2026-04-21 13:37:22 +00:00
fix(twenty-front): add eviction to atom family cache and deduplicate AI chat message storage
https://sonarly.com/issue/27119?type=bug Chrome tab crashes with OOM (error code 5) because the Jotai atom family cache (`atomCache` Map) never evicts entries, and each AI chat message is stored three times (thread-level atom, per-message atom via structuredClone, Apollo InMemoryCache), causing unbounded heap growth during long conversations, rapid workspace switching, or Settings page navigation. Fix: Three changes that address the two largest memory growth vectors in the AI chat module: 1. **Added `evictFamilyKey` to `ComponentFamilyState`** (`ComponentFamilyState.ts` type + `createAtomComponentFamilyState.ts` implementation): Exposes `atomCache.delete()` so consumers can remove atoms they no longer need. This is the missing primitive — the atomCache Map previously had no way to remove entries. 2. **Evict per-message atoms on thread switch** (`AgentChatStreamingPartsDiffSyncEffect.tsx`): Tracks message IDs seen during diff sync in a ref. When `currentAIChatThread` changes, the cleanup effect fires and calls `evictFamilyKey` for every tracked message ID, removing those atoms from the cache. This prevents the biggest accumulation vector: hundreds of per-message atoms (keyed by UUID) holding large ExtendedUIMessage payloads that were never released. 3. **Remove `structuredClone` from `useUpdateStreamingPartsWithDiff.ts`**: The deep clone was creating a full copy of every message object before storing it in the per-message atom. Since `isDeeplyEqual` already gates updates (preventing unnecessary writes), and the message reference from the thread-level atom is stable within a render cycle, the clone is unnecessary. Removing it eliminates one full copy of every message in memory. Together these changes bound per-message atom memory to the current thread's messages (instead of all threads ever viewed), and halve the per-message memory footprint by eliminating deep clones. **Not addressed in this PR** (requires separate, wider-scoped changes): - Apollo InMemoryCache eviction — needs typePolicies configuration in the shared Apollo factory - Thread-level family state atoms (10 per thread) — small scalar values, low priority
This commit is contained in:
parent
d3df58046c
commit
e87fea949e
4 changed files with 39 additions and 4 deletions
|
|
@ -1,11 +1,14 @@
|
|||
import { useCallback, useEffect, useRef } from 'react';
|
||||
|
||||
import { AGENT_CHAT_INSTANCE_ID } from '@/ai/constants/AgentChatInstanceId';
|
||||
import { useUpdateStreamingPartsWithDiff } from '@/ai/hooks/useUpdateStreamingPartsWithDiff';
|
||||
import { agentChatLastDiffSyncedThreadState } from '@/ai/states/agentChatLastDiffSyncedThreadState';
|
||||
import { agentChatMessageComponentFamilyState } from '@/ai/states/agentChatMessageComponentFamilyState';
|
||||
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
|
||||
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
|
||||
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
|
||||
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
|
||||
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
|
||||
import { useEffect } from 'react';
|
||||
|
||||
export const AgentChatStreamingPartsDiffSyncEffect = () => {
|
||||
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
|
||||
|
|
@ -21,11 +24,32 @@ export const AgentChatStreamingPartsDiffSyncEffect = () => {
|
|||
agentChatLastDiffSyncedThreadState,
|
||||
);
|
||||
|
||||
const trackedMessageIdsRef = useRef<Set<string>>(new Set());
|
||||
|
||||
const evictTrackedMessageAtoms = useCallback(() => {
|
||||
for (const messageId of trackedMessageIdsRef.current) {
|
||||
agentChatMessageComponentFamilyState.evictFamilyKey({
|
||||
instanceId: AGENT_CHAT_INSTANCE_ID,
|
||||
familyKey: messageId,
|
||||
});
|
||||
}
|
||||
trackedMessageIdsRef.current.clear();
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
return () => {
|
||||
evictTrackedMessageAtoms();
|
||||
};
|
||||
}, [currentAIChatThread, evictTrackedMessageAtoms]);
|
||||
|
||||
useEffect(() => {
|
||||
if (agentChatMessages.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const currentIds = new Set(agentChatMessages.map((m) => m.id));
|
||||
trackedMessageIdsRef.current = currentIds;
|
||||
|
||||
updateStreamingPartsWithDiff(agentChatMessages);
|
||||
setAgentChatLastDiffSyncedThread(currentAIChatThread);
|
||||
}, [
|
||||
|
|
|
|||
|
|
@ -36,11 +36,9 @@ export const useUpdateStreamingPartsWithDiff = () => {
|
|||
continue;
|
||||
}
|
||||
|
||||
const clonedMessage = structuredClone(incomingMessage);
|
||||
|
||||
jotaiStore.set(
|
||||
agentChatMessageFamilyCallbackState(incomingMessage.id),
|
||||
clonedMessage,
|
||||
incomingMessage,
|
||||
);
|
||||
|
||||
processStreamingMessageUpdate(incomingMessage);
|
||||
|
|
|
|||
|
|
@ -18,4 +18,5 @@ export type ComponentFamilyState<ValueType, FamilyKey> = {
|
|||
atomFamily: (
|
||||
key: ComponentFamilyStateKey<FamilyKey>,
|
||||
) => JotaiWritableAtom<ValueType>;
|
||||
evictFamilyKey: (key: ComponentFamilyStateKey<FamilyKey>) => void;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -49,9 +49,21 @@ export const createAtomComponentFamilyState = <ValueType, FamilyKey>({
|
|||
return baseAtom;
|
||||
};
|
||||
|
||||
const evictFunction = ({
|
||||
instanceId,
|
||||
familyKey,
|
||||
}: ComponentFamilyStateKey<FamilyKey>): void => {
|
||||
const familyKeyStr =
|
||||
typeof familyKey === 'string' ? familyKey : JSON.stringify(familyKey);
|
||||
|
||||
const cacheKey = `${instanceId}__${familyKeyStr}`;
|
||||
atomCache.delete(cacheKey);
|
||||
};
|
||||
|
||||
return {
|
||||
type: 'ComponentFamilyState',
|
||||
key,
|
||||
atomFamily: familyFunction,
|
||||
evictFamilyKey: evictFunction,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in a new issue