优化 MessageConnect 的记忆管理及Cleanup机制

packages/message/extension_message.ts

@@ -1,4 +1,8 @@
+import EventEmitter from "eventemitter3";
 import type { Message, MessageConnect, MessageSend, RuntimeMessageSender, TMessage, TMessageCommAction } from "./types";
+import { uuidv4 } from "@App/pkg/utils/uuid";
+
+const listenerMgr = new EventEmitter<string, any>(); // 单一管理器
 
 export class ExtensionMessage implements Message {
   constructor(private backgroundPrimary = false) {}
@@ -19,8 +23,6 @@ export class ExtensionMessage implements Message {
         if (lastError) {
           console.error("chrome.runtime.lastError in chrome.runtime.sendMessage:", lastError);
           // 通信API出错不回继续对话
-          resolve = null;
-          return;
         }
         resolve!(resp);
         resolve = null;
@@ -146,25 +148,75 @@ export class ExtensionMessage implements Message {
 }
 
 export class ExtensionMessageConnect implements MessageConnect {
-  constructor(private con: chrome.runtime.Port) {}
+  private readonly listenerId = `${uuidv4()}`; // 使用 uuidv4 确保唯一
+  private con: chrome.runtime.Port | null;
+  private isSelfDisconnected = false;
+
+  constructor(con: chrome.runtime.Port) {
+    this.con = con; // 强引用
+    const handler = (msg: TMessage, _con: chrome.runtime.Port) => {
+      listenerMgr.emit(`onMessage:${this.listenerId}`, msg);
+    };
+    const cleanup = (con: chrome.runtime.Port) => {
+      if (this.con) {
+        this.con = null;
+        listenerMgr.removeAllListeners(`cleanup:${this.listenerId}`);
+        con.onMessage.removeListener(handler);
+        con.onDisconnect.removeListener(cleanup);
+        listenerMgr.emit(`onDisconnect:${this.listenerId}`, this.isSelfDisconnected);
+        listenerMgr.removeAllListeners(`onDisconnect:${this.listenerId}`);
+        listenerMgr.removeAllListeners(`onMessage:${this.listenerId}`);
+      }
+    };
+    con.onMessage.addListener(handler);
+    con.onDisconnect.addListener(cleanup);
+    listenerMgr.once(`cleanup:${this.listenerId}`, cleanup);
+  }
 
   sendMessage(data: TMessage) {
-    this.con.postMessage(data);
+    if (!this.con) {
+      console.warn("Attempted to sendMessage on a disconnected port.");
+      // 無法 sendMessage 不应该屏蔽错误
+      throw new Error("Attempted to sendMessage on a disconnected port.");
+    }
+    this.con?.postMessage(data);
   }
 
   onMessage(callback: (data: TMessage) => void) {
-    this.con.onMessage.addListener(callback);
+    if (!this.con) {
+      console.error("onMessage Invalid Port");
+      // 無法監聽的話不应该屏蔽错误
+      throw new Error("onMessage Invalid Port");
+    }
+    listenerMgr.addListener(`onMessage:${this.listenerId}`, callback);
   }
 
   disconnect() {
-    this.con.disconnect();
+    if (!this.con) {
+      console.warn("Attempted to disconnect on a disconnected port.");
+      // 重复 disconnect() 不应该屏蔽错误
+      throw new Error("Attempted to disconnect on a disconnected port.");
+    }
+    this.isSelfDisconnected = true;
+    this.con?.disconnect();
+    // Note: .disconnect() will NOT automatically trigger the 'cleanup' listener
+    listenerMgr.emit(`cleanup:${this.listenerId}`);
   }
 
-  onDisconnect(callback: () => void) {
-    this.con.onDisconnect.addListener(callback);
+  onDisconnect(callback: (isSelfDisconnected: boolean) => void) {
+    if (!this.con) {
+      console.error("onDisconnect Invalid Port");
+      // 無法監聽的話不应该屏蔽错误
+      throw new Error("onDisconnect Invalid Port");
+    }
+    listenerMgr.once(`onDisconnect:${this.listenerId}`, callback);
   }
 
   getPort(): chrome.runtime.Port {
+    if (!this.con) {
+      console.error("Port is already disconnected.");
+      throw new Error("Port is already disconnected.");
+    }
     return this.con;
   }
 }

packages/message/mock_message.ts

@@ -16,11 +16,11 @@ export class MockMessageConnect implements MessageConnect {
   }
 
   disconnect(): void {
-    this.EE.emit("disconnect");
+    this.EE.emit("disconnect", true); // MockMessageConnect 未有模拟由另一端触发 disconnect() 的情况
   }
 
-  onDisconnect(callback: () => void): void {
-    this.EE.on("disconnect", callback);
+  onDisconnect(callback: (isSelfDisconnected: boolean) => void) {
+    this.EE.once("disconnect", callback);
   }
 }
 

packages/message/server.ts

@@ -273,15 +273,15 @@ export function forwardMessage(
   senderTo: MessageSend,
   middleware?: ApiFunctionSync
 ) {
-  const handler = (params: any, fromCon: IGetSender) => {
-    const fromConnect = fromCon.getConnect();
+  const handler = async (params: any, fromCon: IGetSender): Promise<any> => {
+    const fromConnect: MessageConnect | undefined = fromCon.getConnect();
     if (fromConnect) {
-      connect(senderTo, `${prefix}/${path}`, params).then((toCon: MessageConnect) => {
-        fromConnect.onMessage(toCon.sendMessage.bind(toCon));
-        toCon.onMessage(fromConnect.sendMessage.bind(fromConnect));
-        fromConnect.onDisconnect(toCon.disconnect.bind(toCon));
-        toCon.onDisconnect(fromConnect.disconnect.bind(fromConnect));
-      });
+      const toCon: MessageConnect = await connect(senderTo, `${prefix}/${path}`, params);
+      fromConnect.onMessage(toCon.sendMessage.bind(toCon));
+      toCon.onMessage(fromConnect.sendMessage.bind(fromConnect));
+      fromConnect.onDisconnect(toCon.disconnect.bind(toCon));
+      toCon.onDisconnect(fromConnect.disconnect.bind(fromConnect));
+      return undefined;
     } else {
       return sendMessage(senderTo, prefix + "/" + path, params);
     }

packages/message/types.ts

@@ -50,7 +50,7 @@ export interface MessageConnect {
   onMessage(callback: (data: TMessage) => void): void;
   sendMessage(data: TMessage): void;
   disconnect(): void;
-  onDisconnect(callback: () => void): void;
+  onDisconnect(callback: (isSelfDisconnected: boolean) => void): void;
 }
 
 export type ExtMessageSender = {

packages/message/window_message.ts

@@ -2,6 +2,8 @@ import type { Message, MessageConnect, MessageSend, RuntimeMessageSender, TMessa
 import { uuidv4 } from "@App/pkg/utils/uuid";
 import EventEmitter from "eventemitter3";
 
+const listenerMgr = new EventEmitter<string, any>(); // 单一管理器
+
 // 通过 window.postMessage/onmessage 实现通信
 
 export interface PostMessage {
@@ -121,19 +123,41 @@ export class WindowMessage implements Message {
 }
 
 export class WindowMessageConnect implements MessageConnect {
+  private readonly listenerId = `${uuidv4()}`; // 使用 uuidv4 确保唯一
+  private target: PostMessage | null;
+  private isSelfDisconnected = false;
+
   constructor(
     private messageId: string,
-    private EE: EventEmitter<string, any>,
-    private target: PostMessage
+    EE: EventEmitter<string, any>,
+    target: PostMessage
   ) {
-    this.onDisconnect(() => {
-      // 移除所有监听
-      this.EE.removeAllListeners("connectMessage:" + this.messageId);
-      this.EE.removeAllListeners("disconnect:" + this.messageId);
-    });
+    this.target = target; // 强引用
+    const handler = (msg: TMessage) => {
+      listenerMgr.emit(`onMessage:${this.listenerId}`, msg);
+    };
+    const cleanup = () => {
+      if (this.target) {
+        this.target = null;
+        listenerMgr.removeAllListeners(`cleanup:${this.listenerId}`);
+        EE.removeAllListeners("connectMessage:" + this.messageId); // 模拟 con.onMessage.removeListener
+        EE.removeAllListeners("disconnect:" + this.messageId); // 模拟 con.onDisconnect.removeListener
+        listenerMgr.emit(`onDisconnect:${this.listenerId}`, this.isSelfDisconnected);
+        listenerMgr.removeAllListeners(`onDisconnect:${this.listenerId}`);
+        listenerMgr.removeAllListeners(`onMessage:${this.listenerId}`);
+      }
+    };
+    EE.addListener(`connectMessage:${this.messageId}`, handler); // 模拟 con.onMessage.addListener
+    EE.addListener(`disconnect:${this.messageId}`, cleanup); // 模拟 con.onDisconnect.addListener
+    listenerMgr.once(`cleanup:${this.listenerId}`, cleanup);
   }
 
   sendMessage(data: TMessage) {
+    if (!this.target) {
+      console.error("Attempted to sendMessage on a disconnected Target.");
+      // 無法 sendMessage 不应该屏蔽错误
+      throw new Error("Attempted to sendMessage on a disconnected Target.");
+    }
     const body: WindowMessageBody<TMessage> = {
       messageId: this.messageId,
       type: "connectMessage",
@@ -143,20 +167,38 @@ export class WindowMessageConnect implements MessageConnect {
   }
 
   onMessage(callback: (data: TMessage) => void) {
-    this.EE.addListener(`connectMessage:${this.messageId}`, callback);
+    if (!this.target) {
+      console.error("onMessage Invalid Target");
+      // 無法監聽的話不应该屏蔽错误
+      throw new Error("onMessage Invalid Target");
+    }
+    listenerMgr.addListener(`onMessage:${this.listenerId}`, callback);
   }
 
   disconnect() {
+    if (!this.target) {
+      console.warn("Attempted to disconnect on a disconnected Target.");
+      // 重复 disconnect() 不应该屏蔽错误
+      throw new Error("Attempted to disconnect on a disconnected Target.");
+    }
+    this.isSelfDisconnected = true;
     const body: WindowMessageBody<TMessage> = {
       messageId: this.messageId,
       type: "disconnect",
       data: null,
     };
     this.target.postMessage(body);
+    // Note: .disconnect() will NOT automatically trigger the 'cleanup' listener
+    listenerMgr.emit(`cleanup:${this.listenerId}`);
   }
 
-  onDisconnect(callback: () => void) {
-    this.EE.addListener(`disconnect:${this.messageId}`, callback);
+  onDisconnect(callback: (isSelfDisconnected: boolean) => void) {
+    if (!this.target) {
+      console.error("onDisconnect Invalid Target");
+      // 無法監聽的話不应该屏蔽错误
+      throw new Error("onDisconnect Invalid Target");
+    }
+    listenerMgr.once(`onDisconnect:${this.listenerId}`, callback);
   }
 }