feat: handle websocket disconnect gracefully

Noza23 · Apr 16, 2024 · 3decf6a · 3decf6a
1 parent 96d49b6
commit 3decf6a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/backend/routers/inference.py b/backend/routers/inference.py
@@ -1,6 +1,6 @@
 from typing import Union, Any
 
-from fastapi import APIRouter, UploadFile, WebSocket
+from fastapi import APIRouter, UploadFile, WebSocket, WebSocketDisconnect
 from fastapi import Depends, File
 from fastapi import HTTPException, WebSocketException
 from redis import asyncio as aioredis  # type: ignore
@@ -114,8 +114,11 @@ async def inference_ws(
         if len(myotubes) + len(nucleis) == 0:
             await websocket.close()
             break
-
-        data = await websocket.receive_json()
+        try:
+            data = await websocket.receive_json()
+        except WebSocketDisconnect:
+            print("Websocket disconnected.")
+            break
         try:
             point = Point.model_validate(data)
         except ValidationError:

diff --git a/backend/routers/validation.py b/backend/routers/validation.py
@@ -1,5 +1,5 @@
 from fastapi import APIRouter, UploadFile, Depends, WebSocket
-from fastapi import HTTPException, WebSocketException
+from fastapi import HTTPException, WebSocketException, WebSocketDisconnect
 from redis import asyncio as aioredis  # type: ignore
 
 from myo_sam.inference.pipeline import Pipeline
@@ -87,7 +87,6 @@ async def validation_ws(
     """Websocket for validation mode."""
 
     await websocket.accept()
-    print("Hash Received: ", hash_str)
 
     mo = Myotubes.model_validate_json(
         await redis.get(KEYS.result_key(hash_str))
@@ -98,7 +97,7 @@ async def validation_ws(
     i = state.get_next()
 
     if state.done:
-        websocket.close()
+        await websocket.close()
         return
 
     await websocket.send_json(
@@ -110,7 +109,14 @@ async def validation_ws(
     )
 
     while True:
-        data = int(await websocket.receive_text())
+        if state.done:
+            break
+        try:
+            data = int(await websocket.receive_text())
+        except WebSocketDisconnect:
+            print("Websocket disconnected.")
+            break
+
         if data == 0:
             state.invalid.add(i)
         elif data == 1:
@@ -128,15 +134,14 @@ async def validation_ws(
         step = data != 2 if data != -1 else -1
         i = min(max(i + step, 0), len(mo) - 1)
 
-        if i == len(mo):
+        if i + 1 == len(mo):
             state.done = True
             await redis.mset(
                 {
                     KEYS.state_key(hash_str): state.model_dump_json(),
                     KEYS.result_key(hash_str): mo.model_dump_json(),
                 }
             )
-            break
 
         # Send next contour
         await websocket.send_json(