fix(ai-chat) - improvements (#22193)

- remove flickering at assistant message streamed end - add copy code - leave chat history when navigating to settings  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/22193?utm_source=github" target="_blank" rel="noopener noreferrer" data-no-image-dialog="true"><picture><source media="(prefers-color-scheme: dark)" srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img alt="Review in cubic" src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a>
2026-06-26 10:26:40 +02:00
parent e60d790990
commit b625bd1995
8 changed files with 155 additions and 14 deletions
@@ -1,13 +1,14 @@
-import { useCallback, useMemo } from 'react';
 import { useStore } from 'jotai';
+import { useCallback, useMemo } from 'react';
 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 { 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 { agentChatFirstLiveSeqComponentFamilyState } from '@/ai/states/agentChatFirstLiveSeqComponentFamilyState';
+import { agentChatHandleEventCallbackComponentFamilyState } from '@/ai/states/agentChatHandleEventCallbackComponentFamilyState';
+import { agentChatIsAwaitingPersistedRefetchComponentFamilyState } from '@/ai/states/agentChatIsAwaitingPersistedRefetchComponentFamilyState';
 import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
 import { agentChatQueuedMessagesComponentFamilyState } from '@/ai/states/agentChatQueuedMessagesComponentFamilyState';
 import { currentAiChatThreadState } from '@/ai/states/currentAiChatThreadState';
@@ -53,6 +54,11 @@ export const AgentChatMessagesFetchEffect = () => {
    { threadId: currentAiChatThread },
  );

+  const setAgentChatIsAwaitingPersistedRefetch = useSetAtomComponentFamilyState(
+    agentChatIsAwaitingPersistedRefetchComponentFamilyState,
+    { threadId: currentAiChatThread },
+  );
+
  const handleEventCallbackFamilyCallback =
    useAtomComponentFamilyStateCallbackState(
      agentChatHandleEventCallbackComponentFamilyState,
@@ -77,6 +83,7 @@ export const AgentChatMessagesFetchEffect = () => {
      setAgentChatQueuedMessages(
        uiMessages.filter((message) => message.status === 'queued'),
      );
+      setAgentChatIsAwaitingPersistedRefetch(false);

      const catchup = data.chatStreamCatchupChunks;

@@ -119,6 +126,7 @@ export const AgentChatMessagesFetchEffect = () => {
    [
      setAgentChatFetchedMessages,
      setAgentChatQueuedMessages,
+      setAgentChatIsAwaitingPersistedRefetch,
      store,
      handleEventCallbackFamilyCallback,
      firstLiveSeqFamilyCallback,
@@ -128,8 +136,12 @@ export const AgentChatMessagesFetchEffect = () => {
  const handleLoadingChange = useCallback(
    (loading: boolean) => {
      setAgentChatMessagesLoading(loading);
+
+      if (!loading) {
+        setAgentChatIsAwaitingPersistedRefetch(false);
+      }
    },
-    [setAgentChatMessagesLoading],
+    [setAgentChatMessagesLoading, setAgentChatIsAwaitingPersistedRefetch],
  );

  const { refetch: refetchAgentChatMessages } = useQueryWithCallbacks(
@@ -9,6 +9,7 @@ import { useEnsureAgentChatThreadExistsForDraft } from '@/ai/hooks/useEnsureAgen
 import { useEnsureAgentChatThreadIdForSend } from '@/ai/hooks/useEnsureAgentChatThreadIdForSend';
 import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThreadState';
 import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
+import { agentChatIsAwaitingPersistedRefetchComponentFamilyState } from '@/ai/states/agentChatIsAwaitingPersistedRefetchComponentFamilyState';
 import { agentChatIsInitialScrollPendingOnThreadChangeState } from '@/ai/states/agentChatIsInitialScrollPendingOnThreadChangeState';
 import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
 import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
@@ -62,6 +63,11 @@ export const AgentChatStreamSubscriptionEffect = () => {
    { threadId: currentAiChatThread },
  );

+  const agentChatIsAwaitingPersistedRefetch = useAtomComponentFamilyStateValue(
+    agentChatIsAwaitingPersistedRefetchComponentFamilyState,
+    { threadId: currentAiChatThread },
+  );
+
  const agentChatDisplayedThread = useAtomStateValue(
    agentChatDisplayedThreadState,
  );
@@ -79,9 +85,15 @@ export const AgentChatStreamSubscriptionEffect = () => {
      return;
    }

+    const isThreadSwitch = currentAiChatThread !== agentChatDisplayedThread;
+
+    if (!isThreadSwitch && agentChatIsAwaitingPersistedRefetch) {
+      return;
+    }
+
    setAgentChatMessages(agentChatFetchedMessages);

-    if (currentAiChatThread !== agentChatDisplayedThread) {
+    if (isThreadSwitch) {
      if (agentChatFetchedMessages.length > 0) {
        setAgentChatIsInitialScrollPendingOnThreadChange(true);
      }
@@ -90,6 +102,7 @@ export const AgentChatStreamSubscriptionEffect = () => {
  }, [
    agentChatFetchedMessages,
    agentChatIsStreaming,
+    agentChatIsAwaitingPersistedRefetch,
    setAgentChatMessages,
    currentAiChatThread,
    agentChatDisplayedThread,
@@ -10,6 +10,7 @@ import {
  StyledSkeletonContainer,
  StyledTableScrollContainer,
 } from '@/ai/components/LazyMarkdownRendererStyledComponents';
+import { MarkdownCodeBlock } from '@/ai/components/MarkdownCodeBlock';
 import {
  cloneElement,
  isValidElement,
@@ -158,9 +159,7 @@ const MarkdownRenderer = lazy(async () => {
            children?: React.ReactNode;
          }) => <code className={className}>{children}</code>,
          pre: ({ children }) => (
-            <div className="markdown-code-outer-container">
-              <pre className="markdown-block-code">{children}</pre>
-            </div>
+            <MarkdownCodeBlock>{children}</MarkdownCodeBlock>
          ),
        }}
      >
@@ -0,0 +1,80 @@
+import { styled } from '@linaria/react';
+import { useLingui } from '@lingui/react/macro';
+import { isNumber, isString } from '@sniptt/guards';
+import { Children, isValidElement } from 'react';
+import { IconCopy } from 'twenty-ui/icon';
+import { LightIconButton } from 'twenty-ui/input';
+import { themeCssVariables } from 'twenty-ui/theme-constants';
+import { useCopyToClipboard } from '~/hooks/useCopyToClipboard';
+
+const StyledContainer = styled.div`
+  position: relative;
+
+  .markdown-block-code {
+    margin: 0;
+  }
+`;
+
+const StyledCopyButtonContainer = styled.div`
+  align-items: center;
+  border-radius: ${themeCssVariables.border.radius.sm};
+  display: flex;
+  opacity: 0;
+  position: absolute;
+  right: ${themeCssVariables.spacing[2]};
+  top: ${themeCssVariables.spacing[2]};
+  transition: opacity calc(${themeCssVariables.animation.duration.fast} * 1s)
+    ease;
+
+  ${StyledContainer}:hover & {
+    opacity: 1;
+  }
+
+  &:hover {
+    background: ${themeCssVariables.background.primary};
+  }
+`;
+
+const extractTextFromNode = (node: React.ReactNode): string => {
+  if (isString(node) || isNumber(node)) {
+    return node.toString();
+  }
+
+  if (Array.isArray(node)) {
+    return node.map(extractTextFromNode).join('');
+  }
+
+  if (isValidElement<{ children?: React.ReactNode }>(node)) {
+    return Children.toArray(node.props.children)
+      .map(extractTextFromNode)
+      .join('');
+  }
+
+  return '';
+};
+
+export const MarkdownCodeBlock = ({
+  children,
+}: {
+  children: React.ReactNode;
+}) => {
+  const { t } = useLingui();
+  const { copyToClipboard } = useCopyToClipboard();
+
+  const codeText = extractTextFromNode(children);
+
+  return (
+    <StyledContainer className="markdown-code-outer-container">
+      <StyledCopyButtonContainer>
+        <LightIconButton
+          Icon={IconCopy}
+          onClick={() => copyToClipboard(codeText, t`Code copied to clipboard`)}
+          title={t`Copy code`}
+          size="small"
+          accent="tertiary"
+        />
+      </StyledCopyButtonContainer>
+      <pre className="markdown-block-code">{children}</pre>
+    </StyledContainer>
+  );
+};
@@ -15,6 +15,7 @@ import { ON_AGENT_CHAT_EVENT } from '@/ai/graphql/subscriptions/OnAgentChatEvent
 import { agentChatErrorComponentFamilyState } from '@/ai/states/agentChatErrorComponentFamilyState';
 import { agentChatFirstLiveSeqComponentFamilyState } from '@/ai/states/agentChatFirstLiveSeqComponentFamilyState';
 import { agentChatHandleEventCallbackComponentFamilyState } from '@/ai/states/agentChatHandleEventCallbackComponentFamilyState';
+import { agentChatIsAwaitingPersistedRefetchComponentFamilyState } from '@/ai/states/agentChatIsAwaitingPersistedRefetchComponentFamilyState';
 import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState';
 import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
 import { agentChatUsageComponentFamilyState } from '@/ai/states/agentChatUsageComponentFamilyState';
@@ -108,6 +109,10 @@ export const useAgentChatSubscription = (threadId: string | null) => {
  const firstLiveSeqFamilyCallback = useAtomComponentFamilyStateCallbackState(
    agentChatFirstLiveSeqComponentFamilyState,
  );
+  const isAwaitingPersistedRefetchFamilyCallback =
+    useAtomComponentFamilyStateCallbackState(
+      agentChatIsAwaitingPersistedRefetchComponentFamilyState,
+    );
  const handleEventCallbackFamilyCallback =
    useAtomComponentFamilyStateCallbackState(
      agentChatHandleEventCallbackComponentFamilyState,
@@ -132,6 +137,8 @@ export const useAgentChatSubscription = (threadId: string | null) => {
    const errorAtom = errorFamilyCallback(familyKey);
    const isStreamingAtom = isStreamingFamilyCallback(familyKey);
    const firstLiveSeqAtom = firstLiveSeqFamilyCallback(familyKey);
+    const isAwaitingPersistedRefetchAtom =
+      isAwaitingPersistedRefetchFamilyCallback(familyKey);
    const handleEventCallbackAtom =
      handleEventCallbackFamilyCallback(familyKey);
    const messagesAtom = messagesFamilyCallback(familyKey);
@@ -303,6 +310,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {

        case 'message-persisted': {
          closeWriter();
+          store.set(isAwaitingPersistedRefetchAtom, true);
          dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
          break;
        }
@@ -363,6 +371,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {
          store.set(errorAtom, noMoreCreditsError);

          closeWriter();
+          store.set(isAwaitingPersistedRefetchAtom, true);
          dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
          store.set(isStreamingAtom, false);
          break;
@@ -412,6 +421,7 @@ export const useAgentChatSubscription = (threadId: string | null) => {
    errorFamilyCallback,
    isStreamingFamilyCallback,
    firstLiveSeqFamilyCallback,
+    isAwaitingPersistedRefetchFamilyCallback,
    handleEventCallbackFamilyCallback,
    messagesFamilyCallback,
    usageFamilyCallback,
@@ -0,0 +1,9 @@
+import { AgentChatComponentInstanceContext } from '@/ai/contexts/AgentChatComponentInstanceContext';
+import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
+
+export const agentChatIsAwaitingPersistedRefetchComponentFamilyState =
+  createAtomComponentFamilyState<boolean, { threadId: string | null }>({
+    key: 'agentChatIsAwaitingPersistedRefetchComponentFamilyState',
+    defaultValue: false,
+    componentInstanceContext: AgentChatComponentInstanceContext,
+  });
@@ -2,6 +2,8 @@ import { useIsSettingsPage } from '@/navigation/hooks/useIsSettingsPage';
 import { currentMobileNavigationDrawerState } from '@/navigation/states/currentMobileNavigationDrawerState';
 import { isNavigationDrawerExpandedState } from '@/ui/navigation/states/isNavigationDrawerExpanded';
 import { navigationDrawerExpandedMemorizedState } from '@/ui/navigation/states/navigationDrawerExpandedMemorizedState';
+import { navigationDrawerActiveTabState } from '@/ui/navigation/states/navigationDrawerActiveTabState';
+import { NAVIGATION_DRAWER_TABS } from '@/ui/navigation/states/navigationDrawerTabs';
 import { navigationMemorizedUrlState } from '@/ui/navigation/states/navigationMemorizedUrlState';
 import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
 import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
@@ -22,6 +24,9 @@ export const useOpenSettingsMenu = () => {
  const setCurrentMobileNavigationDrawer = useSetAtomState(
    currentMobileNavigationDrawerState,
  );
+  const setNavigationDrawerActiveTab = useSetAtomState(
+    navigationDrawerActiveTabState,
+  );

  const openSettingsMenu = useCallback(() => {
    if (isSettingsPage) {
@@ -32,6 +37,7 @@ export const useOpenSettingsMenu = () => {
    setIsNavigationDrawerExpanded(true);
    setNavigationMemorizedUrl(location.pathname + location.search);
    setCurrentMobileNavigationDrawer('settings');
+    setNavigationDrawerActiveTab(NAVIGATION_DRAWER_TABS.NAVIGATION_MENU);
  }, [
    isSettingsPage,
    isNavigationDrawerExpanded,
@@ -41,6 +47,7 @@ export const useOpenSettingsMenu = () => {
    setIsNavigationDrawerExpanded,
    setNavigationDrawerExpandedMemorized,
    setNavigationMemorizedUrl,
+    setNavigationDrawerActiveTab,
  ]);

  return { openSettingsMenu };
@@ -1,18 +1,17 @@
 import { Logger, Scope } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';

+import { isNonEmptyString } from '@sniptt/guards';
 import { createUIMessageStream } from 'ai';
 import type {
  CodeExecutionData,
  ExtendedUIMessage,
  ExtendedUIMessagePart,
 } from 'twenty-shared/ai';
-import { isNonEmptyString } from '@sniptt/guards';
-import { Repository } from 'typeorm';
 import { isDefined } from 'twenty-shared/utils';
+import { Repository } from 'typeorm';
+import { v4 } from 'uuid';

-import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
-import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';
 import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
 import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
 import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
@@ -30,6 +29,8 @@ import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/service
 import { ChatExecutionService } from 'src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service';
 import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
 import type { AiModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
+import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
+import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository';

 import { STREAM_AGENT_CHAT_JOB_NAME } from './stream-agent-chat-job-name.constant';
 import { type StreamAgentChatJobData } from './stream-agent-chat-job.types';
@@ -184,6 +185,8 @@ export class StreamAgentChatJob {
    abortSignal: AbortSignal;
  }): Promise<void> {
    return new Promise<void>((resolve, reject) => {
+      const assistantMessageId = v4();
+
      let streamUsage = {
        inputTokens: 0,
        outputTokens: 0,
@@ -259,7 +262,8 @@ export class StreamAgentChatJob {

                return error instanceof Error ? error.message : String(error);
              },
-              sendStart: false,
+              sendStart: true,
+              generateMessageId: () => assistantMessageId,
              messageMetadata: ({ part }) => {
                return this.computeMessageMetadata({
                  part,
@@ -280,6 +284,7 @@ export class StreamAgentChatJob {
              onFinish: async ({ responseMessage, isAborted }) => {
                try {
                  await this.handleStreamFinish({
+                    assistantMessageId,
                    responseMessage,
                    isAborted,
                    streamError,
@@ -335,7 +340,10 @@ export class StreamAgentChatJob {
            await this.eventPublisherService.publish({
              threadId: data.threadId,
              workspaceId: data.workspaceId,
-              event: { type: 'message-persisted', messageId: data.threadId },
+              event: {
+                type: 'message-persisted',
+                messageId: assistantMessageId,
+              },
            });
            resolve();
          }
@@ -435,6 +443,7 @@ export class StreamAgentChatJob {
  }

  private async handleStreamFinish({
+    assistantMessageId,
    responseMessage,
    isAborted,
    streamError,
@@ -448,6 +457,7 @@ export class StreamAgentChatJob {
    modelConfig,
    userMessagePromise,
  }: {
+    assistantMessageId: string;
    responseMessage: Omit<ExtendedUIMessage, 'id'>;
    isAborted: boolean;
    streamError: unknown;
@@ -502,6 +512,7 @@ export class StreamAgentChatJob {
    await this.agentChatService.addMessage({
      threadId,
      uiMessage: responseMessage,
+      id: assistantMessageId,
      turnId: userMessage.turnId ?? undefined,
      workspaceId,
    });