fix: allow user to send new messages after generation stop

Author: tfomkin

libs/mobile/chat/features/chat/src/lib/component.tsx

@@ -173,9 +173,9 @@ export function Chat({ chatId, selectedModelId, isNewChat, resetToChatsList }: C
             attachedImages={attachedImages}
             onImageUploaded={handleImageUploaded}
             onDeleteImagePress={handleDeleteImage}
-            chatId={chatId}
             modelId={selectedModelId}
             isResponseGenerating={isResponseGenerating}
+            chat={chat}
           />
         )}
       </View>

libs/mobile/chat/features/form-chat-input/src/lib/component.tsx

@@ -16,6 +16,7 @@ import { AppTextInput, AppInputProps, View, IconButton } from '@open-webui-react
 import {
   appConfigurationApi,
   ChatGenerationOption,
+  ChatResponse,
   tasksApi,
   tasksService,
 } from '@open-webui-react-native/shared/data-access/api';
@@ -36,7 +37,7 @@ interface FormChatInputProps<T extends FieldValues> extends AppInputProps {
   attachedImages: Observable<Array<ImageData>>;
   onImageUploaded: (image: ImageData) => void;
   onDeleteImagePress: (fileName: string) => void;
-  chatId?: string;
+  chat?: ChatResponse;
   modelId?: string;
   onChatCreated?: (id: string) => void;
   isLoading?: boolean;
@@ -58,7 +59,7 @@ export function FormChatInput<T extends FieldValues>({
   attachedImages,
   onImageUploaded,
   onDeleteImagePress,
-  chatId,
+  chat,
   modelId,
   onChatCreated,
   isLoading,
@@ -96,7 +97,7 @@ export function FormChatInput<T extends FieldValues>({
       return ToastService.showError(translate('TEXT_MODEL_NOT_SELECTED'));
     }
     setIsMicrophonePreparing(true);
-    await openVoiceModeModal({ chatId, modelId });
+    await openVoiceModeModal({ chatId: chat?.id, modelId });
     setIsMicrophonePreparing(false);
   };
 
@@ -119,13 +120,16 @@ export function FormChatInput<T extends FieldValues>({
   };
 
   const onStopGenerationPress = async (): Promise<void> => {
-    if (!chatId) return;
+    if (!chat) return;
+
+    const chatId = chat.id;
+    const lastMessageId = chat.chat.history.currentId;
 
     const tasksData = await tasksService.getChatTasks(chatId);
     const taskId = tasksData?.tasksIds[0];
 
     if (taskId) {
-      stopTaskMutation.mutate(taskId);
+      stopTaskMutation.mutate({ taskId, chatId, lastMessageId });
     }
   };
 
@@ -163,6 +167,8 @@ export function FormChatInput<T extends FieldValues>({
               isSubmitDisabled={!isFeatureEnabled(FeatureID.VOICE_MODE) && isInputEmpty}
               onVoiceModePress={onVoiceModePress}
               isVoiceModeAvailable={isFeatureEnabled(FeatureID.VOICE_MODE) && isInputEmpty}
+              onStopGenerationPress={onStopGenerationPress}
+              isResponseGenerating={isResponseGenerating}
               isLoading={isLoading || isMicrophonePreparing}>
               <View className='flex-row flex-1 justify-between'>
                 <View className='gap-16 flex-row '>
@@ -180,20 +186,12 @@ export function FormChatInput<T extends FieldValues>({
                     />
                   )}
                 </View>
-                <View className={'flex-row gap-16 pr-16'}>
-                  {isResponseGenerating && (
-                    <IconButton
-                      iconName='stopCircle'
-                      className='p-0'
-                      onPress={onStopGenerationPress} />
-                  )}
-                  <IconButton
-                    disabled={isLoading}
-                    iconName='microphone'
-                    className='p-0'
-                    onPress={handleDictateModePress}
-                  />
-                </View>
+                <IconButton
+                  disabled={isLoading}
+                  iconName='microphone'
+                  className='p-0 mr-16'
+                  onPress={handleDictateModePress}
+                />
               </View>
             </ChatInputBottomRow>
           }

libs/mobile/chat/features/form-chat-input/src/lib/components/chat-input-bottom-row/component.tsx

@@ -4,6 +4,8 @@ import { IconButton, View } from '@open-webui-react-native/mobile/shared/ui/ui-k
 export interface ChatInputBottomRowProps extends PropsWithChildren {
   onSubmit: () => void;
   onVoiceModePress: () => void;
+  onStopGenerationPress: () => void;
+  isResponseGenerating?: boolean;
   isVoiceModeAvailable?: boolean;
   isSubmitDisabled?: boolean;
   isLoading?: boolean;
@@ -16,18 +18,27 @@ export function ChatInputBottomRow({
   isLoading,
   children,
   isSubmitDisabled,
+  isResponseGenerating,
+  onStopGenerationPress,
 }: ChatInputBottomRowProps): ReactElement {
   return (
     <View className='flex-row justify-between items-center mt-12'>
       {children}
-      <IconButton
-        disabled={isSubmitDisabled}
-        onPress={isVoiceModeAvailable ? onVoiceModePress : onSubmit}
-        iconName={isVoiceModeAvailable ? 'headphones' : 'arrowUp'}
-        className='rounded-full self-end bg-text-primary p-4'
-        iconProps={{ className: 'color-background-primary' }}
-        isLoading={isLoading}
-      />
+      {isResponseGenerating ? (
+        <IconButton
+          iconName='stopCircle'
+          className='p-0'
+          onPress={onStopGenerationPress} />
+      ) : (
+        <IconButton
+          disabled={isSubmitDisabled}
+          onPress={isVoiceModeAvailable ? onVoiceModePress : onSubmit}
+          iconName={isVoiceModeAvailable ? 'headphones' : 'arrowUp'}
+          className='rounded-full self-end bg-text-primary p-4'
+          iconProps={{ className: 'color-background-primary' }}
+          isLoading={isLoading}
+        />
+      )}
     </View>
   );
 }

libs/shared/data-access/api/src/lib/tasks/api.ts

@@ -1,14 +1,38 @@
 import { useMutation, UseMutationOptions, UseMutationResult } from '@tanstack/react-query';
 import { AxiosError } from 'axios';
 import { ApiErrorData } from '@open-webui-react-native/shared/data-access/api-client';
+import { Chat, patchChatQueryData } from '../chats';
 import { StopTaskResponse } from './models';
 import { tasksService } from './service';
 
+type StopTaskArgs = {
+  taskId: string;
+  chatId: string;
+  lastMessageId: string;
+};
+
 function useStopTask(
-  props?: UseMutationOptions<StopTaskResponse, AxiosError<ApiErrorData>, string>,
-): UseMutationResult<StopTaskResponse, AxiosError<ApiErrorData>, string> {
+  props?: UseMutationOptions<StopTaskResponse, AxiosError<ApiErrorData>, StopTaskArgs>,
+): UseMutationResult<StopTaskResponse, AxiosError<ApiErrorData>, StopTaskArgs> {
   return useMutation({
-    mutationFn: (taskId: string) => tasksService.stopTask(taskId),
+    mutationFn: ({ taskId }) => tasksService.stopTask(taskId),
+
+    onSuccess: (_, { chatId, lastMessageId }) => {
+      patchChatQueryData(chatId, {
+        chat: {
+          history: {
+            messages: {
+              [lastMessageId]: {
+                done: true,
+              },
+            },
+          },
+        } as Chat,
+      });
+
+      props?.onSuccess?.(_, { chatId, lastMessageId } as any, undefined);
+    },
+
     ...props,
   });
 }