diff --git a/web/src/lib/ws.ts b/web/src/lib/ws.ts index b730ae9..d0b4d98 100644 --- a/web/src/lib/ws.ts +++ b/web/src/lib/ws.ts @@ -43,6 +43,7 @@ export class WebSocketManager { private errorSubscribers: Set = new Set(); private subscribedEvents: Set = new Set(); + private connectionPromise: Promise | null = null; private constructor() {} @@ -54,42 +55,54 @@ export class WebSocketManager { } // 连接 WebSocket - public connect(): void { - if (this.connected || this.connecting) return; + public connect(): Promise { + if (this.connected) return Promise.resolve(); + if (this.connectionPromise) return this.connectionPromise; - this.connecting = true; - const token = localStorage.getItem('authToken') || ''; + this.connectionPromise = new Promise((resolve, reject) => { + this.connecting = true; + const token = localStorage.getItem('authToken') || ''; - try { - const protocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://'; - this.socket = new WebSocket(`${protocol}${window.location.host}/api/ws`, [token]); + try { + const protocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://'; + this.socket = new WebSocket(`${protocol}${window.location.host}/api/ws`, [token]); - this.socket.onopen = () => { - this.connected = true; - this.connecting = false; - this.reconnectAttempts = 0; - - this.resubscribeEvents(); - }; - - this.socket.onmessage = this.handleMessage.bind(this); - - this.socket.onclose = () => { - this.connected = false; + this.socket.onopen = () => { + this.connected = true; + this.connecting = false; + this.reconnectAttempts = 0; + this.connectionPromise = null; + this.resubscribeEvents(); + resolve(); + }; + + this.socket.onmessage = this.handleMessage.bind(this); + + this.socket.onclose = () => { + this.connected = false; + this.connecting = false; + this.connectionPromise = null; + this.scheduleReconnect(); + }; + + this.socket.onerror = (error) => { + console.error('WebSocket error:', error); + this.connecting = false; + this.connectionPromise = null; + reject(error); + toast.error('WebSocket 连接发生错误,请检查网络或稍后重试'); + }; + } catch (error) { this.connecting = false; + this.connectionPromise = null; + reject(error); + console.error('Failed to create WebSocket:', error); + toast.error('创建 WebSocket 连接失败,请检查网络或稍后重试'); this.scheduleReconnect(); - }; + } + }); - this.socket.onerror = (error) => { - console.error('WebSocket error:', error); - toast.error('WebSocket 连接发生错误,请检查网络或稍后重试'); - }; - } catch (error) { - this.connecting = false; - console.error('Failed to create WebSocket:', error); - toast.error('创建 WebSocket 连接失败,请检查网络或稍后重试'); - this.scheduleReconnect(); - } + return this.connectionPromise; } private handleMessage(event: MessageEvent): void { @@ -111,14 +124,13 @@ export class WebSocketManager { } } - private sendMessage(message: ClientEvent): void { - if (!this.connected || !this.socket) { - console.warn('Cannot send message: WebSocket not connected'); - return; + private async sendMessage(message: ClientEvent): Promise { + if (!this.connected) { + await this.connect(); } try { - this.socket.send(JSON.stringify(message)); + this.socket!.send(JSON.stringify(message)); } catch (error) { console.error('Failed to send message:', error); toast.error('发送 WebSocket 消息失败', { @@ -127,25 +139,27 @@ export class WebSocketManager { } } - private subscribe(eventType: EventType): void { + private async subscribe(eventType: EventType): Promise { if (this.subscribedEvents.has(eventType)) return; - this.sendMessage({ subscribe: eventType }); + await this.sendMessage({ subscribe: eventType }); this.subscribedEvents.add(eventType); } // 取消订阅事件 - private unsubscribe(eventType: EventType): void { + private async unsubscribe(eventType: EventType): Promise { if (!this.subscribedEvents.has(eventType)) return; - this.sendMessage({ unsubscribe: eventType }); + await this.sendMessage({ unsubscribe: eventType }); this.subscribedEvents.delete(eventType); } - private resubscribeEvents(): void { - for (const eventType of this.subscribedEvents) { - this.sendMessage({ subscribe: eventType }); - } + private async resubscribeEvents(): Promise { + await Promise.all( + Array.from(this.subscribedEvents).map(async (eventType) => { + await this.sendMessage({ subscribe: eventType }); + }) + ); } private scheduleReconnect(): void { @@ -168,7 +182,6 @@ export class WebSocketManager { } public subscribeToLogs(callback: LogsCallback): () => void { - this.connect(); this.logsSubscribers.add(callback); if (this.logsSubscribers.size === 1) { @@ -185,7 +198,6 @@ export class WebSocketManager { // 订阅任务状态 public subscribeToTasks(callback: TasksCallback): () => void { - this.connect(); this.tasksSubscribers.add(callback); if (this.tasksSubscribers.size === 1) { @@ -201,7 +213,6 @@ export class WebSocketManager { } public subscribeToSysInfo(callback: SysInfoCallback): () => void { - this.connect(); this.sysInfoSubscribers.add(callback); if (this.sysInfoSubscribers.size === 1) { @@ -259,6 +270,7 @@ export class WebSocketManager { this.connected = false; this.connecting = false; + this.connectionPromise = null; this.subscribedEvents.clear(); } }