From e2c634558f93d3215e4618bc867449171895e0fe Mon Sep 17 00:00:00 2001
From: Gimir <mukhammad.amaev@mail.ru>
Date: Wed, 29 Jan 2025 14:10:50 +0300
Subject: [PATCH] fix(chat): fix custom buttons in playback mode, hide
 templates in schema chats (Issue #3002, #3013) (#3035)

Co-authored-by: Magomed-Elbi Dzhukalaev <magomed-elbi_dzhukalaev@epam.com>
---
 apps/chat/src/components/Chat/Chat.tsx        |  2 +-
 .../ChatMessageContent/UserMessage.tsx        |  1 +
 .../Chat/ChatMessage/MessageButtons.tsx       | 12 +++++-
 .../chat/src/components/Chat/ChatStarters.tsx |  4 ++
 .../Chat/Playback/PlaybackControls.tsx        | 38 ++++++++++++++++---
 .../conversations/conversations.selectors.ts  | 10 ++++-
 6 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/apps/chat/src/components/Chat/Chat.tsx b/apps/chat/src/components/Chat/Chat.tsx
index a2e197c1c4..f0dfc3d26f 100644
--- a/apps/chat/src/components/Chat/Chat.tsx
+++ b/apps/chat/src/components/Chat/Chat.tsx
@@ -755,7 +755,7 @@ export const ChatView = memo(() => {
                         />
                       )}
 
-                      {!isPlayback && <ChatStarters />}
+                      <ChatStarters />
 
                       {!isPlayback && (
                         <ChatInput
diff --git a/apps/chat/src/components/Chat/ChatMessage/ChatMessageContent/UserMessage.tsx b/apps/chat/src/components/Chat/ChatMessage/ChatMessageContent/UserMessage.tsx
index f380d3d427..821e90cf62 100644
--- a/apps/chat/src/components/Chat/ChatMessage/ChatMessageContent/UserMessage.tsx
+++ b/apps/chat/src/components/Chat/ChatMessage/ChatMessageContent/UserMessage.tsx
@@ -112,6 +112,7 @@ export const UserMessage = memo(function UserMessage({
     SettingsSelectors.isFeatureEnabled(state, Feature.MessageTemplates),
   );
   const isChatFullWidth = useAppSelector(UISelectors.selectIsChatFullWidth);
+
   const isMobileOrOverlay = isSmallScreen() || isOverlay;
   const isInputDisabled = isMessageInputDisabled(messageIndex, allMessages);
 
diff --git a/apps/chat/src/components/Chat/ChatMessage/MessageButtons.tsx b/apps/chat/src/components/Chat/ChatMessage/MessageButtons.tsx
index 64a970117b..b9a6ec31ff 100644
--- a/apps/chat/src/components/Chat/ChatMessage/MessageButtons.tsx
+++ b/apps/chat/src/components/Chat/ChatMessage/MessageButtons.tsx
@@ -18,6 +18,7 @@ import { getMessageCustomContent } from '@/src/utils/server/chat';
 
 import { Translation } from '@/src/types/translation';
 
+import { ConversationsSelectors } from '@/src/store/conversations/conversations.reducers';
 import { useAppSelector } from '@/src/store/hooks';
 import { SettingsSelectors } from '@/src/store/settings/settings.reducers';
 
@@ -68,6 +69,9 @@ export const MessageUserButtons = ({
   const { t } = useTranslation(Translation.Chat);
 
   const isOverlay = useAppSelector(SettingsSelectors.selectIsOverlay);
+  const isConversationsWithSchema = useAppSelector(
+    ConversationsSelectors.selectIsSelectedConversationsWithSchema,
+  );
 
   return (
     <div
@@ -78,7 +82,7 @@ export const MessageUserButtons = ({
     >
       {!isMessageStreaming && (
         <>
-          {isEditTemplatesAvailable && (
+          {isEditTemplatesAvailable && !isConversationsWithSchema && (
             <Tooltip
               placement="top"
               isTriggerClickable
@@ -274,6 +278,10 @@ export const MessageMobileButtons = ({
 }: MessageMobileButtonsProps) => {
   const { t } = useTranslation(Translation.Chat);
 
+  const isConversationsWithSchema = useAppSelector(
+    ConversationsSelectors.selectIsSelectedConversationsWithSchema,
+  );
+
   const isAssistant = message.role === Role.Assistant;
 
   if (isAssistant) {
@@ -389,7 +397,7 @@ export const MessageMobileButtons = ({
     !isMessageStreaming &&
     !isConversationInvalid && (
       <>
-        {isEditTemplatesAvailable && (
+        {isEditTemplatesAvailable && !isConversationsWithSchema && (
           <MenuItem
             className="hover:bg-accent-primary-alpha focus:visible disabled:cursor-not-allowed group-hover:visible"
             onClick={() => onToggleTemplatesEditing()}
diff --git a/apps/chat/src/components/Chat/ChatStarters.tsx b/apps/chat/src/components/Chat/ChatStarters.tsx
index 9bf6236423..ea633a53cc 100644
--- a/apps/chat/src/components/Chat/ChatStarters.tsx
+++ b/apps/chat/src/components/Chat/ChatStarters.tsx
@@ -23,6 +23,9 @@ const ChatStartersView = ({ schema }: ChatStartersViewProps) => {
   const dispatch = useAppDispatch();
 
   const formValue = useAppSelector(ChatSelectors.selectChatFormValue);
+  const isPlayback = useAppSelector(
+    ConversationsSelectors.selectIsPlaybackSelectedConversations,
+  );
 
   const handleChange = useCallback(
     (property: string, value: MessageFormValueType, submit?: boolean) => {
@@ -56,6 +59,7 @@ const ChatStartersView = ({ schema }: ChatStartersViewProps) => {
       buttonsWrapperClassName="md:justify-center flex-nowrap overflow-x-auto overflow-y-hidden px-2"
       buttonClassName="shrink-0"
       propertyWrapperClassName="items-center"
+      disabled={isPlayback}
     />
   );
 };
diff --git a/apps/chat/src/components/Chat/Playback/PlaybackControls.tsx b/apps/chat/src/components/Chat/Playback/PlaybackControls.tsx
index 59762260c2..db3632252c 100644
--- a/apps/chat/src/components/Chat/Playback/PlaybackControls.tsx
+++ b/apps/chat/src/components/Chat/Playback/PlaybackControls.tsx
@@ -12,8 +12,13 @@ import { useTranslation } from 'next-i18next';
 
 import classNames from 'classnames';
 
+import {
+  getConfigurationValue,
+  getMessageFormValue,
+} from '@/src/utils/app/form-schema';
 import { hasParentWithFloatingOverlay } from '@/src/utils/app/modals';
 
+import { ChatActions } from '@/src/store/chat/chat.reducer';
 import {
   ConversationsActions,
   ConversationsSelectors,
@@ -26,6 +31,8 @@ import { ScrollDownButton } from '@/src/components/Common/ScrollDownButton';
 import { ChatInputFooter } from '../ChatInput/ChatInputFooter';
 import { PlaybackAttachments } from './PlaybackAttachments';
 
+import { Attachment } from '@epam/ai-dial-shared';
+
 interface Props {
   showScrollDownButton: boolean;
   onScrollDownClick: () => void;
@@ -105,17 +112,21 @@ export const PlaybackControls = ({
       currentMessage && currentMessage?.custom_content?.attachments?.length
         ? currentMessage.custom_content.attachments
         : [];
+    const form_value =
+      getMessageFormValue(currentMessage) ??
+      getConfigurationValue(currentMessage);
     const message = attachments.length
-      ? { content, custom_content: { attachments } }
-      : { content };
+      ? { content, custom_content: { attachments, form_value } }
+      : { content, custom_content: { form_value } };
     return message;
   }, [activeIndex, isActiveIndex, isNextMessageInStack, selectedConversations]);
 
-  const hasAttachments =
+  const hasAttachments = !!(
     activeMessage &&
     activeMessage.custom_content &&
     activeMessage.custom_content.attachments &&
-    activeMessage.custom_content.attachments.length;
+    activeMessage.custom_content.attachments.length
+  );
 
   const handlePlayNextMessage = useCallback(() => {
     if (isMessageStreaming || !isNextMessageInStack) {
@@ -209,6 +220,21 @@ export const PlaybackControls = ({
     };
   }, [controlsContainerRef, onResize]);
 
+  useEffect(() => {
+    if (
+      phase === PlaybackPhases.MESSAGE &&
+      activeMessage?.custom_content?.form_value
+    ) {
+      Object.entries(activeMessage.custom_content.form_value).forEach(
+        ([property, value]) => {
+          dispatch(ChatActions.setFormValue({ property, value }));
+        },
+      );
+    } else if (phase === PlaybackPhases.EMPTY) {
+      dispatch(ChatActions.resetFormValue());
+    }
+  }, [activeMessage?.custom_content?.form_value, dispatch, phase]);
+
   return (
     <div ref={controlsContainerRef} className="w-full pt-3 md:pt-5">
       <div
@@ -254,7 +280,9 @@ export const PlaybackControls = ({
 
                   {phase === PlaybackPhases.MESSAGE && hasAttachments && (
                     <PlaybackAttachments
-                      attachments={activeMessage.custom_content.attachments}
+                      attachments={
+                        activeMessage.custom_content.attachments as Attachment[]
+                      }
                     />
                   )}
                   <button
diff --git a/apps/chat/src/store/conversations/conversations.selectors.ts b/apps/chat/src/store/conversations/conversations.selectors.ts
index 33db36fe9a..11ebcb6f27 100644
--- a/apps/chat/src/store/conversations/conversations.selectors.ts
+++ b/apps/chat/src/store/conversations/conversations.selectors.ts
@@ -24,7 +24,10 @@ import {
   sortByName,
   splitEntityId,
 } from '@/src/utils/app/folders';
-import { isMessageInputDisabled } from '@/src/utils/app/form-schema';
+import {
+  isConversationWithFormSchema,
+  isMessageInputDisabled,
+} from '@/src/utils/app/form-schema';
 import {
   getConversationRootId,
   isEntityIdExternal,
@@ -852,3 +855,8 @@ export const selectIsSelectedConversationBlocksInput = createSelector(
       ),
     ),
 );
+
+export const selectIsSelectedConversationsWithSchema = createSelector(
+  [selectSelectedConversations],
+  (conversations) => conversations.some(isConversationWithFormSchema),
+);