feat: WebSocket connect 使用 Promise,确保 sendMessage 发生在 connect 后 (#390)

This commit is contained in:
ᴀᴍᴛᴏᴀᴇʀ
2025-07-11 20:00:15 +08:00
committed by GitHub
parent 0990a276ff
commit 2837bb5234
+55 -43
View File
@@ -43,6 +43,7 @@ export class WebSocketManager {
private errorSubscribers: Set<ErrorCallback> = new Set(); private errorSubscribers: Set<ErrorCallback> = new Set();
private subscribedEvents: Set<EventType> = new Set(); private subscribedEvents: Set<EventType> = new Set();
private connectionPromise: Promise<void> | null = null;
private constructor() {} private constructor() {}
@@ -54,42 +55,54 @@ export class WebSocketManager {
} }
// 连接 WebSocket // 连接 WebSocket
public connect(): void { public connect(): Promise<void> {
if (this.connected || this.connecting) return; if (this.connected) return Promise.resolve();
if (this.connectionPromise) return this.connectionPromise;
this.connecting = true; this.connectionPromise = new Promise((resolve, reject) => {
const token = localStorage.getItem('authToken') || ''; this.connecting = true;
const token = localStorage.getItem('authToken') || '';
try { try {
const protocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://'; const protocol = window.location.protocol === 'https:' ? 'wss://' : 'ws://';
this.socket = new WebSocket(`${protocol}${window.location.host}/api/ws`, [token]); this.socket = new WebSocket(`${protocol}${window.location.host}/api/ws`, [token]);
this.socket.onopen = () => { this.socket.onopen = () => {
this.connected = true; this.connected = true;
this.connecting = false; this.connecting = false;
this.reconnectAttempts = 0; this.reconnectAttempts = 0;
this.connectionPromise = null;
this.resubscribeEvents();
resolve();
};
this.resubscribeEvents(); this.socket.onmessage = this.handleMessage.bind(this);
};
this.socket.onmessage = this.handleMessage.bind(this); this.socket.onclose = () => {
this.connected = false;
this.connecting = false;
this.connectionPromise = null;
this.scheduleReconnect();
};
this.socket.onclose = () => { this.socket.onerror = (error) => {
this.connected = false; console.error('WebSocket error:', error);
this.connecting = false;
this.connectionPromise = null;
reject(error);
toast.error('WebSocket 连接发生错误,请检查网络或稍后重试');
};
} catch (error) {
this.connecting = false; this.connecting = false;
this.connectionPromise = null;
reject(error);
console.error('Failed to create WebSocket:', error);
toast.error('创建 WebSocket 连接失败,请检查网络或稍后重试');
this.scheduleReconnect(); this.scheduleReconnect();
}; }
});
this.socket.onerror = (error) => { return this.connectionPromise;
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();
}
} }
private handleMessage(event: MessageEvent): void { private handleMessage(event: MessageEvent): void {
@@ -111,14 +124,13 @@ export class WebSocketManager {
} }
} }
private sendMessage(message: ClientEvent): void { private async sendMessage(message: ClientEvent): Promise<void> {
if (!this.connected || !this.socket) { if (!this.connected) {
console.warn('Cannot send message: WebSocket not connected'); await this.connect();
return;
} }
try { try {
this.socket.send(JSON.stringify(message)); this.socket!.send(JSON.stringify(message));
} catch (error) { } catch (error) {
console.error('Failed to send message:', error); console.error('Failed to send message:', error);
toast.error('发送 WebSocket 消息失败', { toast.error('发送 WebSocket 消息失败', {
@@ -127,25 +139,27 @@ export class WebSocketManager {
} }
} }
private subscribe(eventType: EventType): void { private async subscribe(eventType: EventType): Promise<void> {
if (this.subscribedEvents.has(eventType)) return; if (this.subscribedEvents.has(eventType)) return;
this.sendMessage({ subscribe: eventType }); await this.sendMessage({ subscribe: eventType });
this.subscribedEvents.add(eventType); this.subscribedEvents.add(eventType);
} }
// 取消订阅事件 // 取消订阅事件
private unsubscribe(eventType: EventType): void { private async unsubscribe(eventType: EventType): Promise<void> {
if (!this.subscribedEvents.has(eventType)) return; if (!this.subscribedEvents.has(eventType)) return;
this.sendMessage({ unsubscribe: eventType }); await this.sendMessage({ unsubscribe: eventType });
this.subscribedEvents.delete(eventType); this.subscribedEvents.delete(eventType);
} }
private resubscribeEvents(): void { private async resubscribeEvents(): Promise<void> {
for (const eventType of this.subscribedEvents) { await Promise.all(
this.sendMessage({ subscribe: eventType }); Array.from(this.subscribedEvents).map(async (eventType) => {
} await this.sendMessage({ subscribe: eventType });
})
);
} }
private scheduleReconnect(): void { private scheduleReconnect(): void {
@@ -168,7 +182,6 @@ export class WebSocketManager {
} }
public subscribeToLogs(callback: LogsCallback): () => void { public subscribeToLogs(callback: LogsCallback): () => void {
this.connect();
this.logsSubscribers.add(callback); this.logsSubscribers.add(callback);
if (this.logsSubscribers.size === 1) { if (this.logsSubscribers.size === 1) {
@@ -185,7 +198,6 @@ export class WebSocketManager {
// 订阅任务状态 // 订阅任务状态
public subscribeToTasks(callback: TasksCallback): () => void { public subscribeToTasks(callback: TasksCallback): () => void {
this.connect();
this.tasksSubscribers.add(callback); this.tasksSubscribers.add(callback);
if (this.tasksSubscribers.size === 1) { if (this.tasksSubscribers.size === 1) {
@@ -201,7 +213,6 @@ export class WebSocketManager {
} }
public subscribeToSysInfo(callback: SysInfoCallback): () => void { public subscribeToSysInfo(callback: SysInfoCallback): () => void {
this.connect();
this.sysInfoSubscribers.add(callback); this.sysInfoSubscribers.add(callback);
if (this.sysInfoSubscribers.size === 1) { if (this.sysInfoSubscribers.size === 1) {
@@ -259,6 +270,7 @@ export class WebSocketManager {
this.connected = false; this.connected = false;
this.connecting = false; this.connecting = false;
this.connectionPromise = null;
this.subscribedEvents.clear(); this.subscribedEvents.clear();
} }
} }