From 151c2895e05218dcf302018b05158b43591ead33 Mon Sep 17 00:00:00 2001
From: Hiep Le <69354317+hieptl@users.noreply.github.com>
Date: Thu, 20 Nov 2025 01:28:17 +0700
Subject: [PATCH] feat(frontend): disable change-agent button until WebSocket
 connection is ready (#11781)

---
 .../features/chat/change-agent-button.tsx     | 23 +++++++++++++------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/frontend/src/components/features/chat/change-agent-button.tsx b/frontend/src/components/features/chat/change-agent-button.tsx
index 6d41f5cfc1..e609c30b37 100644
--- a/frontend/src/components/features/chat/change-agent-button.tsx
+++ b/frontend/src/components/features/chat/change-agent-button.tsx
@@ -14,6 +14,7 @@ import { AgentState } from "#/types/agent-state";
 import { useActiveConversation } from "#/hooks/query/use-active-conversation";
 import { useCreateConversation } from "#/hooks/mutation/use-create-conversation";
 import { displaySuccessToast } from "#/utils/custom-toast-handlers";
+import { useUnifiedWebSocketStatus } from "#/hooks/use-unified-websocket-status";
 
 export function ChangeAgentButton() {
   const [contextMenuOpen, setContextMenuOpen] = useState<boolean>(false);
@@ -26,6 +27,10 @@ export function ChangeAgentButton() {
     (state) => state.setConversationMode,
   );
 
+  const webSocketStatus = useUnifiedWebSocketStatus();
+
+  const isWebSocketConnected = webSocketStatus === "CONNECTED";
+
   const shouldUsePlanningAgent = USE_PLANNING_AGENT();
 
   const { curAgentState } = useAgentState();
@@ -40,10 +45,10 @@ export function ChangeAgentButton() {
 
   // Close context menu when agent starts running
   useEffect(() => {
-    if (isAgentRunning && contextMenuOpen) {
+    if ((isAgentRunning || !isWebSocketConnected) && contextMenuOpen) {
       setContextMenuOpen(false);
     }
-  }, [isAgentRunning, contextMenuOpen]);
+  }, [isAgentRunning, contextMenuOpen, isWebSocketConnected]);
 
   const handlePlanClick = (
     event: React.MouseEvent<HTMLButtonElement> | KeyboardEvent,
@@ -79,9 +84,15 @@ export function ChangeAgentButton() {
     );
   };
 
+  const isButtonDisabled =
+    isAgentRunning ||
+    isCreatingConversation ||
+    !isWebSocketConnected ||
+    !shouldUsePlanningAgent;
+
   // Handle Shift + Tab keyboard shortcut to cycle through modes
   useEffect(() => {
-    if (!shouldUsePlanningAgent || isAgentRunning) {
+    if (isButtonDisabled) {
       return undefined;
     }
 
@@ -108,10 +119,10 @@ export function ChangeAgentButton() {
       document.removeEventListener("keydown", handleKeyDown);
     };
   }, [
-    shouldUsePlanningAgent,
-    isAgentRunning,
+    isButtonDisabled,
     conversationMode,
     setConversationMode,
+    handlePlanClick,
   ]);
 
   const handleButtonClick = (event: React.MouseEvent<HTMLButtonElement>) => {
@@ -142,8 +153,6 @@ export function ChangeAgentButton() {
     return <LessonPlanIcon width={18} height={18} color="#ffffff" />;
   }, [isExecutionAgent]);
 
-  const isButtonDisabled = isAgentRunning || isCreatingConversation;
-
   if (!shouldUsePlanningAgent) {
     return null;
   }