fix: issue with memory error message display (#16738)

web/src/pages/memory/memory-message/hook.ts

@@ -78,42 +78,44 @@ export const useMessageAction = () => {
         memory_id: memoryId,
         message_id: selectedMessage.message_id,
       })
-      .then(() => {
-        message.success(t('message.deleted'));
-        queryClient.invalidateQueries({
-          queryKey: [MemoryApiAction.FetchMemoryMessage],
-        });
+      .then((data: any) => {
+        if (data.data.code === 0) {
+          message.success(t('message.deleted'));
+          queryClient.invalidateQueries({
+            queryKey: [MemoryApiAction.FetchMemoryMessage],
+          });
+        }
       });
     setShowDeleteDialog(false);
-  }, [selectedMessage.message_id, queryClient]);
+  }, [selectedMessage.message_id, queryClient, memoryId]);
 
   const handleUpdateMessageState = useCallback(
-    (messageInfo: IMessageInfo, enable: boolean) => {
-      // delete message
-      const selectedMessageInfo = messageInfo || selectedMessage;
-      memoryService
-        .updateMessageState({
+    async (messageInfo: IMessageInfo, enable: boolean): Promise<boolean> => {
+      try {
+        const { data } = await memoryService.updateMessageState({
           memory_id: memoryId,
-          message_id: selectedMessageInfo.message_id,
+          message_id: messageInfo.message_id,
           status: enable || false,
-        })
-        .then((data: any) => {
-          if (data.data.code === 0) {
-            message.success(t('message.updated'));
-            queryClient.invalidateQueries({
-              queryKey: [MemoryApiAction.FetchMemoryMessage],
-            });
-          }
         });
-      setShowDeleteDialog(false);
+        if (data.code === 0) {
+          message.success(t('message.updated'));
+          queryClient.invalidateQueries({
+            queryKey: [MemoryApiAction.FetchMemoryMessage],
+          });
+          return true;
+        }
+        return false;
+      } catch {
+        return false;
+      }
     },
-    [selectedMessage, queryClient, memoryId],
+    [queryClient, memoryId],
   );
 
   const handleClickUpdateMessageState = useCallback(
-    (message: IMessageInfo, enable: boolean) => {
-      setSelectedMessage(message);
-      handleUpdateMessageState(message, enable);
+    (messageInfo: IMessageInfo, enable: boolean): Promise<boolean> => {
+      setSelectedMessage(messageInfo);
+      return handleUpdateMessageState(messageInfo, enable);
     },
     [handleUpdateMessageState],
   );

web/src/pages/memory/memory-message/message-table.tsx

@@ -51,7 +51,7 @@ import {
   TextSelect,
 } from 'lucide-react';
 import * as React from 'react';
-import { useMemo, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { CopyToClipboard } from 'react-copy-to-clipboard';
 import { useMessageAction } from './hook';
 import { IMessageInfo } from './interface';
@@ -75,6 +75,48 @@ function getTaskStatus(progress: number) {
   }
 }
 
+type UpdateMessageStateFn = (
+  message: IMessageInfo,
+  enable: boolean,
+) => Promise<boolean>;
+
+function StatusSwitch({
+  row,
+  disabled,
+  onUpdate,
+}: {
+  row: Row<IMessageInfo>;
+  disabled: boolean;
+  onUpdate: UpdateMessageStateFn;
+}) {
+  const serverStatus = row.original.status;
+  const [optimisticStatus, setOptimisticStatus] = useState(serverStatus);
+  const [pending, setPending] = useState(false);
+
+  useEffect(() => {
+    setOptimisticStatus(serverStatus);
+  }, [serverStatus]);
+
+  return (
+    <div className="flex items-center">
+      <Switch
+        disabled={disabled || pending}
+        checked={optimisticStatus}
+        onCheckedChange={(val) => {
+          setOptimisticStatus(val);
+          setPending(true);
+          onUpdate(row.original, val).then((success) => {
+            setPending(false);
+            if (!success) {
+              setOptimisticStatus(serverStatus);
+            }
+          });
+        }}
+      />
+    </div>
+  );
+}
+
 export function MemoryTable({
   messages,
   total,
@@ -216,20 +258,13 @@ export function MemoryTable({
       {
         accessorKey: 'status',
         header: () => <span>{t('memory.messages.enable')}</span>,
-        cell: ({ row }) => {
-          const isEnabled = row.getValue('status') as boolean;
-          return (
-            <div className="flex items-center">
-              <Switch
-                disabled={disabledRowFunc(row)}
-                defaultChecked={isEnabled}
-                onCheckedChange={(val) => {
-                  handleClickUpdateMessageState(row.original, val);
-                }}
-              />
-            </div>
-          );
-        },
+        cell: ({ row }) => (
+          <StatusSwitch
+            row={row}
+            disabled={disabledRowFunc(row)}
+            onUpdate={handleClickUpdateMessageState}
+          />
+        ),
       },
       columnHelper.display({
         id: 'task_progress',