Feature: Add ability to reconnect websockets

frontend/src/components/AgentControlBar.tsx

@@ -3,6 +3,7 @@ import React from "react";
 import { useSelector } from "react-redux";
 import PauseIcon from "#/assets/pause";
 import PlayIcon from "#/assets/play";
+import Refresh from "#/assets/refresh.svg?react";
 import { generateAgentStateChangeEvent } from "#/services/agentStateService";
 import { RootState } from "#/store";
 import AgentState from "#/types/AgentState";
@@ -36,23 +37,21 @@ const IgnoreTaskStateMap: Record<string, AgentState[]> = {
 interface ActionButtonProps {
   isDisabled?: boolean;
   content: string;
-  action: AgentState;
-  handleAction: (action: AgentState) => void;
+  handleClick: () => void;
   large?: boolean;
 }
 
 function ActionButton({
   isDisabled = false,
   content,
-  action,
-  handleAction,
+  handleClick,
   children,
   large = false,
 }: React.PropsWithChildren<ActionButtonProps>) {
   return (
     <Tooltip content={content} closeDelay={100}>
       <button
-        onClick={() => handleAction(action)}
+        onClick={handleClick}
         disabled={isDisabled}
         className={`
           relative overflow-visible cursor-default hover:cursor-pointer group
@@ -74,15 +73,32 @@ function ActionButton({
 function AgentControlBar() {
   const { send } = useSocket();
   const { curAgentState } = useSelector((state: RootState) => state.agent);
+  const socket = useSocket();
 
-  const handleAction = (action: AgentState) => {
+  const handleAction = () => {
+    const action =
+      curAgentState === AgentState.PAUSED
+        ? AgentState.RUNNING
+        : AgentState.PAUSED;
     if (!IgnoreTaskStateMap[action].includes(curAgentState)) {
       send(generateAgentStateChangeEvent(action));
     }
   };
 
-  return (
-    <div className="flex justify-between items-center gap-20">
+  const renderButton = () => {
+    if (!socket.isConnected) {
+      return (
+        <ActionButton
+          content="Reconnect"
+          handleClick={() => window.location.reload()}
+          large
+        >
+          <Refresh width={20} height={20} />
+        </ActionButton>
+      );
+    }
+
+    return (
       <ActionButton
         isDisabled={
           curAgentState !== AgentState.RUNNING &&
@@ -93,16 +109,17 @@ function AgentControlBar() {
             ? "Resume the agent task"
             : "Pause the current task"
         }
-        action={
-          curAgentState === AgentState.PAUSED
-            ? AgentState.RUNNING
-            : AgentState.PAUSED
-        }
-        handleAction={handleAction}
+        handleClick={handleAction}
         large
       >
         {curAgentState === AgentState.PAUSED ? <PlayIcon /> : <PauseIcon />}
       </ActionButton>
+    );
+  };
+
+  return (
+    <div className="flex justify-between items-center gap-20">
+      {renderButton()}
     </div>
   );
 }

frontend/src/components/AgentStatusBar.tsx

@@ -5,6 +5,7 @@ import { I18nKey } from "#/i18n/declaration";
 import { RootState } from "#/store";
 import AgentState from "#/types/AgentState";
 import beep from "#/utils/beep";
+import { useSocket } from "#/context/socket";
 
 enum IndicatorColor {
   BLUE = "bg-blue-500",
@@ -19,6 +20,7 @@ function AgentStatusBar() {
   const { t } = useTranslation();
   const { curAgentState } = useSelector((state: RootState) => state.agent);
   const { curStatusMessage } = useSelector((state: RootState) => state.status);
+  const socket = useSocket();
 
   const AgentStatusMap: {
     [k: string]: { message: string; indicator: IndicatorColor };
@@ -92,8 +94,15 @@ function AgentStatusBar() {
   }, [curAgentState]);
 
   const [statusMessage, setStatusMessage] = React.useState<string>("");
+  const indicator = socket.isConnected
+    ? AgentStatusMap[curAgentState].indicator
+    : IndicatorColor.RED;
 
   React.useEffect(() => {
+    if (!socket.isConnected) {
+      setStatusMessage("Disconnected");
+      return;
+    }
     if (curAgentState === AgentState.LOADING) {
       const trimmedCustomMessage = curStatusMessage.status.trim();
       if (trimmedCustomMessage) {
@@ -102,14 +111,12 @@ function AgentStatusBar() {
       }
     }
     setStatusMessage(AgentStatusMap[curAgentState].message);
-  }, [curAgentState, curStatusMessage.status]);
+  }, [curAgentState, curStatusMessage.status, socket.isConnected]);
 
   return (
     <div className="flex flex-col items-center">
       <div className="flex items-center bg-neutral-800 px-2 py-1 text-gray-400 rounded-[100px] text-sm gap-[6px]">
-        <div
-          className={`w-2 h-2 rounded-full animate-pulse ${AgentStatusMap[curAgentState].indicator}`}
-        />
+        <div className={`w-2 h-2 rounded-full animate-pulse ${indicator}`} />
         <span className="text-sm text-stone-400">{statusMessage}</span>
       </div>
     </div>

frontend/src/context/socket.tsx

@@ -92,6 +92,7 @@ function SocketProvider({ children }: SocketProviderProps) {
   const send = React.useCallback(
     (data: string | ArrayBufferLike | Blob | ArrayBufferView) => {
       if (!wsRef.current) {
+        // maybe reconnect with exponential backoff here?
         EventLogger.error("WebSocket is not connected.");
         return;
       }