Revert "Feature: Add ability to reconnect websockets" (All-Hands-AI#4801

)

frontend/src/context/socket.tsx

@@ -2,11 +2,9 @@ import React from "react";
 import { Data } from "ws";
 import EventLogger from "#/utils/event-logger";
 
-const RECONNECT_RETRIES = 5;
-
 interface WebSocketClientOptions {
   token: string | null;
-  onOpen?: (event: Event, isNewSession: boolean) => void;
+  onOpen?: (event: Event) => void;
   onMessage?: (event: MessageEvent<Data>) => void;
   onError?: (event: Event) => void;
   onClose?: (event: Event) => void;
@@ -16,8 +14,8 @@ interface WebSocketContextType {
   send: (data: string | ArrayBufferLike | Blob | ArrayBufferView) => void;
   start: (options?: WebSocketClientOptions) => void;
   stop: () => void;
-  setRuntimeIsInitialized: (runtimeIsInitialized: boolean) => void;
-  runtimeIsInitialized: boolean;
+  setRuntimeIsInitialized: () => void;
+  runtimeActive: boolean;
   isConnected: boolean;
   events: Record<string, unknown>[];
 }
@@ -32,11 +30,14 @@ interface SocketProviderProps {
 
 function SocketProvider({ children }: SocketProviderProps) {
   const wsRef = React.useRef<WebSocket | null>(null);
-  const wsReconnectRetries = React.useRef<number>(RECONNECT_RETRIES);
   const [isConnected, setIsConnected] = React.useState(false);
-  const [runtimeIsInitialized, setRuntimeIsInitialized] = React.useState(false);
+  const [runtimeActive, setRuntimeActive] = React.useState(false);
   const [events, setEvents] = React.useState<Record<string, unknown>[]>([]);
 
+  const setRuntimeIsInitialized = () => {
+    setRuntimeActive(true);
+  };
+
   const start = React.useCallback((options?: WebSocketClientOptions): void => {
     if (wsRef.current) {
       EventLogger.warning(
@@ -58,9 +59,7 @@ function SocketProvider({ children }: SocketProviderProps) {
 
     ws.addEventListener("open", (event) => {
       setIsConnected(true);
-      const isNewSession = sessionToken === "NO_JWT";
-      wsReconnectRetries.current = RECONNECT_RETRIES;
-      options?.onOpen?.(event, isNewSession);
+      options?.onOpen?.(event);
     });
 
     ws.addEventListener("message", (event) => {
@@ -77,22 +76,17 @@ function SocketProvider({ children }: SocketProviderProps) {
 
     ws.addEventListener("close", (event) => {
       EventLogger.event(event, "SOCKET CLOSE");
+
       setIsConnected(false);
-      setRuntimeIsInitialized(false);
+      setRuntimeActive(false);
       wsRef.current = null;
       options?.onClose?.(event);
-      if (wsReconnectRetries.current) {
-        wsReconnectRetries.current -= 1;
-        const token = localStorage.getItem("token");
-        setTimeout(() => start({ ...(options || {}), token }), 1);
-      }
     });
 
     wsRef.current = ws;
   }, []);
 
   const stop = React.useCallback((): void => {
-    wsReconnectRetries.current = 0;
     if (wsRef.current) {
       wsRef.current.close();
       wsRef.current = null;
@@ -117,7 +111,7 @@ function SocketProvider({ children }: SocketProviderProps) {
       start,
       stop,
       setRuntimeIsInitialized,
-      runtimeIsInitialized,
+      runtimeActive,
       isConnected,
       events,
     }),
@@ -126,7 +120,7 @@ function SocketProvider({ children }: SocketProviderProps) {
       start,
       stop,
       setRuntimeIsInitialized,
-      runtimeIsInitialized,
+      runtimeActive,
       isConnected,
       events,
     ],

frontend/src/routes/_oh.app.tsx

@@ -124,8 +124,7 @@ function App() {
   const { files, importedProjectZip } = useSelector(
     (state: RootState) => state.initalQuery,
   );
-  const { start, send, setRuntimeIsInitialized, runtimeIsInitialized } =
-    useSocket();
+  const { start, send, setRuntimeIsInitialized, runtimeActive } = useSocket();
   const { settings, token, ghToken, repo, q, lastCommit } =
     useLoaderData<typeof clientLoader>();
   const fetcher = useFetcher();
@@ -162,32 +161,21 @@ function App() {
     );
   };
 
-  const doSendInitialQuery = React.useRef<boolean>(true);
-
   const sendInitialQuery = (query: string, base64Files: string[]) => {
     const timestamp = new Date().toISOString();
     send(createChatMessage(query, base64Files, timestamp));
   };
 
-  const handleOpen = React.useCallback(
-    (event: Event, isNewSession: boolean) => {
-      if (!isNewSession) {
-        dispatch(clearMessages());
-        dispatch(clearTerminal());
-        dispatch(clearJupyter());
-      }
-      doSendInitialQuery.current = isNewSession;
-      const initEvent = {
-        action: ActionType.INIT,
-        args: settings,
-      };
-      send(JSON.stringify(initEvent));
-
-      // display query in UI, but don't send it to the server
-      if (q && isNewSession) addIntialQueryToChat(q, files);
-    },
-    [settings],
-  );
+  const handleOpen = React.useCallback(() => {
+    const initEvent = {
+      action: ActionType.INIT,
+      args: settings,
+    };
+    send(JSON.stringify(initEvent));
+
+    // display query in UI, but don't send it to the server
+    if (q) addIntialQueryToChat(q, files);
+  }, [settings]);
 
   const handleMessage = React.useCallback(
     (message: MessageEvent<WebSocket.Data>) => {
@@ -230,7 +218,7 @@ function App() {
         isAgentStateChange(parsed) &&
         parsed.extras.agent_state === AgentState.INIT
       ) {
-        setRuntimeIsInitialized(true);
+        setRuntimeIsInitialized();
 
         // handle new session
         if (!token) {
@@ -245,7 +233,7 @@ function App() {
             additionalInfo = `Files have been uploaded. Please check the /workspace for files.`;
           }
 
-          if (q && doSendInitialQuery.current) {
+          if (q) {
             if (additionalInfo) {
               sendInitialQuery(`${q}\n\n[${additionalInfo}]`, files);
             } else {
@@ -277,15 +265,15 @@ function App() {
   });
 
   React.useEffect(() => {
-    if (runtimeIsInitialized && userId && ghToken) {
+    if (runtimeActive && userId && ghToken) {
       // Export if the user valid, this could happen mid-session so it is handled here
       send(getGitHubTokenCommand(ghToken));
     }
-  }, [userId, ghToken, runtimeIsInitialized]);
+  }, [userId, ghToken, runtimeActive]);
 
   React.useEffect(() => {
     (async () => {
-      if (runtimeIsInitialized && importedProjectZip) {
+      if (runtimeActive && importedProjectZip) {
         // upload files action
         try {
           const blob = base64ToBlob(importedProjectZip);
@@ -299,7 +287,7 @@ function App() {
         }
       }
     })();
-  }, [runtimeIsInitialized, importedProjectZip]);
+  }, [runtimeActive, importedProjectZip]);
 
   const {
     isOpen: securityModalIsOpen,
@@ -315,7 +303,7 @@ function App() {
             className={cn(
               "w-2 h-2 rounded-full border",
               "absolute left-3 top-3",
-              runtimeIsInitialized
+              runtimeActive
                 ? "bg-green-800 border-green-500"
                 : "bg-red-800 border-red-500",
             )}

openhands/controller/agent_controller.py

@@ -117,7 +117,6 @@ def __init__(
             max_iterations=max_iterations,
             confirmation_mode=confirmation_mode,
         )
-
         self.max_budget_per_task = max_budget_per_task
         self.agent_to_llm_config = agent_to_llm_config if agent_to_llm_config else {}
         self.agent_configs = agent_configs if agent_configs else {}

openhands/server/listen.py

@@ -329,12 +329,11 @@ async def websocket_endpoint(websocket: WebSocket):
             await websocket.send_json({'error': 'Invalid token', 'error_code': 401})
             await websocket.close()
             return
-        logger.info(f'Existing session: {sid}')
     else:
         sid = str(uuid.uuid4())
         jwt_token = sign_token({'sid': sid}, config.jwt_secret)
-        logger.info(f'New session: {sid}')
 
+    logger.info(f'New session: {sid}')
     session = session_manager.add_or_restart_session(sid, websocket)
     await websocket.send_json({'token': jwt_token, 'status': 'ok'})
 

openhands/server/session/agent_session.py

@@ -100,9 +100,9 @@ async def _start(
         config: AppConfig,
         agent: Agent,
         max_iterations: int,
-        max_budget_per_task: float | None,
-        agent_to_llm_config: dict[str, LLMConfig] | None,
-        agent_configs: dict[str, AgentConfig] | None,
+        max_budget_per_task: float | None = None,
+        agent_to_llm_config: dict[str, LLMConfig] | None = None,
+        agent_configs: dict[str, AgentConfig] | None = None,
     ):
         self._create_security_analyzer(config.security.security_analyzer)
         await self._create_runtime(

openhands/server/session/session.py

@@ -30,7 +30,7 @@ class Session:
     sid: str
     websocket: WebSocket | None
     last_active_ts: int = 0
-    is_alive: bool = False
+    is_alive: bool = True
     agent_session: AgentSession
     loop: asyncio.AbstractEventLoop
 
@@ -109,7 +109,6 @@ async def _initialize_agent(self, data: dict):
 
         # Create the agent session
         try:
-            self.is_alive = True
             await self.agent_session.start(
                 runtime_name=self.config.runtime,
                 config=self.config,
@@ -139,7 +138,9 @@ async def on_event(self, event: Event):
             return
         if event.source == EventSource.AGENT:
             await self.send(event_to_dict(event))
-        elif event.source == EventSource.USER and isinstance(event, CmdOutputObservation):
+        elif event.source == EventSource.USER and isinstance(
+            event, CmdOutputObservation
+        ):
             await self.send(event_to_dict(event))
         # NOTE: ipython observations are not sent here currently
         elif event.source == EventSource.ENVIRONMENT and isinstance(
@@ -158,8 +159,7 @@ async def on_event(self, event: Event):
     async def dispatch(self, data: dict):
         action = data.get('action', '')
         if action == ActionType.INIT:
-            if not self.is_alive:
-                await self._initialize_agent(data)
+            await self._initialize_agent(data)
             return
         event = event_from_dict(data.copy())
         # This checks if the model supports images