diff --git a/frontend/__tests__/components/chat/chat-interface.test.tsx b/frontend/__tests__/components/chat/chat-interface.test.tsx index c018e010e9d9..4a5c80b26e22 100644 --- a/frontend/__tests__/components/chat/chat-interface.test.tsx +++ b/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 + useOptimisticUserMessageStore as unknown as ReturnType ).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 + useOptimisticUserMessageStore as unknown as ReturnType ).mockReturnValue({ setOptimisticUserMessage: vi.fn(), getOptimisticUserMessage: vi.fn(() => "Optimistic message"), @@ -246,7 +246,7 @@ describe("ChatInterface - Empty state", () => { parsedEvents: [], }); ( - useOptimisticUserMessage as unknown as ReturnType + useOptimisticUserMessageStore as unknown as ReturnType ).mockReturnValue({ setOptimisticUserMessage: vi.fn(), getOptimisticUserMessage: vi.fn(() => null), diff --git a/frontend/src/components/features/chat/chat-interface.tsx b/frontend/src/components/features/chat/chat-interface.tsx index b7217594b87d..a61201d64779 100644 --- a/frontend/src/components/features/chat/chat-interface.tsx +++ b/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(null); const { diff --git a/frontend/src/components/features/chat/messages.tsx b/frontend/src/components/features/chat/messages.tsx index 6187583506e3..289a6bdb58f1 100644 --- a/frontend/src/components/features/chat/messages.tsx +++ b/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 = React.memo( isPending, unsubscribeFromConversation, } = useCreateConversationAndSubscribeMultiple(); - const { getOptimisticUserMessage } = useOptimisticUserMessage(); + const { getOptimisticUserMessage } = useOptimisticUserMessageStore(); const { conversationId } = useConversationId(); const { data: conversation } = useUserConversation(conversationId); diff --git a/frontend/src/components/features/home/tasks/task-card.tsx b/frontend/src/components/features/home/tasks/task-card.tsx index 38b361acae64..39652e880c8d 100644 --- a/frontend/src/components/features/home/tasks/task-card.tsx +++ b/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(); diff --git a/frontend/src/context/ws-client-provider.tsx b/frontend/src/context/ws-client-provider.tsx index 77a02e2d3997..9fa523206456 100644 --- a/frontend/src/context/ws-client-provider.tsx +++ b/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) { - const { removeOptimisticUserMessage } = useOptimisticUserMessage(); + const { removeOptimisticUserMessage } = useOptimisticUserMessageStore(); const { setErrorMessage, removeErrorMessage } = useWSErrorMessage(); const queryClient = useQueryClient(); const sioRef = React.useRef(null); diff --git a/frontend/src/hooks/use-optimistic-user-message.ts b/frontend/src/hooks/use-optimistic-user-message.ts deleted file mode 100644 index 33cbad7d9331..000000000000 --- a/frontend/src/hooks/use-optimistic-user-message.ts +++ /dev/null @@ -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(queryKey, message); - }; - - const getOptimisticUserMessage = () => - queryClient.getQueryData(queryKey); - - const removeOptimisticUserMessage = () => { - queryClient.removeQueries({ queryKey }); - }; - - return { - setOptimisticUserMessage, - getOptimisticUserMessage, - removeOptimisticUserMessage, - }; -}; diff --git a/frontend/src/stores/optimistic-user-message-store.ts b/frontend/src/stores/optimistic-user-message-store.ts new file mode 100644 index 000000000000..2b6e40e90498 --- /dev/null +++ b/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( + (set, get) => ({ + ...initialState, + + setOptimisticUserMessage: (message: string) => + set(() => ({ + optimisticUserMessage: message, + })), + + getOptimisticUserMessage: () => get().optimisticUserMessage, + + removeOptimisticUserMessage: () => + set(() => ({ + optimisticUserMessage: null, + })), + }), +);