diff --git a/packages/twenty-front/src/modules/ai/components/AgentChatStreamingPartsDiffSyncEffect.tsx b/packages/twenty-front/src/modules/ai/components/AgentChatStreamingPartsDiffSyncEffect.tsx index 7761ab11ec5..560b776c7f2 100644 --- a/packages/twenty-front/src/modules/ai/components/AgentChatStreamingPartsDiffSyncEffect.tsx +++ b/packages/twenty-front/src/modules/ai/components/AgentChatStreamingPartsDiffSyncEffect.tsx @@ -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>(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); }, [ diff --git a/packages/twenty-front/src/modules/ai/hooks/useUpdateStreamingPartsWithDiff.ts b/packages/twenty-front/src/modules/ai/hooks/useUpdateStreamingPartsWithDiff.ts index 07dede10f86..63e1ec4fdfa 100644 --- a/packages/twenty-front/src/modules/ai/hooks/useUpdateStreamingPartsWithDiff.ts +++ b/packages/twenty-front/src/modules/ai/hooks/useUpdateStreamingPartsWithDiff.ts @@ -36,11 +36,9 @@ export const useUpdateStreamingPartsWithDiff = () => { continue; } - const clonedMessage = structuredClone(incomingMessage); - jotaiStore.set( agentChatMessageFamilyCallbackState(incomingMessage.id), - clonedMessage, + incomingMessage, ); processStreamingMessageUpdate(incomingMessage); diff --git a/packages/twenty-front/src/modules/ui/utilities/state/jotai/types/ComponentFamilyState.ts b/packages/twenty-front/src/modules/ui/utilities/state/jotai/types/ComponentFamilyState.ts index 1eb3702bcb4..7a8d0c3c992 100644 --- a/packages/twenty-front/src/modules/ui/utilities/state/jotai/types/ComponentFamilyState.ts +++ b/packages/twenty-front/src/modules/ui/utilities/state/jotai/types/ComponentFamilyState.ts @@ -18,4 +18,5 @@ export type ComponentFamilyState = { atomFamily: ( key: ComponentFamilyStateKey, ) => JotaiWritableAtom; + evictFamilyKey: (key: ComponentFamilyStateKey) => void; }; diff --git a/packages/twenty-front/src/modules/ui/utilities/state/jotai/utils/createAtomComponentFamilyState.ts b/packages/twenty-front/src/modules/ui/utilities/state/jotai/utils/createAtomComponentFamilyState.ts index fe056f51e81..f5245e1bd56 100644 --- a/packages/twenty-front/src/modules/ui/utilities/state/jotai/utils/createAtomComponentFamilyState.ts +++ b/packages/twenty-front/src/modules/ui/utilities/state/jotai/utils/createAtomComponentFamilyState.ts @@ -49,9 +49,21 @@ export const createAtomComponentFamilyState = ({ return baseAtom; }; + const evictFunction = ({ + instanceId, + familyKey, + }: ComponentFamilyStateKey): 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, }; };