Feat Tightening up Timeouts and interrupt conditions. (#3926)

openhands/server/mock/listen.py

@@ -2,6 +2,7 @@ import uvicorn
 from fastapi import FastAPI, WebSocket
 
 from openhands.core.schema import ActionType
+from openhands.runtime.utils.shutdown_listener import should_continue
 
 app = FastAPI()
 
@@ -15,7 +16,7 @@ async def websocket_endpoint(websocket: WebSocket):
     )
 
     try:
-        while True:
+        while should_continue():
             # receive message
             data = await websocket.receive_json()
             print(f'Received message: {data}')

openhands/server/session/manager.py

@@ -5,6 +5,7 @@ from fastapi import WebSocket
 
 from openhands.core.config import AppConfig
 from openhands.core.logger import openhands_logger as logger
+from openhands.runtime.utils.shutdown_listener import should_continue
 from openhands.server.session.session import Session
 from openhands.storage.files import FileStore
 
@@ -47,7 +48,7 @@ class SessionManager:
         return await self.send(sid, {'message': message})
 
     async def _cleanup_sessions(self):
-        while True:
+        while should_continue():
             current_time = time.time()
             session_ids_to_remove = []
             for sid, session in list(self._sessions.items()):

openhands/server/session/session.py

@@ -20,6 +20,7 @@ from openhands.events.observation import (
 from openhands.events.serialization import event_from_dict, event_to_dict
 from openhands.events.stream import EventStreamSubscriber
 from openhands.llm.llm import LLM
+from openhands.runtime.utils.shutdown_listener import should_continue
 from openhands.server.session.agent import AgentSession
 from openhands.storage.files import FileStore
 
@@ -53,7 +54,7 @@ class Session:
         try:
             if self.websocket is None:
                 return
-            while True:
+            while should_continue():
                 try:
                     data = await self.websocket.receive_json()
                 except ValueError: