feat: Update WebSocket endpoint for checkpoint progress and adjust related routes

This commit is contained in:
Will Miao
2025-07-02 18:38:02 +08:00
parent d7cb546c5f
commit 971cd56a4a
3 changed files with 1 additions and 18 deletions

View File

@@ -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)

View File

@@ -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:

View File

@@ -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);