openmcp-client/renderer/src/api/message-bridge.ts

import { pinkLog, redLog } from '@/views/setting/util';
import { acquireVsCodeApi, electronApi, getPlatform } from './platform';
import { isReactive } from 'vue';

export interface VSCodeMessage {
	command: string;
	data?: unknown;
	callbackId?: string;
}

export interface RestFulResponse<T = any> {
	code: number;
	msg: T;
}

export type MessageHandler = (message: VSCodeMessage) => void;
export type CommandHandler = (data: any) => void;

interface AddCommandListenerOption {
	once: boolean // 只调用一次就销毁
}

export interface ICommandRequestData {
	clientId?: string;
	[key: string]: any;
}

export class MessageBridge {
	private ws: WebSocket | null = null;
	private handlers = new Map<string, Set<CommandHandler>>();
	private isConnected: Promise<boolean> | null = null;

	constructor(
		private setupSignature: any
	) {

		const platform = getPlatform();

		switch (platform) {
			case 'vscode':
				this.setupVsCodeListener();
				pinkLog('current platform: vscode');
				break;

			case 'electron':
				this.setupElectronListener();
				pinkLog('current platform: electron');
				break;
			
			case 'nodejs':
				this.setupNodejsListener();
				pinkLog('current platform: nodejs');
				break;
			
			case 'web':
				this.setupWebSocket();
				pinkLog('current platform: web');
				break;
		}
	}

	// VS Code 环境监听
	private setupVsCodeListener() {
		const vscode = acquireVsCodeApi();

		window.addEventListener('message', (event: MessageEvent<VSCodeMessage>) => {
			this.dispatchMessage(event.data);
		});

		this.postMessage = (message) => vscode.postMessage(message);
	}

	// WebSocket 环境连接
	public setupWebSocket(setupSignature?: string) {
		const wsUrl = setupSignature || this.setupSignature;

		if (typeof wsUrl !== 'string') {
			throw new Error('setupSignature must be a string');
		}
		
		console.log(wsUrl);
		
		this.ws = new WebSocket(wsUrl);
		const ws = this.ws;

		this.isConnected = new Promise<boolean>((resolve, reject) => {
			ws.onopen = () => {
				resolve(true);
			};

			ws.onmessage = (event) => {
				try {				
					const message = JSON.parse(event.data) as VSCodeMessage;
					this.dispatchMessage(message);
				} catch (err) {
					console.error('Message parse error:', err);
					console.log(event);
				}
			};
	
			ws.onerror = (err) => {
				redLog('WebSocket error:');				
				resolve(false);
			};
	
			ws.onclose = () => {
				redLog('WebSocket connection closed');
				resolve(false);
			};
	
			this.postMessage = (message) => {
				if (this.ws?.readyState === WebSocket.OPEN) {
					console.log('send', message);
					this.ws.send(JSON.stringify(message));
				}
			};
		});

		return this.isConnected;
	}

	public async awaitForWebsocket() {
		
		if (this.isConnected) {
			return await this.isConnected;
		}
		return false;
	}

	private setupElectronListener() {
		electronApi.onReply((event: MessageEvent<VSCodeMessage>) => {
			console.log(event);
			this.dispatchMessage(event.data);
		});

		this.postMessage = (message) => {
			console.log(message);
			electronApi.sendToMain(message);
		};		
	}

	private setupNodejsListener() {
		
		const emitter = this.setupSignature;
		if (!emitter.on || !emitter.emit) {
			return;
		}

		emitter.on('message/service', (message: VSCodeMessage) => {
			this.dispatchMessage(message);
		});

		this.postMessage = (message) => {
			emitter.emit('message/renderer', message);
		};
	}

	/**
	 * @description 对 message 发起调度，根据 command 类型获取调取器
	 * @param message 
	 */
	private dispatchMessage(message: VSCodeMessage) {
		const command = message.command;
		const data = message.data;

		const handlers = this.handlers.get(command) || [];
		handlers.forEach(handler => handler(data));
	}

	public postMessage(message: VSCodeMessage) {
		throw new Error('PostMessage not initialized');
	}

	/**
	 * @description 注册一个命令的执行器（支持一次性或持久监听）
	 * @example
	 * // 基本用法（持久监听）
	 * const removeListener = bridge.addCommandListener('message', (data) => {
	 *   console.log('收到消息:', data.msg.text);
	 * }, { once: false });
	 * 
	 * // 稍后取消监听
	 * removeListener();
	 * 
	 * @example
	 * // 一次性监听（自动移除）
	 * bridge.addCommandListener('connect', (data) => {
	 *   const { code, msg } = data;
	 *   console.log(`连接结果: ${code === 200 ? '成功' : '失败'}`);
	 * }, { once: true });
	 */
	public addCommandListener(
		command: string,
		commandHandler: CommandHandler,
		option: AddCommandListenerOption
	): () => boolean {
		if (!this.handlers.has(command)) {
			this.handlers.set(command, new Set<CommandHandler>());
		}
		// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
		const commandHandlers = this.handlers.get(command)!;

		const wrapperCommandHandler = option.once ? (data: any) => {
			commandHandler(data);
			commandHandlers.delete(wrapperCommandHandler);
		} : commandHandler;

		commandHandlers.add(wrapperCommandHandler);
		return () => commandHandlers.delete(wrapperCommandHandler);
	}

	private deserializeReactiveData(data: any) {
		if (isReactive(data)) {
			return JSON.parse(JSON.stringify(data));
		}

		// 只对第一层进行遍历
		for (const key in data) {
			if (isReactive(data[key])) {
				data[key] = JSON.parse(JSON.stringify(data[key]));
			}
		}

		return data;
	}

	/**
	 * @description do as axios does
	 * @param command 
	 * @param data 
	 * @returns 
	 */
	public commandRequest<T = any>(command: string, data?: ICommandRequestData): Promise<RestFulResponse<T>>  {

		return new Promise<RestFulResponse>((resolve, reject) => {
			this.addCommandListener(command, (data) => {
				resolve(data as RestFulResponse);
			}, { once: true });

			this.postMessage({
				command,
				data: this.deserializeReactiveData(data)
			});
		});
	}

	public destroy() {
		this.ws?.close();
		this.handlers.clear();
	}
}

// 单例实例
let messageBridge: MessageBridge;

export function createMessageBridge(setupSignature: any) {
	messageBridge = new MessageBridge(setupSignature);
}

// 向外暴露一个独立函数，保证 MessageBridge 是单例的
export function useMessageBridge() {
	if (!messageBridge && getPlatform() !== 'nodejs') {		
		messageBridge = new MessageBridge(import.meta.env.VITE_WEBSOCKET_URL);
	}
	const bridge = messageBridge;

	return bridge;
}