Add draft message persistence for AI chat threads (#18371)

2026-03-10 06:12:02 +05:30
parent c433b2b73f
commit ae122f4bb1
27 changed files with 733 additions and 304 deletions
@@ -61,8 +61,8 @@ const jestConfig = {
  extensionsToTreatAsEsm: ['.ts', '.tsx'],
  coverageThreshold: {
    global: {
-      statements: 49.3,
-      lines: 47.9,
+      statements: 49.1,
+      lines: 47.7,
      functions: 39.5,
    },
  },
@@ -0,0 +1,152 @@
+import { styled } from '@linaria/react';
+import { EditorContent } from '@tiptap/react';
+import { LightButton } from 'twenty-ui/input';
+import { themeCssVariables } from 'twenty-ui/theme-constants';
+
+import { AIChatEmptyState } from '@/ai/components/AIChatEmptyState';
+import { AIChatStandaloneError } from '@/ai/components/AIChatStandaloneError';
+import { AgentChatContextPreview } from '@/ai/components/internal/AgentChatContextPreview';
+import { AgentChatFileUploadButton } from '@/ai/components/internal/AgentChatFileUploadButton';
+import { AIChatContextUsageButton } from '@/ai/components/internal/AIChatContextUsageButton';
+import { AIChatEditorFocusEffect } from '@/ai/components/internal/AIChatEditorFocusEffect';
+import { AIChatSkeletonLoader } from '@/ai/components/internal/AIChatSkeletonLoader';
+import { SendMessageButton } from '@/ai/components/internal/SendMessageButton';
+import { useAIChatEditor } from '@/ai/hooks/useAIChatEditor';
+import { useAiModelLabel } from '@/ai/hooks/useAiModelOptions';
+import { currentWorkspaceState } from '@/auth/states/currentWorkspaceState';
+import { useIsMobile } from '@/ui/utilities/responsive/hooks/useIsMobile';
+import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
+
+const StyledInputArea = styled.div<{ isMobile: boolean }>`
+  align-items: flex-end;
+  background: ${themeCssVariables.background.primary};
+  display: flex;
+  flex-direction: column;
+  flex-shrink: 0;
+  gap: ${themeCssVariables.spacing[2]};
+  padding-block: ${({ isMobile }) =>
+    isMobile ? '0' : themeCssVariables.spacing[3]};
+  padding-inline: ${themeCssVariables.spacing[3]};
+`;
+
+const StyledInputBox = styled.div`
+  background-color: ${themeCssVariables.background.transparent.lighter};
+  border: 1px solid ${themeCssVariables.border.color.medium};
+  border-radius: ${themeCssVariables.border.radius.sm};
+  box-sizing: border-box;
+  display: flex;
+  flex-direction: column;
+  gap: ${themeCssVariables.spacing[2]};
+  min-height: 140px;
+  padding: ${themeCssVariables.spacing[2]};
+  width: 100%;
+
+  &:focus-within {
+    border-color: ${themeCssVariables.color.blue};
+    box-shadow: 0px 0px 0px 3px ${themeCssVariables.color.transparent.blue2};
+  }
+`;
+
+const StyledEditorWrapper = styled.div`
+  display: flex;
+  flex: 1;
+  flex-direction: column;
+  min-height: 0;
+
+  .tiptap {
+    background: transparent;
+    border: none;
+    box-shadow: none;
+    color: ${themeCssVariables.font.color.primary};
+    font-family: inherit;
+    font-size: ${themeCssVariables.font.size.md};
+    font-weight: ${themeCssVariables.font.weight.regular};
+    line-height: 16px;
+    max-height: 320px;
+    min-height: 48px;
+    outline: none;
+    overflow-y: auto;
+    padding: 0;
+
+    p {
+      margin: 0;
+    }
+
+    p.is-editor-empty:first-of-type::before {
+      color: ${themeCssVariables.font.color.light};
+      content: attr(data-placeholder);
+      float: left;
+      font-weight: ${themeCssVariables.font.weight.regular};
+      height: 0;
+      pointer-events: none;
+    }
+  }
+`;
+
+const StyledButtonsContainer = styled.div`
+  align-items: center;
+  display: flex;
+  justify-content: space-between;
+  width: 100%;
+`;
+
+const StyledLeftButtonsContainer = styled.div`
+  align-items: center;
+  display: flex;
+  gap: ${themeCssVariables.spacing['0.5']};
+`;
+
+const StyledRightButtonsContainer = styled.div`
+  align-items: center;
+  display: flex;
+  gap: ${themeCssVariables.spacing[1]};
+`;
+
+const StyledReadOnlyModelButtonContainer = styled.div`
+  > * {
+    cursor: default;
+
+    &:hover,
+    &:active {
+      background: transparent;
+    }
+  }
+`;
+
+export const AIChatEditorSection = () => {
+  const isMobile = useIsMobile();
+  const currentWorkspace = useAtomStateValue(currentWorkspaceState);
+  const smartModelLabel = useAiModelLabel(currentWorkspace?.smartModel, false);
+
+  const { editor, handleSendAndClear } = useAIChatEditor();
+
+  return (
+    <>
+      <AIChatEditorFocusEffect editor={editor} />
+      <AIChatEmptyState editor={editor} />
+      <AIChatStandaloneError />
+      <AIChatSkeletonLoader />
+
+      <StyledInputArea isMobile={isMobile}>
+        <AgentChatContextPreview />
+        <StyledInputBox>
+          <StyledEditorWrapper>
+            <EditorContent editor={editor} />
+          </StyledEditorWrapper>
+          <StyledButtonsContainer>
+            <StyledLeftButtonsContainer>
+              <AgentChatFileUploadButton />
+              <AIChatContextUsageButton />
+            </StyledLeftButtonsContainer>
+            <StyledRightButtonsContainer>
+              <StyledReadOnlyModelButtonContainer>
+                <LightButton accent="tertiary" title={smartModelLabel} />
+              </StyledReadOnlyModelButtonContainer>
+              <SendMessageButton onSend={handleSendAndClear} />
+            </StyledRightButtonsContainer>
+          </StyledButtonsContainer>
+        </StyledInputBox>
+      </StyledInputArea>
+    </>
+  );
+};
@@ -1,13 +1,16 @@
 import { styled } from '@linaria/react';
 import { type Editor } from '@tiptap/react';
+import { isDefined } from 'twenty-shared/utils';

 import { AIChatSuggestedPrompts } from '@/ai/components/suggested-prompts/AIChatSuggestedPrompts';
+import { useAgentChatContext } from '@/ai/contexts/AgentChatContext';
+import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
 import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
 import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
-import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
+import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
+import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
 import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
 import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
-import { isDefined } from 'twenty-shared/utils';

 const StyledEmptyState = styled.div`
  display: flex;
@@ -22,16 +25,25 @@ type AIChatEmptyStateProps = {
 };

 export const AIChatEmptyState = ({ editor }: AIChatEmptyStateProps) => {
-  const agentChatIsLoading = useAtomStateValue(agentChatIsLoadingState);
-
  const agentChatError = useAtomStateValue(agentChatErrorState);
+  const { threadsLoading, messagesLoading } = useAgentChatContext();
+  const skipMessagesSkeletonUntilLoaded = useAtomStateValue(
+    skipMessagesSkeletonUntilLoadedState,
+  );
+  const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);

  const hasMessages = useAtomComponentSelectorValue(
    agentChatHasMessageComponentSelector,
  );

+  const isOnNewChatSlot =
+    !isDefined(currentAIChatThread) ||
+    currentAIChatThread === AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+  const skeletonShowing =
+    (threadsLoading && isOnNewChatSlot) ||
+    (messagesLoading && !skipMessagesSkeletonUntilLoaded);
  const shouldRender =
-    !hasMessages && !isDefined(agentChatError) && !agentChatIsLoading;
+    !hasMessages && !isDefined(agentChatError) && !skeletonShowing;

  if (!shouldRender) {
    return null;
@@ -1,26 +1,17 @@
 import { styled } from '@linaria/react';
-import { EditorContent } from '@tiptap/react';
 import { useState } from 'react';
-import { LightButton } from 'twenty-ui/input';
 import { themeCssVariables } from 'twenty-ui/theme-constants';

 import { DropZone } from '@/activities/files/components/DropZone';
-import { AgentChatFileUploadButton } from '@/ai/components/internal/AgentChatFileUploadButton';
-import { useAiModelLabel } from '@/ai/hooks/useAiModelOptions';
-
-import { AIChatEmptyState } from '@/ai/components/AIChatEmptyState';
-import { AIChatStandaloneError } from '@/ai/components/AIChatStandaloneError';
-import { AIChatTabMessageList } from '@/ai/components/AIChatTabMessageList';
-import { AIChatContextUsageButton } from '@/ai/components/internal/AIChatContextUsageButton';
-import { AIChatSkeletonLoader } from '@/ai/components/internal/AIChatSkeletonLoader';
-import { AgentChatContextPreview } from '@/ai/components/internal/AgentChatContextPreview';
-import { SendMessageButton } from '@/ai/components/internal/SendMessageButton';
-import { useAIChatEditor } from '@/ai/hooks/useAIChatEditor';
+import { AIChatEditorSection } from '@/ai/components/AIChatEditorSection';
 import { useAIChatFileUpload } from '@/ai/hooks/useAIChatFileUpload';
-import { currentWorkspaceState } from '@/auth/states/currentWorkspaceState';
-import { useIsMobile } from '@/ui/utilities/responsive/hooks/useIsMobile';
+import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
+import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
+import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
 import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';

+import { AIChatTabMessageList } from '@/ai/components/AIChatTabMessageList';
+
 const StyledContainer = styled.div<{ isDraggingFile: boolean }>`
  background: ${themeCssVariables.background.primary};
  display: flex;
@@ -31,110 +22,20 @@ const StyledContainer = styled.div<{ isDraggingFile: boolean }>`
    isDraggingFile ? themeCssVariables.spacing[3] : '0'};
 `;

-const StyledInputArea = styled.div<{ isMobile: boolean }>`
-  align-items: flex-end;
-  background: ${themeCssVariables.background.primary};
-  display: flex;
-  flex-direction: column;
-  gap: ${themeCssVariables.spacing[2]};
-  padding-block: ${({ isMobile }) =>
-    isMobile ? '0' : themeCssVariables.spacing[3]};
-  padding-inline: ${themeCssVariables.spacing[3]};
-`;
-
-const StyledInputBox = styled.div`
-  background-color: ${themeCssVariables.background.transparent.lighter};
-  border: 1px solid ${themeCssVariables.border.color.medium};
-  border-radius: ${themeCssVariables.border.radius.sm};
-  box-sizing: border-box;
-  display: flex;
-  flex-direction: column;
-  gap: ${themeCssVariables.spacing[2]};
-  min-height: 140px;
-  padding: ${themeCssVariables.spacing[2]};
-  width: 100%;
-
-  &:focus-within {
-    border-color: ${themeCssVariables.color.blue};
-    box-shadow: 0px 0px 0px 3px ${themeCssVariables.color.transparent.blue2};
-  }
-`;
-
-const StyledEditorWrapper = styled.div`
-  display: flex;
-  flex: 1;
-  flex-direction: column;
-  min-height: 0;
-
-  .tiptap {
-    background: transparent;
-    border: none;
-    box-shadow: none;
-    color: ${themeCssVariables.font.color.primary};
-    font-family: inherit;
-    font-size: ${themeCssVariables.font.size.md};
-    font-weight: ${themeCssVariables.font.weight.regular};
-    line-height: 16px;
-    max-height: 320px;
-    min-height: 48px;
-    outline: none;
-    overflow-y: auto;
-    padding: 0;
-
-    p {
-      margin: 0;
-    }
-
-    p.is-editor-empty:first-of-type::before {
-      color: ${themeCssVariables.font.color.light};
-      content: attr(data-placeholder);
-      float: left;
-      font-weight: ${themeCssVariables.font.weight.regular};
-      height: 0;
-      pointer-events: none;
-    }
-  }
-`;
-
-const StyledButtonsContainer = styled.div`
-  align-items: center;
-  display: flex;
-  justify-content: space-between;
-  width: 100%;
-`;
-
-const StyledLeftButtonsContainer = styled.div`
-  align-items: center;
-  display: flex;
-  gap: ${themeCssVariables.spacing['0.5']};
-`;
-
-const StyledRightButtonsContainer = styled.div`
-  align-items: center;
-  display: flex;
-  gap: ${themeCssVariables.spacing[1]};
-`;
-
-const StyledReadOnlyModelButtonContainer = styled.div`
-  > * {
-    cursor: default;
-
-    &:hover,
-    &:active {
-      background: transparent;
-    }
-  }
-`;
-
 export const AIChatTab = () => {
  const [isDraggingFile, setIsDraggingFile] = useState(false);
-  const isMobile = useIsMobile();
+  const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
+  const threadIdCreatedFromDraft = useAtomStateValue(
+    threadIdCreatedFromDraftState,
+  );
+  const draftKey = currentAIChatThread ?? AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+  const editorSectionKey =
+    draftKey !== AGENT_CHAT_NEW_THREAD_DRAFT_KEY &&
+    draftKey === threadIdCreatedFromDraft
+      ? AGENT_CHAT_NEW_THREAD_DRAFT_KEY
+      : draftKey;

  const { uploadFiles } = useAIChatFileUpload();
-  const currentWorkspace = useAtomStateValue(currentWorkspaceState);
-  const smartModelLabel = useAiModelLabel(currentWorkspace?.smartModel, false);
-
-  const { editor, handleSendAndClear } = useAIChatEditor();

  return (
    <StyledContainer
@@ -150,29 +51,7 @@ export const AIChatTab = () => {
      {!isDraggingFile && (
        <>
          <AIChatTabMessageList />
-          <AIChatEmptyState editor={editor} />
-          <AIChatStandaloneError />
-          <AIChatSkeletonLoader />
-          <StyledInputArea isMobile={isMobile}>
-            <AgentChatContextPreview />
-            <StyledInputBox>
-              <StyledEditorWrapper>
-                <EditorContent editor={editor} />
-              </StyledEditorWrapper>
-              <StyledButtonsContainer>
-                <StyledLeftButtonsContainer>
-                  <AgentChatFileUploadButton />
-                  <AIChatContextUsageButton />
-                </StyledLeftButtonsContainer>
-                <StyledRightButtonsContainer>
-                  <StyledReadOnlyModelButtonContainer>
-                    <LightButton accent="tertiary" title={smartModelLabel} />
-                  </StyledReadOnlyModelButtonContainer>
-                  <SendMessageButton onSend={handleSendAndClear} />
-                </StyledRightButtonsContainer>
-              </StyledButtonsContainer>
-            </StyledInputBox>
-          </StyledInputArea>
+          <AIChatEditorSection key={editorSectionKey} />
        </>
      )}
    </StyledContainer>
@@ -36,15 +36,15 @@ const StyledButtonsContainer = styled.div`
 `;

 export const AIChatThreadsList = () => {
-  const { createChatThread } = useCreateNewAIChatThread();
+  const { switchToNewChat } = useCreateNewAIChatThread();

  const focusId = 'threads-list';

  useHotkeysOnFocusedElement({
    keys: [`${Key.Control}+${Key.Enter}`, `${Key.Meta}+${Key.Enter}`],
-    callback: () => createChatThread(),
+    callback: () => switchToNewChat(),
    focusId,
-    dependencies: [createChatThread],
+    dependencies: [switchToNewChat],
  });

  const { threads, hasNextPage, loading, fetchMoreRef } = useChatThreads();
@@ -77,7 +77,7 @@ export const AIChatThreadsList = () => {
            accent="blue"
            size="medium"
            title={t`New chat`}
-            onClick={() => createChatThread()}
+            onClick={() => switchToNewChat()}
            hotkeys={[getOsControlSymbol(), '⏎']}
          />
        </StyledButtonsContainer>
@@ -14,14 +14,13 @@ import { useEffect } from 'react';
 import { Temporal } from 'temporal-polyfill';

 export const AgentChatDataEffect = () => {
-  const { uiMessages, isLoading } = useAgentChatData();
-  const chatState = useAgentChat(uiMessages);
+  const { uiMessages, isLoading, ensureThreadIdForSend } = useAgentChatData();
+  const chatState = useAgentChat(uiMessages, ensureThreadIdForSend);

  const combinedIsLoading = chatState.isLoading || isLoading;
  const isStreaming = chatState.status === 'streaming';

  const setAgentChatIsLoading = useSetAtomState(agentChatIsLoadingState);
-
  const setAgentChatError = useSetAtomState(agentChatErrorState);

  const [agentChatUISessionStartTime, setAgentChatUISessionStartTime] =
@@ -1,8 +1,5 @@
-import { AgentChatDataEffect } from '@/ai/components/AgentChatDataEffect';
-import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
-
+import { AgentChatProviderContent } from '@/ai/components/AgentChatProviderContent';
 import { useIsFeatureEnabled } from '@/workspace/hooks/useIsFeatureEnabled';
-import { Suspense } from 'react';
 import { FeatureFlagKey } from '~/generated-metadata/graphql';

 export const AgentChatProvider = ({
@@ -16,14 +13,5 @@ export const AgentChatProvider = ({
    return <>{children}</>;
  }

-  return (
-    <Suspense fallback={null}>
-      <AgentChatComponentInstanceContext.Provider
-        value={{ instanceId: 'agentChatComponentInstance' }}
-      >
-        <AgentChatDataEffect />
-        {children}
-      </AgentChatComponentInstanceContext.Provider>
-    </Suspense>
-  );
+  return <AgentChatProviderContent>{children}</AgentChatProviderContent>;
 };
@@ -0,0 +1,33 @@
+import { AgentChatDataEffect } from '@/ai/components/AgentChatDataEffect';
+import { AgentChatContext } from '@/ai/contexts/AgentChatContext';
+import { useAgentChatData } from '@/ai/hooks/useAgentChatData';
+import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
+import { Suspense } from 'react';
+
+export const AgentChatProviderContent = ({
+  children,
+}: {
+  children: React.ReactNode;
+}) => {
+  const { ensureThreadForDraft, threadsLoading, messagesLoading } =
+    useAgentChatData();
+
+  const contextValue = {
+    ensureThreadForDraft,
+    threadsLoading,
+    messagesLoading,
+  };
+
+  return (
+    <Suspense fallback={null}>
+      <AgentChatContext.Provider value={contextValue}>
+        <AgentChatComponentInstanceContext.Provider
+          value={{ instanceId: 'agentChatComponentInstance' }}
+        >
+          <AgentChatDataEffect />
+          {children}
+        </AgentChatComponentInstanceContext.Provider>
+      </AgentChatContext.Provider>
+    </Suspense>
+  );
+};
@@ -0,0 +1,34 @@
+import { type Editor } from '@tiptap/react';
+import { useEffect } from 'react';
+
+import { focusEditorAfterMigrateState } from '@/ai/states/focusEditorAfterMigrateState';
+import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
+
+type AIChatEditorFocusEffectProps = {
+  editor: Editor | null;
+};
+
+export const AIChatEditorFocusEffect = ({
+  editor,
+}: AIChatEditorFocusEffectProps) => {
+  const [focusEditorAfterMigrate, setFocusEditorAfterMigrate] = useAtomState(
+    focusEditorAfterMigrateState,
+  );
+
+  useEffect(() => {
+    if (!focusEditorAfterMigrate || !editor) {
+      return;
+    }
+
+    const rafId = requestAnimationFrame(() => {
+      editor.commands.focus('end');
+      setFocusEditorAfterMigrate(false);
+    });
+
+    return () => {
+      cancelAnimationFrame(rafId);
+    };
+  }, [focusEditorAfterMigrate, editor, setFocusEditorAfterMigrate]);
+
+  return null;
+};
@@ -1,12 +1,17 @@
-import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
-import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
-import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
-import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
 import { styled } from '@linaria/react';
 import { useContext } from 'react';
 import Skeleton, { SkeletonTheme } from 'react-loading-skeleton';
+import { isDefined } from 'twenty-shared/utils';
 import { ThemeContext, themeCssVariables } from 'twenty-ui/theme-constants';

+import { useAgentChatContext } from '@/ai/contexts/AgentChatContext';
+import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
+import { agentChatHasMessageComponentSelector } from '@/ai/states/agentChatHasMessageComponentSelector';
+import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
+import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
+import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
+import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
+
 const StyledSkeletonContainer = styled.div`
  display: flex;
  flex: 1;
@@ -29,14 +34,24 @@ const NUMBER_OF_SKELETONS = 6;

 export const AIChatSkeletonLoader = () => {
  const { theme } = useContext(ThemeContext);
-
-  const agentChatIsLoading = useAtomStateValue(agentChatIsLoadingState);
+  const { threadsLoading, messagesLoading } = useAgentChatContext();
+  const skipMessagesSkeletonUntilLoaded = useAtomStateValue(
+    skipMessagesSkeletonUntilLoadedState,
+  );
+  const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);

  const hasMessages = useAtomComponentSelectorValue(
    agentChatHasMessageComponentSelector,
  );

-  const shouldRender = agentChatIsLoading && !hasMessages;
+  const isOnNewChatSlot =
+    !isDefined(currentAIChatThread) ||
+    currentAIChatThread === AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+  const showForMessagesLoading =
+    messagesLoading && !skipMessagesSkeletonUntilLoaded;
+  const shouldRender =
+    !hasMessages &&
+    ((threadsLoading && isOnNewChatSlot) || showForMessagesLoading);

  if (!shouldRender) {
    return null;
@@ -0,0 +1,15 @@
+import { createContext, useContext } from 'react';
+
+export type AgentChatContextValue = {
+  ensureThreadForDraft: (() => void) | undefined;
+  threadsLoading: boolean;
+  messagesLoading: boolean;
+};
+
+export const AgentChatContext = createContext<AgentChatContextValue>({
+  ensureThreadForDraft: undefined,
+  threadsLoading: false,
+  messagesLoading: false,
+});
+
+export const useAgentChatContext = () => useContext(AgentChatContext);
@@ -5,11 +5,17 @@ import { Paragraph } from '@tiptap/extension-paragraph';
 import { Text } from '@tiptap/extension-text';
 import { Placeholder } from '@tiptap/extensions/placeholder';
 import { useEditor } from '@tiptap/react';
-import { useCallback, useMemo } from 'react';
+import { useMemo } from 'react';
 import { isDefined } from 'twenty-shared/utils';

+import { useAgentChatContext } from '@/ai/contexts/AgentChatContext';
 import { AI_CHAT_INPUT_ID } from '@/ai/constants/AiChatInputId';
+import {
+  AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
+  agentChatDraftsByThreadIdState,
+} from '@/ai/states/agentChatDraftsByThreadIdState';
 import { agentChatInputState } from '@/ai/states/agentChatInputState';
+import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
 import { dispatchAgentChatSendMessageEvent } from '@/ai/utils/dispatchAgentChatSendMessageEvent';
 import { MENTION_SUGGESTION_PLUGIN_KEY } from '@/mention/constants/MentionSuggestionPluginKey';
 import { MentionSuggestion } from '@/mention/extensions/MentionSuggestion';
@@ -18,16 +24,37 @@ import { useMentionSearch } from '@/mention/hooks/useMentionSearch';
 import { usePushFocusItemToFocusStack } from '@/ui/utilities/focus/hooks/usePushFocusItemToFocusStack';
 import { useRemoveFocusItemFromFocusStackById } from '@/ui/utilities/focus/hooks/useRemoveFocusItemFromFocusStackById';
 import { FocusComponentType } from '@/ui/utilities/focus/types/FocusComponentType';
+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';
 import { turnIntoEmptyStringIfWhitespacesOnly } from '~/utils/string/turnIntoEmptyStringIfWhitespacesOnly';

+const textToTiptapContent = (text: string) => ({
+  type: 'doc',
+  content: [
+    {
+      type: 'paragraph',
+      content: text ? [{ type: 'text', text }] : [],
+    },
+  ],
+});
+
 export const useAIChatEditor = () => {
  const setAgentChatInput = useSetAtomState(agentChatInputState);
+  const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
+  const [agentChatDraftsByThreadId, setAgentChatDraftsByThreadId] =
+    useAtomState(agentChatDraftsByThreadIdState);
+  const { ensureThreadForDraft } = useAgentChatContext();
+
  const { searchMentionRecords } = useMentionSearch();
  const { pushFocusItemToFocusStack } = usePushFocusItemToFocusStack();
  const { removeFocusItemFromFocusStackById } =
    useRemoveFocusItemFromFocusStackById();

+  const draftKey = currentAIChatThread ?? AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+  const initialDraft = agentChatDraftsByThreadId[draftKey] ?? '';
+  const initialContent = textToTiptapContent(initialDraft);
+
  const extensions = useMemo(
    () => [
      Document,
@@ -46,6 +73,7 @@ export const useAIChatEditor = () => {
  );

  const editor = useEditor({
+    content: initialContent,
    extensions,
    editorProps: {
      handleKeyDown: (view, event) => {
@@ -72,6 +100,10 @@ export const useAIChatEditor = () => {
        currentEditor.getText({ blockSeparator: '\n' }),
      );
      setAgentChatInput(text);
+      setAgentChatDraftsByThreadId((prev) => ({ ...prev, [draftKey]: text }));
+      if (draftKey === AGENT_CHAT_NEW_THREAD_DRAFT_KEY && text.trim() !== '') {
+        ensureThreadForDraft?.();
+      }
    },
    onFocus: () => {
      pushFocusItemToFocusStack({
@@ -104,10 +136,10 @@ export const useAIChatEditor = () => {
    mentionStorage.searchMentionRecords = searchMentionRecords;
  }

-  const handleSendAndClear = useCallback(() => {
+  const handleSendAndClear = () => {
    dispatchAgentChatSendMessageEvent();
    editor?.commands.clearContent();
-  }, [editor]);
+  };

  return { editor, handleSendAndClear };
 };
@@ -1,9 +1,16 @@
+import {
+  AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
+  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 { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
 import { useOpenAskAIPageInSidePanel } from '@/side-panel/hooks/useOpenAskAIPageInSidePanel';
 import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
 import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
+import { useStore } from 'jotai';
 import { isDefined } from 'twenty-shared/utils';
 import { type AgentChatThread } from '~/generated-metadata/graphql';

@@ -15,21 +22,46 @@ export const useAIChatThreadClick = (
  options: UseAIChatThreadClickOptions = {},
 ) => {
  const { resetNavigationStack = false } = options;
-  const [, setCurrentAIChatThread] = useAtomState(currentAIChatThreadState);
+  const setThreadIdCreatedFromDraft = useSetAtomState(
+    threadIdCreatedFromDraftState,
+  );
+  const [currentAIChatThread, setCurrentAIChatThread] = useAtomState(
+    currentAIChatThreadState,
+  );
+  const setAgentChatInput = useSetAtomState(agentChatInputState);
  const setCurrentAIChatThreadTitle = useSetAtomState(
    currentAIChatThreadTitleState,
  );
  const setAgentChatUsage = useSetAtomState(agentChatUsageState);
+  const setAgentChatDraftsByThreadId = useSetAtomState(
+    agentChatDraftsByThreadIdState,
+  );
+  const store = useStore();
  const { openAskAIPage } = useOpenAskAIPageInSidePanel();

  const handleThreadClick = (thread: AgentChatThread) => {
+    setThreadIdCreatedFromDraft(null);
+    const previousDraftKey =
+      currentAIChatThread ?? AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+    const isSameThread = thread.id === currentAIChatThread;
+
+    setAgentChatDraftsByThreadId((prev) => ({
+      ...prev,
+      [previousDraftKey]: store.get(agentChatInputState.atom),
+    }));
    setCurrentAIChatThread(thread.id);
+
+    if (!isSameThread) {
+      const newDraft =
+        store.get(agentChatDraftsByThreadIdState.atom)[thread.id] ?? '';
+      setAgentChatInput(newDraft);
+    }
+
    setCurrentAIChatThreadTitle(thread.title ?? null);

    const hasUsageData =
      (thread.conversationSize ?? 0) > 0 &&
      isDefined(thread.contextWindowTokens);
-
    setAgentChatUsage(
      hasUsageData
        ? {
@@ -10,6 +10,10 @@ import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTi

 import { AGENT_CHAT_RETRY_EVENT_NAME } from '@/ai/constants/AgentChatRetryEventName';
 import { AGENT_CHAT_STOP_EVENT_NAME } from '@/ai/constants/AgentChatStopEventName';
+import {
+  AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
+  agentChatDraftsByThreadIdState,
+} from '@/ai/states/agentChatDraftsByThreadIdState';
 import { agentChatInputState } from '@/ai/states/agentChatInputState';
 import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
 import { getTokenPair } from '@/apollo/utils/getTokenPair';
@@ -21,13 +25,17 @@ import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomState
 import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
 import { useChat } from '@ai-sdk/react';
 import { DefaultChatTransport } from 'ai';
-import { useCallback } from 'react';
+import { useStore } from 'jotai';
+import { useCallback, useState } from 'react';
 import { type ExtendedUIMessage } from 'twenty-shared/ai';
 import { isDefined } from 'twenty-shared/utils';
 import { REACT_APP_SERVER_BASE_URL } from '~/config';
 import { cookieStorage } from '~/utils/cookie-storage';

-export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
+export const useAgentChat = (
+  uiMessages: ExtendedUIMessage[],
+  ensureThreadIdForSend: () => Promise<string | null>,
+) => {
  const setTokenPair = useSetAtomState(tokenPairState);
  const setAgentChatUsage = useSetAtomState(agentChatUsageState);

@@ -35,17 +43,24 @@ export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
  const setCurrentAIChatThreadTitle = useSetAtomState(
    currentAIChatThreadTitleState,
  );
+  const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
  const apolloClient = useApolloClient();
+  const store = useStore();

  const agentChatSelectedFiles = useAtomStateValue(agentChatSelectedFilesState);

  const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);

+  const [, setPendingThreadIdAfterFirstSend] = useState<string | null>(null);
+
  const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useAtomState(
    agentChatUploadedFilesState,
  );

-  const [agentChatInput, setAgentChatInput] = useAtomState(agentChatInputState);
+  const [, setAgentChatInput] = useAtomState(agentChatInputState);
+  const setAgentChatDraftsByThreadId = useSetAtomState(
+    agentChatDraftsByThreadIdState,
+  );

  const retryFetchWithRenewedToken = async (
    input: RequestInfo | URL,
@@ -167,23 +182,30 @@ export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
        (part) => part.type === 'data-thread-title',
      );

-      if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
-        setCurrentAIChatThreadTitle(titlePart.data.title);
-        if (isDefined(currentAIChatThread)) {
-          const threadRef = apolloClient.cache.identify({
-            __typename: 'AgentChatThread',
-            id: currentAIChatThread,
-          });
-          if (isDefined(threadRef)) {
-            apolloClient.cache.modify({
-              id: threadRef,
-              fields: {
-                title: () => titlePart.data.title,
-              },
+      setPendingThreadIdAfterFirstSend((pendingId) => {
+        const threadIdForTitle = pendingId ?? currentAIChatThread;
+        if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
+          setCurrentAIChatThreadTitle(titlePart.data.title);
+          if (isDefined(threadIdForTitle)) {
+            const threadRef = apolloClient.cache.identify({
+              __typename: 'AgentChatThread',
+              id: threadIdForTitle,
            });
+            if (isDefined(threadRef)) {
+              apolloClient.cache.modify({
+                id: threadRef,
+                fields: {
+                  title: () => titlePart.data.title,
+                },
+              });
+            }
          }
        }
-      }
+        if (isDefined(pendingId)) {
+          setCurrentAIChatThread(pendingId);
+        }
+        return null;
+      });
    },
  });

@@ -191,24 +213,47 @@ export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {
  const isLoading = isStreaming || agentChatSelectedFiles.length > 0;

  const handleSendMessage = useCallback(async () => {
-    if (agentChatInput.trim() === '' || isLoading || !currentAIChatThread) {
+    const draftKey =
+      store.get(currentAIChatThreadState.atom) ??
+      AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+    const contentToSend =
+      draftKey === AGENT_CHAT_NEW_THREAD_DRAFT_KEY
+        ? (
+            store.get(agentChatDraftsByThreadIdState.atom)[
+              AGENT_CHAT_NEW_THREAD_DRAFT_KEY
+            ] ?? store.get(agentChatInputState.atom)
+          ).trim()
+        : store.get(agentChatInputState.atom).trim();
+
+    if (contentToSend === '' || isLoading) {
      return;
    }

-    const content = agentChatInput.trim();
+    const threadId = await ensureThreadIdForSend();
+    if (!threadId) {
+      return;
+    }
+
+    if (draftKey === AGENT_CHAT_NEW_THREAD_DRAFT_KEY) {
+      setPendingThreadIdAfterFirstSend(threadId);
+    }

    setAgentChatInput('');
+    setAgentChatDraftsByThreadId((prev) => ({
+      ...prev,
+      [draftKey]: '',
+    }));

    const browsingContext = getBrowsingContext();

    sendMessage(
      {
-        text: content,
+        text: contentToSend,
        files: agentChatUploadedFiles,
      },
      {
        body: {
-          threadId: currentAIChatThread,
+          threadId,
          browsingContext,
        },
      },
@@ -216,14 +261,15 @@ export const useAgentChat = (uiMessages: ExtendedUIMessage[]) => {

    setAgentChatUploadedFiles([]);
  }, [
-    agentChatInput,
+    store,
    isLoading,
-    currentAIChatThread,
+    ensureThreadIdForSend,
    setAgentChatInput,
    getBrowsingContext,
    sendMessage,
    agentChatUploadedFiles,
    setAgentChatUploadedFiles,
+    setAgentChatDraftsByThreadId,
  ]);

  useListenToBrowserEvent({
@@ -1,75 +1,135 @@
+import { useApolloClient } from '@apollo/client';
 import { getOperationName } from '@apollo/client/utilities';
-import { type SetStateAction } from 'jotai';
+import { useStore } from 'jotai';
 import { isDefined } from 'twenty-shared/utils';

 import { CHAT_THREADS_PAGE_SIZE } from '@/ai/constants/ChatThreads';
 import { useAgentChatScrollToBottom } from '@/ai/hooks/useAgentChatScrollToBottom';
 import {
-  agentChatUsageState,
-  type AgentChatUsageState,
-} from '@/ai/states/agentChatUsageState';
+  AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
+  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 { focusEditorAfterMigrateState } from '@/ai/states/focusEditorAfterMigrateState';
+import { hasTriggeredCreateForDraftState } from '@/ai/states/hasTriggeredCreateForDraftState';
 import { isCreatingChatThreadState } from '@/ai/states/isCreatingChatThreadState';
+import { isCreatingForFirstSendState } from '@/ai/states/isCreatingForFirstSendState';
+import { pendingCreateFromDraftPromiseState } from '@/ai/states/pendingCreateFromDraftPromiseState';
+import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
+import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
 import { mapDBMessagesToUIMessages } from '@/ai/utils/mapDBMessagesToUIMessages';
 import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
 import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';

 import {
-  type AgentChatThread,
+  type GetChatThreadsQuery,
  GetChatThreadsDocument,
  useCreateChatThreadMutation,
  useGetChatMessagesQuery,
  useGetChatThreadsQuery,
 } from '~/generated-metadata/graphql';

-const setUsageFromThread = (
-  thread: AgentChatThread,
-  setAgentChatUsage: (
-    update: SetStateAction<AgentChatUsageState | null>,
-  ) => void,
-) => {
-  const hasUsageData =
-    (thread.conversationSize ?? 0) > 0 && isDefined(thread.contextWindowTokens);
-
-  setAgentChatUsage(
-    hasUsageData
-      ? {
-          lastMessage: null,
-          conversationSize: thread.conversationSize ?? 0,
-          contextWindowTokens: thread.contextWindowTokens ?? 0,
-          inputTokens: thread.totalInputTokens,
-          outputTokens: thread.totalOutputTokens,
-          inputCredits: thread.totalInputCredits,
-          outputCredits: thread.totalOutputCredits,
-        }
-      : null,
-  );
-};
-
 export const useAgentChatData = () => {
  const [currentAIChatThread, setCurrentAIChatThread] = useAtomState(
    currentAIChatThreadState,
  );
+  const setAgentChatInput = useSetAtomState(agentChatInputState);
  const setAgentChatUsage = useSetAtomState(agentChatUsageState);
  const setCurrentAIChatThreadTitle = useSetAtomState(
    currentAIChatThreadTitleState,
  );
-  const [isCreatingChatThread, setIsCreatingChatThread] = useAtomState(
-    isCreatingChatThreadState,
+  const [, setIsCreatingChatThread] = useAtomState(isCreatingChatThreadState);
+  const setAgentChatDraftsByThreadId = useSetAtomState(
+    agentChatDraftsByThreadIdState,
  );
+  const setPendingCreateFromDraftPromise = useSetAtomState(
+    pendingCreateFromDraftPromiseState,
+  );
+  const store = useStore();
+  const apolloClient = useApolloClient();

  const { scrollToBottom } = useAgentChatScrollToBottom();

  const [createChatThread] = useCreateChatThreadMutation({
    onCompleted: (data) => {
+      if (store.get(isCreatingForFirstSendState.atom)) {
+        store.set(isCreatingForFirstSendState.atom, false);
+        setIsCreatingChatThread(false);
+        return;
+      }
+
+      const newThreadId = data.createChatThread.id;
+      const previousDraftKey =
+        store.get(currentAIChatThreadState.atom) ??
+        AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+      const draftsSnapshot = store.get(agentChatDraftsByThreadIdState.atom);
+      const newDraft = draftsSnapshot[AGENT_CHAT_NEW_THREAD_DRAFT_KEY] ?? '';
+
      setIsCreatingChatThread(false);
-      setCurrentAIChatThread(data.createChatThread.id);
+      if (previousDraftKey === AGENT_CHAT_NEW_THREAD_DRAFT_KEY) {
+        store.set(hasTriggeredCreateForDraftState.atom, true);
+        setAgentChatDraftsByThreadId((prev) => ({
+          ...prev,
+          [newThreadId]: newDraft,
+          [AGENT_CHAT_NEW_THREAD_DRAFT_KEY]: '',
+        }));
+        store.set(focusEditorAfterMigrateState.atom, true);
+        store.set(skipMessagesSkeletonUntilLoadedState.atom, true);
+        store.set(threadIdCreatedFromDraftState.atom, newThreadId);
+      } else {
+        setAgentChatDraftsByThreadId((prev) => ({
+          ...prev,
+          [previousDraftKey]: store.get(agentChatInputState.atom),
+        }));
+      }
+      setCurrentAIChatThread(newThreadId);
+      setAgentChatInput(newDraft);
      setCurrentAIChatThreadTitle(null);
      setAgentChatUsage(null);
+
+      const newThread = data.createChatThread;
+      const threadListVariables = {
+        paging: { first: CHAT_THREADS_PAGE_SIZE },
+      };
+      const existing = apolloClient.cache.readQuery<GetChatThreadsQuery>({
+        query: GetChatThreadsDocument,
+        variables: threadListVariables,
+      });
+      if (isDefined(existing) && isDefined(existing.chatThreads)) {
+        const newNode = {
+          __typename: 'AgentChatThread' as const,
+          ...newThread,
+          totalInputTokens: 0,
+          totalOutputTokens: 0,
+          contextWindowTokens: null,
+          conversationSize: 0,
+          totalInputCredits: 0,
+          totalOutputCredits: 0,
+        };
+        const newEdge = {
+          __typename: 'AgentChatThreadEdge' as const,
+          node: newNode,
+          cursor: newThread.id,
+        };
+        apolloClient.cache.writeQuery({
+          query: GetChatThreadsDocument,
+          variables: threadListVariables,
+          data: {
+            chatThreads: {
+              ...existing.chatThreads,
+              edges: [newEdge, ...existing.chatThreads.edges],
+            },
+          },
+        });
+      }
    },
    onError: () => {
      setIsCreatingChatThread(false);
+      store.set(isCreatingForFirstSendState.atom, false);
+      store.set(hasTriggeredCreateForDraftState.atom, false);
    },
    refetchQueries: [
      getOperationName(GetChatThreadsDocument) ?? 'GetChatThreads',
@@ -80,32 +140,124 @@ export const useAgentChatData = () => {
    variables: { paging: { first: CHAT_THREADS_PAGE_SIZE } },
    skip: isDefined(currentAIChatThread),
    onCompleted: (data) => {
-      const edges = data?.chatThreads?.edges ?? [];
-      const threads = edges.map((edge) => edge.node);
+      const threads = data.chatThreads.edges.map((edge) => edge.node);
+
      if (threads.length > 0) {
        const firstThread = threads[0];
+        const newDraft =
+          store.get(agentChatDraftsByThreadIdState.atom)[firstThread.id] ?? '';

        setCurrentAIChatThread(firstThread.id);
+        setAgentChatInput(newDraft);
        setCurrentAIChatThreadTitle(firstThread.title ?? null);
-        setUsageFromThread(firstThread, setAgentChatUsage);
-      } else if (!isCreatingChatThread) {
-        setIsCreatingChatThread(true);
-        createChatThread();
+
+        const hasUsageData =
+          (firstThread.conversationSize ?? 0) > 0 &&
+          isDefined(firstThread.contextWindowTokens);
+        setAgentChatUsage(
+          hasUsageData
+            ? {
+                lastMessage: null,
+                conversationSize: firstThread.conversationSize ?? 0,
+                contextWindowTokens: firstThread.contextWindowTokens ?? 0,
+                inputTokens: firstThread.totalInputTokens,
+                outputTokens: firstThread.totalOutputTokens,
+                inputCredits: firstThread.totalInputCredits,
+                outputCredits: firstThread.totalOutputCredits,
+              }
+            : null,
+        );
+      } else {
+        store.set(hasTriggeredCreateForDraftState.atom, false);
+        setCurrentAIChatThread(AGENT_CHAT_NEW_THREAD_DRAFT_KEY);
+        setAgentChatInput(
+          store.get(agentChatDraftsByThreadIdState.atom)[
+            AGENT_CHAT_NEW_THREAD_DRAFT_KEY
+          ] ?? '',
+        );
+        setCurrentAIChatThreadTitle(null);
+        setAgentChatUsage(null);
      }
    },
  });

+  const isNewThread = currentAIChatThread === AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
  const { loading: messagesLoading, data } = useGetChatMessagesQuery({
    variables: { threadId: currentAIChatThread! },
-    skip: !isDefined(currentAIChatThread),
-    onCompleted: scrollToBottom,
+    skip: !isDefined(currentAIChatThread) || isNewThread,
+    onCompleted: () => {
+      store.set(skipMessagesSkeletonUntilLoadedState.atom, false);
+      scrollToBottom();
+    },
  });

+  const ensureThreadForDraft = () => {
+    const current = store.get(currentAIChatThreadState.atom);
+    if (current !== AGENT_CHAT_NEW_THREAD_DRAFT_KEY) {
+      return;
+    }
+    const draft =
+      store.get(agentChatDraftsByThreadIdState.atom)[
+        AGENT_CHAT_NEW_THREAD_DRAFT_KEY
+      ] ?? '';
+    if (draft.trim() === '') {
+      return;
+    }
+    if (store.get(hasTriggeredCreateForDraftState.atom)) {
+      return;
+    }
+    if (store.get(isCreatingChatThreadState.atom)) {
+      return;
+    }
+    setIsCreatingChatThread(true);
+    const createPromise = createChatThread();
+    const threadIdPromise = createPromise.then(
+      (result) => result?.data?.createChatThread?.id ?? null,
+    );
+    setPendingCreateFromDraftPromise(threadIdPromise);
+    threadIdPromise.finally(() => {
+      setPendingCreateFromDraftPromise(null);
+    });
+  };
+
+  const ensureThreadIdForSend = async (): Promise<string | null> => {
+    const current = store.get(currentAIChatThreadState.atom);
+    if (current !== AGENT_CHAT_NEW_THREAD_DRAFT_KEY) {
+      return current;
+    }
+    const inFlightCreate = store.get(pendingCreateFromDraftPromiseState.atom);
+    if (
+      store.get(isCreatingChatThreadState.atom) &&
+      isDefined(inFlightCreate)
+    ) {
+      try {
+        const threadId = await inFlightCreate;
+        return threadId;
+      } catch {
+        return null;
+      }
+    }
+    store.set(isCreatingForFirstSendState.atom, true);
+    setIsCreatingChatThread(true);
+    try {
+      const result = await createChatThread();
+      return result?.data?.createChatThread?.id ?? null;
+    } catch {
+      return null;
+    } finally {
+      setIsCreatingChatThread(false);
+    }
+  };
+
  const uiMessages = mapDBMessagesToUIMessages(data?.chatMessages || []);
  const isLoading = messagesLoading || threadsLoading;

  return {
    uiMessages,
    isLoading,
+    threadsLoading,
+    messagesLoading,
+    ensureThreadForDraft,
+    ensureThreadIdForSend,
  };
 };
@@ -1,74 +1,56 @@
-import { useApolloClient } from '@apollo/client';
+import { useStore } from 'jotai';

-import { CHAT_THREADS_PAGE_SIZE } from '@/ai/constants/ChatThreads';
+import {
+  AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
+  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 { hasTriggeredCreateForDraftState } from '@/ai/states/hasTriggeredCreateForDraftState';
+import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
 import { useOpenAskAIPageInSidePanel } from '@/side-panel/hooks/useOpenAskAIPageInSidePanel';
 import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
 import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
-import { isDefined } from 'twenty-shared/utils';
-
-import {
-  type GetChatThreadsQuery,
-  GetChatThreadsDocument,
-  useCreateChatThreadMutation,
-} from '~/generated-metadata/graphql';

 export const useCreateNewAIChatThread = () => {
-  const apolloClient = useApolloClient();
-  const [, setCurrentAIChatThread] = useAtomState(currentAIChatThreadState);
+  const setThreadIdCreatedFromDraft = useSetAtomState(
+    threadIdCreatedFromDraftState,
+  );
+  const [currentAIChatThread, setCurrentAIChatThread] = useAtomState(
+    currentAIChatThreadState,
+  );
+  const setAgentChatInput = useSetAtomState(agentChatInputState);
  const setAgentChatUsage = useSetAtomState(agentChatUsageState);
  const setCurrentAIChatThreadTitle = useSetAtomState(
    currentAIChatThreadTitleState,
  );
-
+  const setAgentChatDraftsByThreadId = useSetAtomState(
+    agentChatDraftsByThreadIdState,
+  );
+  const store = useStore();
  const { openAskAIPage } = useOpenAskAIPageInSidePanel();
-  const [createChatThread] = useCreateChatThreadMutation({
-    onCompleted: (data) => {
-      setCurrentAIChatThread(data.createChatThread.id);
-      setCurrentAIChatThreadTitle(null);
-      setAgentChatUsage(null);

-      openAskAIPage({ resetNavigationStack: false });
+  const switchToNewChat = () => {
+    setThreadIdCreatedFromDraft(null);
+    const previousDraftKey =
+      currentAIChatThread ?? AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
+    const newChatDraft =
+      store.get(agentChatDraftsByThreadIdState.atom)[
+        AGENT_CHAT_NEW_THREAD_DRAFT_KEY
+      ] ?? '';
+    setAgentChatDraftsByThreadId((prev) => ({
+      ...prev,
+      [previousDraftKey]: store.get(agentChatInputState.atom),
+    }));
+    store.set(hasTriggeredCreateForDraftState.atom, false);
+    setCurrentAIChatThread(AGENT_CHAT_NEW_THREAD_DRAFT_KEY);
+    setAgentChatInput(newChatDraft);
+    setCurrentAIChatThreadTitle(null);
+    setAgentChatUsage(null);
+    openAskAIPage({ resetNavigationStack: false });
+  };

-      const newThread = data.createChatThread;
-      const threadListVariables = {
-        paging: { first: CHAT_THREADS_PAGE_SIZE },
-      };
-      const existing = apolloClient.cache.readQuery<GetChatThreadsQuery>({
-        query: GetChatThreadsDocument,
-        variables: threadListVariables,
-      });
-      if (isDefined(existing) && isDefined(existing.chatThreads)) {
-        const newNode = {
-          __typename: 'AgentChatThread' as const,
-          ...newThread,
-          totalInputTokens: 0,
-          totalOutputTokens: 0,
-          contextWindowTokens: null,
-          conversationSize: 0,
-          totalInputCredits: 0,
-          totalOutputCredits: 0,
-        };
-        const newEdge = {
-          __typename: 'AgentChatThreadEdge' as const,
-          node: newNode,
-          cursor: newThread.id,
-        };
-        apolloClient.cache.writeQuery({
-          query: GetChatThreadsDocument,
-          variables: threadListVariables,
-          data: {
-            chatThreads: {
-              ...existing.chatThreads,
-              edges: [newEdge, ...existing.chatThreads.edges],
-            },
-          },
-        });
-      }
-    },
-  });
-
-  return { createChatThread };
+  return { switchToNewChat };
 };
@@ -0,0 +1,14 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const AGENT_CHAT_NEW_THREAD_DRAFT_KEY = '__new__';
+
+const DRAFTS_STORAGE_KEY = 'ai/agentChatDraftsByThreadIdState';
+
+export const agentChatDraftsByThreadIdState = createAtomState<
+  Record<string, string>
+>({
+  key: DRAFTS_STORAGE_KEY,
+  defaultValue: {},
+  useLocalStorage: true,
+  localStorageOptions: { getOnInit: true },
+});
@@ -0,0 +1,6 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const focusEditorAfterMigrateState = createAtomState<boolean>({
+  key: 'ai/focusEditorAfterMigrateState',
+  defaultValue: false,
+});
@@ -0,0 +1,6 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const hasTriggeredCreateForDraftState = createAtomState<boolean>({
+  key: 'ai/hasTriggeredCreateForDraftState',
+  defaultValue: false,
+});
@@ -0,0 +1,6 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const isCreatingForFirstSendState = createAtomState<boolean>({
+  key: 'ai/isCreatingForFirstSendState',
+  defaultValue: false,
+});
@@ -0,0 +1,8 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const pendingCreateFromDraftPromiseState = createAtomState<Promise<
+  string | null
+> | null>({
+  key: 'ai/pendingCreateFromDraftPromiseState',
+  defaultValue: null,
+});
@@ -0,0 +1,6 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const skipMessagesSkeletonUntilLoadedState = createAtomState<boolean>({
+  key: 'ai/skipMessagesSkeletonUntilLoadedState',
+  defaultValue: false,
+});
@@ -0,0 +1,6 @@
+import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
+
+export const threadIdCreatedFromDraftState = createAtomState<string | null>({
+  key: 'ai/threadIdCreatedFromDraftState',
+  defaultValue: null,
+});
@@ -127,7 +127,7 @@ export const MainNavigationDrawerTabsRow = () => {
  );
  const [navigationDrawerActiveTab, setNavigationDrawerActiveTab] =
    useAtomState(navigationDrawerActiveTabState);
-  const { createChatThread } = useCreateNewAIChatThread();
+  const { switchToNewChat } = useCreateNewAIChatThread();
  const isAiEnabled = useIsFeatureEnabled(FeatureFlagKey.IS_AI_ENABLED);
  const setIsNavigationDrawerExpanded = useSetAtomState(
    isNavigationDrawerExpandedState,
@@ -155,7 +155,7 @@ export const MainNavigationDrawerTabsRow = () => {
    if (isMobile) {
      setIsNavigationDrawerExpanded(false);
    }
-    createChatThread();
+    switchToNewChat();
  };

  const handleNewChatKeyDown = (event: React.KeyboardEvent) => {
@@ -36,7 +36,7 @@ export const MobileNavigationBar = () => {
    useAtomState(isNavigationDrawerExpandedState);
  const [currentMobileNavigationDrawer, setCurrentMobileNavigationDrawer] =
    useAtomState(currentMobileNavigationDrawerState);
-  const { createChatThread } = useCreateNewAIChatThread();
+  const { switchToNewChat } = useCreateNewAIChatThread();
  const isAiEnabled = useIsFeatureEnabled(FeatureFlagKey.IS_AI_ENABLED);
  const { alphaSortedActiveNonSystemObjectMetadataItems } =
    useFilteredObjectMetadataItems();
@@ -100,7 +100,7 @@ export const MobileNavigationBar = () => {
            onClick: () => {
              setIsNavigationDrawerExpanded(false);
              closeSidePanelMenu();
-              createChatThread();
+              switchToNewChat();
            },
          },
        ]
@@ -22,7 +22,7 @@ export const SidePanelTopBarRightCornerIcon = () => {
  const isAiEnabled = useIsFeatureEnabled(FeatureFlagKey.IS_AI_ENABLED);
  const sidePanelPage = useAtomStateValue(sidePanelPageState);
  const { openAskAIPage } = useOpenAskAIPageInSidePanel();
-  const { createChatThread } = useCreateNewAIChatThread();
+  const { switchToNewChat } = useCreateNewAIChatThread();

  if (isMobile || !isAiEnabled) {
    return null;
@@ -52,7 +52,7 @@ export const SidePanelTopBarRightCornerIcon = () => {
        Icon={IconEdit}
        size="small"
        variant="tertiary"
-        onClick={() => createChatThread()}
+        onClick={() => switchToNewChat()}
        ariaLabel={t`New conversation`}
      />
    </StyledIconButtonContainer>
@@ -22,15 +22,19 @@ type StateAtom<ValueType> = WritableAtom<
  void
 >;

+type LocalStorageOptions = { getOnInit?: boolean };
+
 export const createAtomState = <ValueType>({
  key,
  defaultValue,
  useLocalStorage = false,
+  localStorageOptions,
  useCookieStorage,
 }: {
  key: string;
  defaultValue: ValueType;
  useLocalStorage?: boolean;
+  localStorageOptions?: LocalStorageOptions;
  useCookieStorage?: CookieStorageConfig<ValueType>;
 }): State<ValueType> => {
  let baseAtom: StateAtom<ValueType>;
@@ -51,6 +55,8 @@ export const createAtomState = <ValueType>({
    baseAtom = atomWithStorage<ValueType>(
      key,
      defaultValue,
+      undefined,
+      localStorageOptions ?? undefined,
    ) as StateAtom<ValueType>;
  } else {
    baseAtom = atom(defaultValue);