houseassassin/twenty
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Convert AI chat state atoms to component family states (#19585)

Browse Source 
## Summary
Refactored AI chat state management to use component family states
instead of global atoms. This change enables proper state isolation per
chat thread, preventing state leakage between different chat instances.

## Key Changes

- **Converted global atoms to component family states:**
  - `agentChatErrorState` → `agentChatErrorComponentFamilyState`
- `agentChatIsStreamingState` →
`agentChatIsStreamingComponentFamilyState`
- `agentChatFirstLiveSeqState` →
`agentChatFirstLiveSeqComponentFamilyState`
- `agentChatHandleEventCallbackState` →
`agentChatHandleEventCallbackComponentFamilyState`
  - `agentChatUsageState` → `agentChatUsageComponentFamilyState`
- `currentAIChatThreadTitleState` →
`currentAIChatThreadTitleComponentFamilyState`

- **Updated state access patterns:**
- Introduced `useAtomComponentFamilyStateCallbackState` hook to get
family state callbacks
- Changed from direct atom access to family key-based access using `{
threadId }` as the family key
- Updated all components and hooks to use the new family state patterns

- **Removed instance ID dependency:**
  - Eliminated `AGENT_CHAT_INSTANCE_ID` constant usage
- Simplified state access by using thread ID directly as the family key

- **Updated affected components and hooks:**
- `useAgentChatSubscription`: Now creates family state atoms per thread
- `AgentChatMessagesFetchEffect`: Uses family state callbacks for event
handling
- `AgentChatThreadInitializationEffect`: Sets family state values per
thread
  - `useAIChatThreadClick`: Updates family states when switching threads
- `AIChatErrorUnderMessageList`, `AIChatLastMessageWithStreamingState`,
`AIChatEmptyState`, `AIChatStandaloneError`, `SendMessageButton`,
`AIChatContextUsageButton`, `SidePanelAskAIInfo`: Updated to use family
state values with thread ID

## Implementation Details

- All new component family states use
`AgentChatComponentInstanceContext` for proper component-level isolation
- Family key structure: `{ threadId: string | null }`
- Removed `disposed` flag logic as it's no longer needed with proper
state isolation
- Updated dependency arrays in hooks to include family callback
functions

https://claude.ai/code/session_01D8syiJtCXu5bEHSr5KYkCt

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Félix Malfait
2026-04-12 07:39:12 +02:00
committed by GitHub
parent c268a50e4e
commit 55b1624210
26 changed files with 294 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/twenty-front/src/modules/ai/components/AIChatEmptyState.tsx
+8 -4
View File
@@ -4,12 +4,13 @@ import { isDefined } from 'twenty-shared/utils';
import { AIChatSuggestedPrompts } from '@/ai/components/suggested-prompts/AIChatSuggestedPrompts';
import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatErrorComponentFamilyState } from '@/ai/states/agentChatErrorComponentFamilyState';
import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
@@ -26,7 +27,11 @@ type AIChatEmptyStateProps = {
};
export const AIChatEmptyState = ({ editor }: AIChatEmptyStateProps) => {
const agentChatError = useAtomStateValue(agentChatErrorState);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatError = useAtomComponentFamilyStateValue(
agentChatErrorComponentFamilyState,
{ threadId: currentAIChatThread },
);
const agentChatThreadsLoading = useAtomStateValue(
agentChatThreadsLoadingState,
);
@@ -36,7 +41,6 @@ export const AIChatEmptyState = ({ editor }: AIChatEmptyStateProps) => {
const skipMessagesSkeletonUntilLoaded = useAtomStateValue(
skipMessagesSkeletonUntilLoadedState,
);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const hasMessages = useAtomComponentSelectorValue(
agentChatHasMessageComponentSelector,
packages/twenty-front/src/modules/ai/components/AIChatErrorUnderMessageList.tsx
+15 -4
View File
@@ -1,10 +1,12 @@
import { AIChatErrorRenderer } from '@/ai/components/AIChatErrorRenderer';
import { AgentMessageRole } from '@/ai/constants/AgentMessageRole';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThreadState';
import { agentChatErrorComponentFamilyState } from '@/ai/states/agentChatErrorComponentFamilyState';
import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
import { agentChatMessageComponentFamilySelector } from '@/ai/states/agentChatMessageComponentFamilySelector';
import { agentChatMessageIdsComponentSelector } from '@/ai/states/agentChatMessageIdsComponentSelector';
import { useAtomComponentFamilySelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilySelectorValue';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { styled } from '@linaria/react';
@@ -15,8 +17,17 @@ const StyledErrorWrapper = styled.div`
`;
export const AIChatErrorUnderMessageList = () => {
const agentChatError = useAtomStateValue(agentChatErrorState);
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const agentChatDisplayedThread = useAtomStateValue(
agentChatDisplayedThreadState,
);
const agentChatError = useAtomComponentFamilyStateValue(
agentChatErrorComponentFamilyState,
{ threadId: agentChatDisplayedThread },
);
const agentChatIsStreaming = useAtomComponentFamilyStateValue(
agentChatIsStreamingComponentFamilyState,
{ threadId: agentChatDisplayedThread },
);
const agentChatMessageIds = useAtomComponentSelectorValue(
agentChatMessageIdsComponentSelector,
packages/twenty-front/src/modules/ai/components/AIChatLastMessageWithStreamingState.tsx
+15 -4
View File
@@ -1,7 +1,9 @@
import { AIChatMessage } from '@/ai/components/AIChatMessage';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThreadState';
import { agentChatErrorComponentFamilyState } from '@/ai/states/agentChatErrorComponentFamilyState';
import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
import { agentChatLastMessageIdComponentSelector } from '@/ai/states/agentChatLastMessageIdComponentSelector';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { isDefined } from 'twenty-shared/utils';
@@ -11,8 +13,17 @@ export const AIChatLastMessageWithStreamingState = () => {
agentChatLastMessageIdComponentSelector,
);
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const agentChatError = useAtomStateValue(agentChatErrorState);
const agentChatDisplayedThread = useAtomStateValue(
agentChatDisplayedThreadState,
);
const agentChatIsStreaming = useAtomComponentFamilyStateValue(
agentChatIsStreamingComponentFamilyState,
{ threadId: agentChatDisplayedThread },
);
const agentChatError = useAtomComponentFamilyStateValue(
agentChatErrorComponentFamilyState,
{ threadId: agentChatDisplayedThread },
);
if (!isDefined(lastMessageId)) {
return null;
packages/twenty-front/src/modules/ai/components/AIChatStandaloneError.tsx
+8 -2
View File
@@ -1,9 +1,11 @@
import { styled } from '@linaria/react';
import { AIChatErrorRenderer } from '@/ai/components/AIChatErrorRenderer';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatErrorComponentFamilyState } from '@/ai/states/agentChatErrorComponentFamilyState';
import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { isDefined } from 'twenty-shared/utils';
@@ -18,7 +20,11 @@ const StyledErrorContainer = styled.div`
export const AIChatStandaloneError = () => {
const agentChatIsLoading = useAtomStateValue(agentChatIsLoadingState);
const agentChatError = useAtomStateValue(agentChatErrorState);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatError = useAtomComponentFamilyStateValue(
agentChatErrorComponentFamilyState,
{ threadId: currentAIChatThread },
);
const hasMessages = useAtomComponentSelectorValue(
agentChatHasMessageComponentSelector,
packages/twenty-front/src/modules/ai/components/AgentChatMessagesFetchEffect.tsx
+30 -5
View File
@@ -4,8 +4,8 @@ import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { agentChatFirstLiveSeqState } from '@/ai/states/agentChatFirstLiveSeqState';
import { agentChatHandleEventCallbackState } from '@/ai/states/agentChatHandleEventCallbackState';
import { agentChatFirstLiveSeqComponentFamilyState } from '@/ai/states/agentChatFirstLiveSeqComponentFamilyState';
import { agentChatHandleEventCallbackComponentFamilyState } from '@/ai/states/agentChatHandleEventCallbackComponentFamilyState';
import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
@@ -15,6 +15,7 @@ import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSk
import { mapDBMessagesToUIMessages } from '@/ai/utils/mapDBMessagesToUIMessages';
import { useQueryWithCallbacks } from '@/apollo/hooks/useQueryWithCallbacks';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
import { useAtomComponentFamilyStateCallbackState } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateCallbackState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomComponentFamilyState } from '@/ui/utilities/state/jotai/hooks/useSetAtomComponentFamilyState';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
@@ -52,6 +53,14 @@ export const AgentChatMessagesFetchEffect = () => {
{ threadId: currentAIChatThread },
);
const handleEventCallbackFamilyCallback =
useAtomComponentFamilyStateCallbackState(
agentChatHandleEventCallbackComponentFamilyState,
);
const firstLiveSeqFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatFirstLiveSeqComponentFamilyState,
);
const handleFirstLoad = useCallback(
(_data: GetChatMessagesQuery) => {
setSkipMessagesSkeletonUntilLoaded(false);
@@ -75,13 +84,23 @@ export const AgentChatMessagesFetchEffect = () => {
return;
}
const handleEvent = store.get(agentChatHandleEventCallbackState.atom);
const threadId = store.get(currentAIChatThreadState.atom);
if (!isDefined(threadId)) {
return;
}
const familyKey = { threadId };
const handleEvent = store.get(
handleEventCallbackFamilyCallback(familyKey),
);
if (!isDefined(handleEvent)) {
return;
}
const firstLiveSeq = store.get(agentChatFirstLiveSeqState.atom);
const firstLiveSeq = store.get(firstLiveSeqFamilyCallback(familyKey));
for (let index = 0; index < catchup.chunks.length; index++) {
const chunkSeq = index + 1;
@@ -97,7 +116,13 @@ export const AgentChatMessagesFetchEffect = () => {
} as AgentChatSubscriptionEvent);
}
},
[setAgentChatFetchedMessages, setAgentChatQueuedMessages, store],
[
setAgentChatFetchedMessages,
setAgentChatQueuedMessages,
store,
handleEventCallbackFamilyCallback,
firstLiveSeqFamilyCallback,
],
);
const handleLoadingChange = useCallback(
packages/twenty-front/src/modules/ai/components/AgentChatStreamSubscriptionEffect.tsx
+5 -2
View File
@@ -11,7 +11,7 @@ import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThr
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatIsInitialScrollPendingOnThreadChangeState } from '@/ai/states/agentChatIsInitialScrollPendingOnThreadChangeState';
import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
@@ -57,7 +57,10 @@ export const AgentChatStreamSubscriptionEffect = () => {
{ threadId: currentAIChatThread },
);
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const agentChatIsStreaming = useAtomComponentFamilyStateValue(
agentChatIsStreamingComponentFamilyState,
{ threadId: currentAIChatThread },
);
const agentChatDisplayedThread = useAtomStateValue(
agentChatDisplayedThreadState,
packages/twenty-front/src/modules/ai/components/AgentChatThreadInitializationEffect.tsx
+20 -12
View File
@@ -9,15 +9,16 @@ import {
import { agentChatInputState } from '@/ai/states/agentChatInputState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
import { agentChatThreadsSelector } from '@/ai/states/agentChatThreadsSelector';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { agentChatUsageComponentFamilyState } from '@/ai/states/agentChatUsageComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { currentAIChatThreadTitleComponentFamilyState } from '@/ai/states/currentAIChatThreadTitleComponentFamilyState';
import { hasInitializedAgentChatThreadsState } from '@/ai/states/hasInitializedAgentChatThreadsState';
import { hasTriggeredCreateForDraftState } from '@/ai/states/hasTriggeredCreateForDraftState';
import { useUpdateMetadataStoreDraft } from '@/metadata-store/hooks/useUpdateMetadataStoreDraft';
import { metadataStoreState } from '@/metadata-store/states/metadataStoreState';
import { type FlatAgentChatThread } from '@/metadata-store/types/FlatAgentChatThread';
import { useHasPermissionFlag } from '@/settings/roles/hooks/useHasPermissionFlag';
import { useAtomComponentFamilyStateCallbackState } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateCallbackState';
import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
@@ -37,13 +38,15 @@ export const AgentChatThreadInitializationEffect = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
const setAgentChatInput = useSetAtomState(agentChatInputState);
const setAgentChatUsage = useSetAtomState(agentChatUsageState);
const setCurrentAIChatThreadTitle = useSetAtomState(
currentAIChatThreadTitleState,
);
const setAgentChatThreadsLoading = useSetAtomState(
agentChatThreadsLoadingState,
);
const threadTitleFamilyCallback = useAtomComponentFamilyStateCallbackState(
currentAIChatThreadTitleComponentFamilyState,
);
const agentChatUsageFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatUsageComponentFamilyState,
);
const store = useStore();
const agentChatThreads = useAtomStateValue(agentChatThreadsSelector);
const storeEntry = useAtomValue(
@@ -113,13 +116,20 @@ export const AgentChatThreadInitializationEffect = () => {
setCurrentAIChatThread(firstThread.id);
setAgentChatInput(draftForThread);
setCurrentAIChatThreadTitle(firstThread.title ?? null);
const firstThreadFamilyKey = { threadId: firstThread.id };
store.set(
threadTitleFamilyCallback(firstThreadFamilyKey),
firstThread.title ?? null,
);
const hasUsageData =
(firstThread.conversationSize ?? 0) > 0 &&
isDefined(firstThread.contextWindowTokens);
setAgentChatUsage(
store.set(
agentChatUsageFamilyCallback(firstThreadFamilyKey),
hasUsageData
? {
lastMessage: null,
@@ -140,8 +150,6 @@ export const AgentChatThreadInitializationEffect = () => {
AGENT_CHAT_NEW_THREAD_DRAFT_KEY
] ?? '',
);
setCurrentAIChatThreadTitle(null);
setAgentChatUsage(null);
}
}, [
agentChatThreads,
@@ -152,9 +160,9 @@ export const AgentChatThreadInitializationEffect = () => {
storeEntry.status,
setCurrentAIChatThread,
setAgentChatInput,
setCurrentAIChatThreadTitle,
setAgentChatUsage,
store,
threadTitleFamilyCallback,
agentChatUsageFamilyCallback,
]);
return null;
packages/twenty-front/src/modules/ai/components/internal/AIChatContextUsageButton.tsx
+9 -3
View File
@@ -10,11 +10,13 @@ import { themeCssVariables } from 'twenty-ui/theme-constants';
import { ContextUsageProgressRing } from '@/ai/components/internal/ContextUsageProgressRing';
import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
import {
agentChatUsageState,
agentChatUsageComponentFamilyState,
type AgentChatLastMessageUsage,
} from '@/ai/states/agentChatUsageState';
} from '@/ai/states/agentChatUsageComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { SettingsBillingLabelValueItem } from '@/settings/billing/components/internal/SettingsBillingLabelValueItem';
import { billingState } from '@/client-config/states/billingState';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { formatNumber } from '~/utils/format/formatNumber';
@@ -94,7 +96,11 @@ const getCachedLabel = (lastMessage: AgentChatLastMessageUsage): string => {
export const AIChatContextUsageButton = () => {
const { t } = useLingui();
const [isHovered, setIsHovered] = useState(false);
const agentChatUsage = useAtomStateValue(agentChatUsageState);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatUsage = useAtomComponentFamilyStateValue(
agentChatUsageComponentFamilyState,
{ threadId: currentAIChatThread },
);
const billing = useAtomStateValue(billingState);
const isBillingEnabled = billing?.isBillingEnabled ?? false;
packages/twenty-front/src/modules/ai/components/internal/SendMessageButton.tsx
+8 -2
View File
@@ -1,8 +1,10 @@
import { AGENT_CHAT_STOP_EVENT_NAME } from '@/ai/constants/AgentChatStopEventName';
import { agentChatInputIsEmptySelector } from '@/ai/states/agentChatInputIsEmptySelector';
import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { IconArrowUp, IconPlayerStop } from 'twenty-ui/display';
import { RoundedIconButton } from 'twenty-ui/input';
@@ -18,7 +20,11 @@ export const SendMessageButton = ({ onSend }: SendMessageButtonProps) => {
const agentChatIsLoading = useAtomStateValue(agentChatIsLoadingState);
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatIsStreaming = useAtomComponentFamilyStateValue(
agentChatIsStreamingComponentFamilyState,
{ threadId: currentAIChatThread },
);
const handleStopClick = () => {
dispatchBrowserEvent(AGENT_CHAT_STOP_EVENT_NAME);
packages/twenty-front/src/modules/ai/hooks/useAIChatThreadClick.ts
+17 -8
View File
@@ -1,10 +1,11 @@
import { agentChatDraftsByThreadIdState } from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatInputState } from '@/ai/states/agentChatInputState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { agentChatUsageComponentFamilyState } from '@/ai/states/agentChatUsageComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { currentAIChatThreadTitleComponentFamilyState } from '@/ai/states/currentAIChatThreadTitleComponentFamilyState';
import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
import { useOpenAskAIPageInSidePanel } from '@/side-panel/hooks/useOpenAskAIPageInSidePanel';
import { useAtomComponentFamilyStateCallbackState } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateCallbackState';
import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { useStore } from 'jotai';
@@ -26,13 +27,15 @@ export const useAIChatThreadClick = (
currentAIChatThreadState,
);
const setAgentChatInput = useSetAtomState(agentChatInputState);
const setCurrentAIChatThreadTitle = useSetAtomState(
currentAIChatThreadTitleState,
);
const setAgentChatUsage = useSetAtomState(agentChatUsageState);
const setAgentChatDraftsByThreadId = useSetAtomState(
agentChatDraftsByThreadIdState,
);
const threadTitleFamilyCallback = useAtomComponentFamilyStateCallbackState(
currentAIChatThreadTitleComponentFamilyState,
);
const agentChatUsageFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatUsageComponentFamilyState,
);
const store = useStore();
const { openAskAIPage } = useOpenAskAIPageInSidePanel();
@@ -54,12 +57,18 @@ export const useAIChatThreadClick = (
setAgentChatInput(newDraft);
}
setCurrentAIChatThreadTitle(thread.title ?? null);
const clickedFamilyKey = { threadId: thread.id };
store.set(
threadTitleFamilyCallback(clickedFamilyKey),
thread.title ?? null,
);
const hasUsageData =
(thread.conversationSize ?? 0) > 0 &&
isDefined(thread.contextWindowTokens);
setAgentChatUsage(
store.set(
agentChatUsageFamilyCallback(clickedFamilyKey),
hasUsageData
? {
lastMessage: null,
packages/twenty-front/src/modules/ai/hooks/useAgentChatSubscription.ts
+72 -42
View File
@@ -10,17 +10,17 @@ import {
import { isDefined } from 'twenty-shared/utils';
import { v4 } from 'uuid';
import { AGENT_CHAT_INSTANCE_ID } from '@/ai/constants/AgentChatInstanceId';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { ON_AGENT_CHAT_EVENT } from '@/ai/graphql/subscriptions/OnAgentChatEvent';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatFirstLiveSeqState } from '@/ai/states/agentChatFirstLiveSeqState';
import { agentChatHandleEventCallbackState } from '@/ai/states/agentChatHandleEventCallbackState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatErrorComponentFamilyState } from '@/ai/states/agentChatErrorComponentFamilyState';
import { agentChatFirstLiveSeqComponentFamilyState } from '@/ai/states/agentChatFirstLiveSeqComponentFamilyState';
import { agentChatHandleEventCallbackComponentFamilyState } from '@/ai/states/agentChatHandleEventCallbackComponentFamilyState';
import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { agentChatUsageComponentFamilyState } from '@/ai/states/agentChatUsageComponentFamilyState';
import { currentAIChatThreadTitleComponentFamilyState } from '@/ai/states/currentAIChatThreadTitleComponentFamilyState';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomComponentFamilyStateCallbackState } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateCallbackState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { sseClientState } from '@/sse-db-event/states/sseClientState';
@@ -96,18 +96,52 @@ export const useAgentChatSubscription = (threadId: string | null) => {
const store = useStore();
const sseClient = useAtomStateValue(sseClientState);
const errorFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatErrorComponentFamilyState,
);
const isStreamingFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatIsStreamingComponentFamilyState,
);
const firstLiveSeqFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatFirstLiveSeqComponentFamilyState,
);
const handleEventCallbackFamilyCallback =
useAtomComponentFamilyStateCallbackState(
agentChatHandleEventCallbackComponentFamilyState,
);
const messagesFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatMessagesComponentFamilyState,
);
const usageFamilyCallback = useAtomComponentFamilyStateCallbackState(
agentChatUsageComponentFamilyState,
);
const threadTitleFamilyCallback = useAtomComponentFamilyStateCallbackState(
currentAIChatThreadTitleComponentFamilyState,
);
useEffect(() => {
if (!isDefined(threadId) || !isDefined(sseClient)) {
return;
}
const familyKey = { threadId };
const errorAtom = errorFamilyCallback(familyKey);
const isStreamingAtom = isStreamingFamilyCallback(familyKey);
const firstLiveSeqAtom = firstLiveSeqFamilyCallback(familyKey);
const handleEventCallbackAtom =
handleEventCallbackFamilyCallback(familyKey);
const messagesAtom = messagesFamilyCallback(familyKey);
const usageAtom = usageFamilyCallback(familyKey);
const threadTitleAtom = threadTitleFamilyCallback(familyKey);
let bridge: TransformStream<UIMessageChunk> | null = null;
let throttleTimer: ReturnType<typeof setTimeout> | null = null;
let latestMessage: ExtendedUIMessage | null = null;
let writer: WritableStreamDefaultWriter<UIMessageChunk> | null = null;
let disposed = false;
store.set(agentChatFirstLiveSeqState.atom, null);
store.set(firstLiveSeqAtom, null);
const closeWriter = () => {
if (isDefined(writer)) {
@@ -119,8 +153,8 @@ export const useAgentChatSubscription = (threadId: string | null) => {
const cleanupStream = () => {
closeWriter();
if (store.get(agentChatIsStreamingState.atom)) {
store.set(agentChatIsStreamingState.atom, false);
if (store.get(isStreamingAtom)) {
store.set(isStreamingAtom, false);
}
};
@@ -131,14 +165,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {
return;
}
const atomKey = {
instanceId: AGENT_CHAT_INSTANCE_ID,
familyKey: { threadId },
};
const currentMessages = store.get(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
);
const currentMessages = store.get(messagesAtom);
const streamingMsgIndex = currentMessages.findIndex(
(message) => message.id === messageToFlush.id,
@@ -148,15 +175,9 @@ export const useAgentChatSubscription = (threadId: string | null) => {
const updatedMessages = [...currentMessages];
updatedMessages[streamingMsgIndex] = messageToFlush;
store.set(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
updatedMessages,
);
store.set(messagesAtom, updatedMessages);
} else {
store.set(agentChatMessagesComponentFamilyState.atomFamily(atomKey), [
...currentMessages,
messageToFlush,
]);
store.set(messagesAtom, [...currentMessages, messageToFlush]);
}
};
@@ -184,7 +205,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {
);
if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
store.set(currentAIChatThreadTitleState.atom, titlePart.data.title);
store.set(threadTitleAtom, titlePart.data.title);
}
const metadata = extendedMessage.metadata as
@@ -207,7 +228,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {
const usage = metadata.usage;
const model = metadata.model;
store.set(agentChatUsageState.atom, (prev) => ({
store.set(usageAtom, (prev) => ({
lastMessage: {
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
@@ -234,22 +255,20 @@ export const useAgentChatSubscription = (threadId: string | null) => {
flushToAtom();
if (!disposed) {
store.set(agentChatIsStreamingState.atom, false);
store.set(isStreamingAtom, false);
}
};
const handleEvent = (event: AgentChatSubscriptionEvent) => {
switch (event.type) {
case 'stream-chunk': {
if (
isDefined(event.seq) &&
store.get(agentChatFirstLiveSeqState.atom) === null
) {
store.set(agentChatFirstLiveSeqState.atom, event.seq);
if (isDefined(event.seq) && store.get(firstLiveSeqAtom) === null) {
store.set(firstLiveSeqAtom, event.seq);
}
if (!store.get(agentChatIsStreamingState.atom)) {
store.set(agentChatIsStreamingState.atom, true);
if (!store.get(isStreamingAtom)) {
store.set(isStreamingAtom, true);
store.set(errorAtom, null);
bridge = new TransformStream<UIMessageChunk>();
writer = bridge.writable.getWriter();
@@ -260,7 +279,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {
startReadLoop(adaptedReadable).catch(() => {
if (!disposed) {
store.set(agentChatIsStreamingState.atom, false);
store.set(isStreamingAtom, false);
}
});
}
@@ -288,16 +307,16 @@ export const useAgentChatSubscription = (threadId: string | null) => {
};
streamError.code = event.code;
store.set(agentChatErrorState.atom, streamError);
store.set(errorAtom, streamError);
closeWriter();
store.set(agentChatIsStreamingState.atom, false);
store.set(isStreamingAtom, false);
break;
}
}
};
store.set(agentChatHandleEventCallbackState.atom, () => handleEvent);
store.set(handleEventCallbackAtom, () => handleEvent);
const dispose = sseClient.subscribe<AgentChatEventPayload>(
{
@@ -325,12 +344,23 @@ export const useAgentChatSubscription = (threadId: string | null) => {
return () => {
disposed = true;
store.set(agentChatHandleEventCallbackState.atom, null);
store.set(handleEventCallbackAtom, null);
if (isDefined(throttleTimer)) {
clearTimeout(throttleTimer);
}
cleanupStream();
dispose();
};
}, [threadId, sseClient, store]);
}, [
threadId,
sseClient,
store,
errorFamilyCallback,
isStreamingFamilyCallback,
firstLiveSeqFamilyCallback,
handleEventCallbackFamilyCallback,
messagesFamilyCallback,
usageFamilyCallback,
threadTitleFamilyCallback,
]);
};
packages/twenty-front/src/modules/ai/hooks/useCreateAgentChatThread.ts
-8
View File
@@ -5,9 +5,7 @@ import {
agentChatDraftsByThreadIdState,
} from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatInputState } from '@/ai/states/agentChatInputState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { shouldFocusChatEditorState } from '@/ai/states/shouldFocusChatEditorState';
import { hasTriggeredCreateForDraftState } from '@/ai/states/hasTriggeredCreateForDraftState';
import { isCreatingChatThreadState } from '@/ai/states/isCreatingChatThreadState';
@@ -24,10 +22,6 @@ import { CreateChatThreadDocument } from '~/generated-metadata/graphql';
export const useCreateAgentChatThread = () => {
const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
const setAgentChatInput = useSetAtomState(agentChatInputState);
const setAgentChatUsage = useSetAtomState(agentChatUsageState);
const setCurrentAIChatThreadTitle = useSetAtomState(
currentAIChatThreadTitleState,
);
const setIsCreatingChatThread = useSetAtomState(isCreatingChatThreadState);
const setAgentChatDraftsByThreadId = useSetAtomState(
agentChatDraftsByThreadIdState,
@@ -87,8 +81,6 @@ export const useCreateAgentChatThread = () => {
setCurrentAIChatThread(newThreadId);
setAgentChatInput(newDraft);
setCurrentAIChatThreadTitle(null);
setAgentChatUsage(null);
},
onError: () => {
setIsCreatingChatThread(false);
packages/twenty-front/src/modules/ai/hooks/useSwitchToNewAIChat.ts
-8
View File
@@ -5,9 +5,7 @@ import {
agentChatDraftsByThreadIdState,
} from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatInputState } from '@/ai/states/agentChatInputState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { shouldFocusChatEditorState } from '@/ai/states/shouldFocusChatEditorState';
import { hasTriggeredCreateForDraftState } from '@/ai/states/hasTriggeredCreateForDraftState';
import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
@@ -23,10 +21,6 @@ export const useSwitchToNewAIChat = () => {
currentAIChatThreadState,
);
const setAgentChatInput = useSetAtomState(agentChatInputState);
const setAgentChatUsage = useSetAtomState(agentChatUsageState);
const setCurrentAIChatThreadTitle = useSetAtomState(
currentAIChatThreadTitleState,
);
const setAgentChatDraftsByThreadId = useSetAtomState(
agentChatDraftsByThreadIdState,
);
@@ -48,8 +42,6 @@ export const useSwitchToNewAIChat = () => {
store.set(hasTriggeredCreateForDraftState.atom, false);
setCurrentAIChatThread(AGENT_CHAT_NEW_THREAD_DRAFT_KEY);
setAgentChatInput(newChatDraft);
setCurrentAIChatThreadTitle(null);
setAgentChatUsage(null);
openAskAIPage();
store.set(shouldFocusChatEditorState.atom, true);
};
packages/twenty-front/src/modules/ai/states/agentChatErrorComponentFamilyState.ts
+9
View File
@@ -0,0 +1,9 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
export const agentChatErrorComponentFamilyState =
createAtomComponentFamilyState<Error | null, { threadId: string | null }>({
key: 'agentChatErrorComponentFamilyState',
defaultValue: null,
componentInstanceContext: AgentChatComponentInstanceContext,
});
packages/twenty-front/src/modules/ai/states/agentChatErrorState.ts
-6
View File
@@ -1,6 +0,0 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatErrorState = createAtomState<Error | undefined | null>({
key: 'agentChatErrorState',
defaultValue: null,
});
packages/twenty-front/src/modules/ai/states/agentChatFirstLiveSeqComponentFamilyState.ts
+9
View File
@@ -0,0 +1,9 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
export const agentChatFirstLiveSeqComponentFamilyState =
createAtomComponentFamilyState<number | null, { threadId: string | null }>({
key: 'agentChatFirstLiveSeqComponentFamilyState',
defaultValue: null,
componentInstanceContext: AgentChatComponentInstanceContext,
});
packages/twenty-front/src/modules/ai/states/agentChatFirstLiveSeqState.ts
-6
View File
@@ -1,6 +0,0 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatFirstLiveSeqState = createAtomState<number | null>({
key: 'agentChatFirstLiveSeqState',
defaultValue: null,
});
packages/twenty-front/src/modules/ai/states/agentChatHandleEventCallbackComponentFamilyState.ts
+14
View File
@@ -0,0 +1,14 @@
import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
export const agentChatHandleEventCallbackComponentFamilyState =
createAtomComponentFamilyState<
((event: AgentChatSubscriptionEvent) => void) | null,
{ threadId: string | null }
>({
key: 'agentChatHandleEventCallbackComponentFamilyState',
defaultValue: null,
componentInstanceContext: AgentChatComponentInstanceContext,
});
packages/twenty-front/src/modules/ai/states/agentChatHandleEventCallbackState.ts
-10
View File
@@ -1,10 +0,0 @@
import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatHandleEventCallbackState = createAtomState<
((event: AgentChatSubscriptionEvent) => void) | null
>({
key: 'agentChatHandleEventCallbackState',
defaultValue: null,
});
packages/twenty-front/src/modules/ai/states/agentChatIsStreamingComponentFamilyState.ts
+9
View File
@@ -0,0 +1,9 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
export const agentChatIsStreamingComponentFamilyState =
createAtomComponentFamilyState<boolean, { threadId: string | null }>({
key: 'agentChatIsStreamingComponentFamilyState',
defaultValue: false,
componentInstanceContext: AgentChatComponentInstanceContext,
});
packages/twenty-front/src/modules/ai/states/agentChatIsStreamingState.ts
-6
View File
@@ -1,6 +0,0 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatIsStreamingState = createAtomState({
key: 'agentChatIsStreamingState',
defaultValue: false,
});
packages/twenty-front/src/modules/ai/states/agentChatUsageComponentFamilyState.ts
+30
View File
@@ -0,0 +1,30 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
export type AgentChatLastMessageUsage = {
inputTokens: number;
outputTokens: number;
cachedInputTokens: number;
inputCredits: number;
outputCredits: number;
};
export type AgentChatUsageState = {
lastMessage: AgentChatLastMessageUsage | null;
conversationSize: number;
contextWindowTokens: number;
inputTokens: number;
outputTokens: number;
inputCredits: number;
outputCredits: number;
};
export const agentChatUsageComponentFamilyState =
createAtomComponentFamilyState<
AgentChatUsageState | null,
{ threadId: string | null }
>({
key: 'agentChatUsageComponentFamilyState',
defaultValue: null,
componentInstanceContext: AgentChatComponentInstanceContext,
});
packages/twenty-front/src/modules/ai/states/agentChatUsageState.ts
-24
View File
@@ -1,24 +0,0 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export type AgentChatLastMessageUsage = {
inputTokens: number;
outputTokens: number;
cachedInputTokens: number;
inputCredits: number;
outputCredits: number;
};
export type AgentChatUsageState = {
lastMessage: AgentChatLastMessageUsage | null;
conversationSize: number;
contextWindowTokens: number;
inputTokens: number;
outputTokens: number;
inputCredits: number;
outputCredits: number;
};
export const agentChatUsageState = createAtomState<AgentChatUsageState | null>({
key: 'agentChatUsageState',
defaultValue: null,
});
packages/twenty-front/src/modules/ai/states/currentAIChatThreadTitleComponentFamilyState.ts
+9
View File
@@ -0,0 +1,9 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
export const currentAIChatThreadTitleComponentFamilyState =
createAtomComponentFamilyState<string | null, { threadId: string | null }>({
key: 'currentAIChatThreadTitleComponentFamilyState',
defaultValue: null,
componentInstanceContext: AgentChatComponentInstanceContext,
});
packages/twenty-front/src/modules/ai/states/currentAIChatThreadTitleState.ts
-6
View File
@@ -1,6 +0,0 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const currentAIChatThreadTitleState = createAtomState<string | null>({
key: 'ai/currentAIChatThreadTitleState',
defaultValue: null,
});
packages/twenty-front/src/modules/side-panel/components/SidePanelAskAIInfo.tsx
+7 -3
View File
@@ -1,4 +1,6 @@
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleComponentFamilyState } from '@/ai/states/currentAIChatThreadTitleComponentFamilyState';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { styled } from '@linaria/react';
import { t } from '@lingui/core/macro';
@@ -12,8 +14,10 @@ const StyledPageTitle = styled.div`
`;
export const SidePanelAskAIInfo = () => {
const currentAIChatThreadTitle = useAtomStateValue(
currentAIChatThreadTitleState,
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const currentAIChatThreadTitle = useAtomComponentFamilyStateValue(
currentAIChatThreadTitleComponentFamilyState,
{ threadId: currentAIChatThread },
);
return (