theta-health-mcp-client/node_mcp_client.js

Theta Health MCP Client for WebSocket connection - Works like @notionhq/notion-mcp-server

#!/usr/bin/env node

const WebSocket = require('ws');
const readline = require('readline');

class ThetaHealthMCPClient {
  constructor(endpoint, token) {
    this.endpoint = endpoint;
    this.token = token;
    this.ws = null;
    this.connected = false;
    this.rl = null;
    this.cleanupTimeout = null;
  }

  async connect() {
    return new Promise((resolve, reject) => {
      const wsUrl = `${this.endpoint}?token=${this.token}`;
      
      try {
        this.ws = new WebSocket(wsUrl, ['mcp']);
        
        this.ws.on('open', () => {
          this.connected = true;
          this.setupMessageHandlers();
          resolve();
        });

        this.ws.on('message', (data) => {
          // 直接输出到 stdout（MCP 协议要求）
          console.log(data.toString());
        });

        this.ws.on('error', (error) => {
          console.error(`❌ WebSocket 错误: ${error.message}`);
          reject(error);
        });

        this.ws.on('close', () => {
          this.cleanup();
          process.exit(0);
        });

        // 连接超时处理
        setTimeout(() => {
          if (!this.connected) {
            reject(new Error('连接超时'));
          }
        }, 5000);

      } catch (error) {
        console.error(`❌ 连接失败: ${error.message}`);
        reject(error);
      }
    });
  }

  setupMessageHandlers() {
    // 设置 stdin/stdout 接口
    this.rl = readline.createInterface({
      input: process.stdin,
      output: process.stdout,
      terminal: false
    });

    // 处理来自 stdin 的 MCP 消息
    this.rl.on('line', (line) => {
      if (line.trim() && this.connected) {
        try {
          // 验证 JSON 格式
          JSON.parse(line);
          if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            this.ws.send(line);
          }
        } catch (e) {
          console.error(`⚠️ 无效的JSON: ${line}`);
        }
      }
    });

    this.rl.on('close', () => {
      // stdin 关闭后，等待2秒让服务器有时间响应
      this.cleanupTimeout = setTimeout(() => {
        this.cleanup();
      }, 2000);
    });

    // 处理 stdin 结束（EOF）
    process.stdin.on('end', () => {
      // stdin 结束后，等待2秒让服务器有时间响应
      if (!this.cleanupTimeout) {
        this.cleanupTimeout = setTimeout(() => {
          this.cleanup();
        }, 2000);
      }
    });

    // 处理进程信号
    process.on('SIGINT', () => {
      this.cleanup();
      process.exit(0);
    });

    process.on('SIGTERM', () => {
      this.cleanup();
      process.exit(0);
    });
  }

  cleanup() {
    if (this.cleanupTimeout) {
      clearTimeout(this.cleanupTimeout);
      this.cleanupTimeout = null;
    }
    if (this.rl) {
      this.rl.close();
      this.rl = null;
    }
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.close();
    }
  }
}

// 主函数
async function main() {
  const endpoint = process.env.THETA_HEALTH_ENDPOINT || 'ws://localhost:18022';
  const token = process.env.THETA_HEALTH_TOKEN;
  
  if (!token) {
    console.error('❌ 缺少 THETA_HEALTH_TOKEN 环境变量');
    process.exit(1);
  }

  const client = new ThetaHealthMCPClient(endpoint, token);
  
  try {
    await client.connect();
    
    // 如果没有stdin输入，设置超时退出（防止在某些环境下挂起）
    if (process.stdin.isTTY === false) {
      // 非交互模式下，如果15秒内没有输入就退出
      const timeout = setTimeout(() => {
        console.error('⏰ 非交互模式下等待输入超时');
        client.cleanup();
        process.exit(1);
      }, 15000);
      
      // 如果收到输入，清除超时
      process.stdin.once('data', () => {
        clearTimeout(timeout);
      });
    }
    
  } catch (error) {
    console.error(`❌ 连接失败: ${error.message}`);
    process.exit(1);
  }
}

if (require.main === module) {
  main();
}

module.exports = ThetaHealthMCPClient;