From 971cd56a4a78c5c2ef57eb528f8a215c5a604f83 Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Wed, 2 Jul 2025 18:38:02 +0800
Subject: [PATCH] feat: Update WebSocket endpoint for checkpoint progress and
 adjust related routes

---
 py/routes/checkpoints_routes.py                 |  3 ---
 py/services/websocket_manager.py                | 14 --------------
 static/js/managers/CheckpointDownloadManager.js |  2 +-
 3 files changed, 1 insertion(+), 18 deletions(-)

diff --git a/py/routes/checkpoints_routes.py b/py/routes/checkpoints_routes.py
index 321f82d6..873139f2 100644
--- a/py/routes/checkpoints_routes.py
+++ b/py/routes/checkpoints_routes.py
@@ -57,9 +57,6 @@ class CheckpointsRoutes:
         app.router.add_post('/api/checkpoints/download', self.download_checkpoint)
         app.router.add_post('/api/checkpoints/save-metadata', self.save_metadata) # Add new route
         app.router.add_post('/api/checkpoints/rename', self.rename_checkpoint)  # Add new rename endpoint
-        
-        # Add new WebSocket endpoint for checkpoint progress
-        app.router.add_get('/ws/checkpoint-progress', ws_manager.handle_checkpoint_connection)
 
         # Add new routes for finding duplicates and filename conflicts
         app.router.add_get('/api/checkpoints/find-duplicates', self.find_duplicate_checkpoints)
diff --git a/py/services/websocket_manager.py b/py/services/websocket_manager.py
index c85aa3a2..1887ee1f 100644
--- a/py/services/websocket_manager.py
+++ b/py/services/websocket_manager.py
@@ -40,20 +40,6 @@ class WebSocketManager:
             self._init_websockets.discard(ws)
         return ws
         
-    async def handle_checkpoint_connection(self, request: web.Request) -> web.WebSocketResponse:
-        """Handle new WebSocket connection for checkpoint download progress"""
-        ws = web.WebSocketResponse()
-        await ws.prepare(request)
-        self._checkpoint_websockets.add(ws)
-        
-        try:
-            async for msg in ws:
-                if msg.type == web.WSMsgType.ERROR:
-                    logger.error(f'Checkpoint WebSocket error: {ws.exception()}')
-        finally:
-            self._checkpoint_websockets.discard(ws)
-        return ws
-        
     async def broadcast(self, data: Dict):
         """Broadcast message to all connected clients"""
         if not self._websockets:
diff --git a/static/js/managers/CheckpointDownloadManager.js b/static/js/managers/CheckpointDownloadManager.js
index 76ead14e..0f69a484 100644
--- a/static/js/managers/CheckpointDownloadManager.js
+++ b/static/js/managers/CheckpointDownloadManager.js
@@ -303,7 +303,7 @@ export class CheckpointDownloadManager {
 
             // Setup WebSocket for progress updates using checkpoint-specific endpoint
             const wsProtocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
-            const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/checkpoint-progress`);
+            const ws = new WebSocket(`${wsProtocol}${window.location.host}/ws/fetch-progress`);
             
             ws.onmessage = (event) => {
                 const data = JSON.parse(event.data);