refactor(frontend): migration of use-optimistic-user-message.ts to zustand (#11148)

frontend/__tests__/components/chat/chat-interface.test.tsx

@@ -17,7 +17,7 @@ import type { Message } from "#/message";
 import { SUGGESTIONS } from "#/utils/suggestions";
 import { ChatInterface } from "#/components/features/chat/chat-interface";
 import { useWsClient } from "#/context/ws-client-provider";
-import { useOptimisticUserMessage } from "#/hooks/use-optimistic-user-message";
+import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store";
 import { useWSErrorMessage } from "#/hooks/use-ws-error-message";
 import { useConfig } from "#/hooks/query/use-config";
 import { useGetTrajectory } from "#/hooks/mutation/use-get-trajectory";
@@ -26,7 +26,7 @@ import { OpenHandsAction } from "#/types/core/actions";
 
 // Mock the hooks
 vi.mock("#/context/ws-client-provider");
-vi.mock("#/hooks/use-optimistic-user-message");
+vi.mock("#/stores/optimistic-user-message-store");
 vi.mock("#/hooks/use-ws-error-message");
 vi.mock("#/hooks/query/use-config");
 vi.mock("#/hooks/mutation/use-get-trajectory");
@@ -109,7 +109,7 @@ describe("ChatInterface - Chat Suggestions", () => {
       parsedEvents: [],
     });
     (
-      useOptimisticUserMessage as unknown as ReturnType<typeof vi.fn>
+      useOptimisticUserMessageStore as unknown as ReturnType<typeof vi.fn>
     ).mockReturnValue({
       setOptimisticUserMessage: vi.fn(),
       getOptimisticUserMessage: vi.fn(() => null),
@@ -203,7 +203,7 @@ describe("ChatInterface - Chat Suggestions", () => {
 
   test("should hide chat suggestions when there is an optimistic user message", () => {
     (
-      useOptimisticUserMessage as unknown as ReturnType<typeof vi.fn>
+      useOptimisticUserMessageStore as unknown as ReturnType<typeof vi.fn>
     ).mockReturnValue({
       setOptimisticUserMessage: vi.fn(),
       getOptimisticUserMessage: vi.fn(() => "Optimistic message"),
@@ -246,7 +246,7 @@ describe("ChatInterface - Empty state", () => {
       parsedEvents: [],
     });
     (
-      useOptimisticUserMessage as unknown as ReturnType<typeof vi.fn>
+      useOptimisticUserMessageStore as unknown as ReturnType<typeof vi.fn>
     ).mockReturnValue({
       setOptimisticUserMessage: vi.fn(),
       getOptimisticUserMessage: vi.fn(() => null),

frontend/src/components/features/chat/chat-interface.tsx

@@ -22,7 +22,7 @@ import { useAgentStore } from "#/stores/agent-store";
 import { ScrollToBottomButton } from "#/components/shared/buttons/scroll-to-bottom-button";
 import { LoadingSpinner } from "#/components/shared/loading-spinner";
 import { displayErrorToast } from "#/utils/custom-toast-handlers";
-import { useOptimisticUserMessage } from "#/hooks/use-optimistic-user-message";
+import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store";
 import { useWSErrorMessage } from "#/hooks/use-ws-error-message";
 import { ErrorMessageBanner } from "./error-message-banner";
 import {
@@ -49,7 +49,7 @@ export function ChatInterface() {
   const { getErrorMessage } = useWSErrorMessage();
   const { send, isLoadingMessages, parsedEvents } = useWsClient();
   const { setOptimisticUserMessage, getOptimisticUserMessage } =
-    useOptimisticUserMessage();
+    useOptimisticUserMessageStore();
   const { t } = useTranslation();
   const scrollRef = React.useRef<HTMLDivElement>(null);
   const {

frontend/src/components/features/chat/messages.tsx

@@ -12,7 +12,7 @@ import {
 } from "#/types/core/guards";
 import { EventMessage } from "./event-message";
 import { ChatMessage } from "./chat-message";
-import { useOptimisticUserMessage } from "#/hooks/use-optimistic-user-message";
+import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store";
 import { LaunchMicroagentModal } from "./microagent/launch-microagent-modal";
 import { useUserConversation } from "#/hooks/query/use-user-conversation";
 import { useConversationId } from "#/hooks/use-conversation-id";
@@ -48,7 +48,7 @@ export const Messages: React.FC<MessagesProps> = React.memo(
       isPending,
       unsubscribeFromConversation,
     } = useCreateConversationAndSubscribeMultiple();
-    const { getOptimisticUserMessage } = useOptimisticUserMessage();
+    const { getOptimisticUserMessage } = useOptimisticUserMessageStore();
     const { conversationId } = useConversationId();
     const { data: conversation } = useUserConversation(conversationId);
 

frontend/src/components/features/home/tasks/task-card.tsx

@@ -4,7 +4,7 @@ import { SuggestedTask } from "#/utils/types";
 import { useIsCreatingConversation } from "#/hooks/use-is-creating-conversation";
 import { useCreateConversation } from "#/hooks/mutation/use-create-conversation";
 import { TaskIssueNumber } from "./task-issue-number";
-import { useOptimisticUserMessage } from "#/hooks/use-optimistic-user-message";
+import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store";
 import { cn } from "#/utils/utils";
 
 const getTaskTypeMap = (
@@ -21,7 +21,7 @@ interface TaskCardProps {
 }
 
 export function TaskCard({ task }: TaskCardProps) {
-  const { setOptimisticUserMessage } = useOptimisticUserMessage();
+  const { setOptimisticUserMessage } = useOptimisticUserMessageStore();
   const { mutate: createConversation } = useCreateConversation();
   const isCreatingConversation = useIsCreatingConversation();
   const { t } = useTranslation();

frontend/src/context/ws-client-provider.tsx

@@ -26,7 +26,7 @@ import {
   isStatusUpdate,
   isUserMessage,
 } from "#/types/core/guards";
-import { useOptimisticUserMessage } from "#/hooks/use-optimistic-user-message";
+import { useOptimisticUserMessageStore } from "#/stores/optimistic-user-message-store";
 import { useWSErrorMessage } from "#/hooks/use-ws-error-message";
 
 export type WebSocketStatus = "CONNECTING" | "CONNECTED" | "DISCONNECTED";
@@ -131,7 +131,7 @@ export function WsClientProvider({
   conversationId,
   children,
 }: React.PropsWithChildren<WsClientProviderProps>) {
-  const { removeOptimisticUserMessage } = useOptimisticUserMessage();
+  const { removeOptimisticUserMessage } = useOptimisticUserMessageStore();
   const { setErrorMessage, removeErrorMessage } = useWSErrorMessage();
   const queryClient = useQueryClient();
   const sioRef = React.useRef<Socket | null>(null);

frontend/src/hooks/use-optimistic-user-message.ts

@@ -1,23 +0,0 @@
-import { useQueryClient } from "@tanstack/react-query";
-
-export const useOptimisticUserMessage = () => {
-  const queryKey = ["optimistic_user_message"] as const;
-  const queryClient = useQueryClient();
-
-  const setOptimisticUserMessage = (message: string) => {
-    queryClient.setQueryData<string>(queryKey, message);
-  };
-
-  const getOptimisticUserMessage = () =>
-    queryClient.getQueryData<string>(queryKey);
-
-  const removeOptimisticUserMessage = () => {
-    queryClient.removeQueries({ queryKey });
-  };
-
-  return {
-    setOptimisticUserMessage,
-    getOptimisticUserMessage,
-    removeOptimisticUserMessage,
-  };
-};

frontend/src/stores/optimistic-user-message-store.ts

@@ -0,0 +1,36 @@
+import { create } from "zustand";
+
+interface OptimisticUserMessageState {
+  optimisticUserMessage: string | null;
+}
+
+interface OptimisticUserMessageActions {
+  setOptimisticUserMessage: (message: string) => void;
+  getOptimisticUserMessage: () => string | null;
+  removeOptimisticUserMessage: () => void;
+}
+
+type OptimisticUserMessageStore = OptimisticUserMessageState &
+  OptimisticUserMessageActions;
+
+const initialState: OptimisticUserMessageState = {
+  optimisticUserMessage: null,
+};
+
+export const useOptimisticUserMessageStore = create<OptimisticUserMessageStore>(
+  (set, get) => ({
+    ...initialState,
+
+    setOptimisticUserMessage: (message: string) =>
+      set(() => ({
+        optimisticUserMessage: message,
+      })),
+
+    getOptimisticUserMessage: () => get().optimisticUserMessage,
+
+    removeOptimisticUserMessage: () =>
+      set(() => ({
+        optimisticUserMessage: null,
+      })),
+  }),
+);