guanghulab/server/app/modules/agent-handshake.js

'use strict';

const crypto = require('crypto');
const COS = require('cos-nodejs-sdk-v5');
const { execSync } = require('child_process');

// ─── 常量 ───
const HANDSHAKE_EXPIRE_MS = 5 * 60 * 1000; // 5分钟
const HEARTBEAT_INTERVAL = 30 * 1000; // 30秒
const HEARTBEAT_TIMEOUT = 3 * HEARTBEAT_INTERVAL; // 90秒

// ─── COS 客户端 ───
const cos = new COS({
  SecretId: process.env.COS_SECRET_ID,
  SecretKey: process.env.COS_SECRET_KEY
});

// ─── 内存存储 ───
const pendingHandshakes = new Map(); // handshakeId → { initiatorId, targetId, createdAt, tempKey }
const activeConnections = new Map(); // connectionId → { initiatorId, targetId, lastHeartbeat, modelVersion }

// ─── 新增: 连接状态持久化 ───
const CONNECTION_BUCKET = 'agent-connections-1250000000';
const CONNECTION_REGION = 'ap-guangzhou';

/**
 * 保存连接状态到COS
 */
async function saveConnectionState(connectionId, state) {
  await cos.putObject({
    Bucket: CONNECTION_BUCKET,
    Region: CONNECTION_REGION,
    Key: `connections/${connectionId}.json`,
    Body: JSON.stringify(state)
  });
}

/**
 * 加载连接状态
 */
async function loadConnectionState(connectionId) {
  try {
    const { Body } = await cos.getObject({
      Bucket: CONNECTION_BUCKET,
      Region: CONNECTION_REGION,
      Key: `connections/${connectionId}.json`
    });
    return JSON.parse(Body.toString());
  } catch (err) {
    if (err.statusCode === 404) return null;
    throw err;
  }
}

// ─── 原有握手逻辑保持不变 ───
function generateHandshakeId() {
  return crypto.randomBytes(16).toString('hex');
}

function generateTempKey() {
  return crypto.randomBytes(32).toString('hex');
}

function generateConnectionId() {
  return crypto.randomBytes(16).toString('hex');
}

/**
 * 初始化握手
 */
function initiateHandshake(initiatorId, targetId, token) {
  const emailAuth = require('./email-auth');
  if (!emailAuth.validateSession(token).valid) {
    throw new Error('无效的会话令牌');
  }

  const handshakeId = generateHandshakeId();
  const tempKey = generateTempKey();

  pendingHandshakes.set(handshakeId, {
    initiatorId,
    targetId,
    createdAt: Date.now(),
    tempKey
  });

  return { handshakeId, tempKey };
}

/**
 * 完成握手
 */
function completeHandshake(handshakeId, targetId, targetKey, token) {
  const emailAuth = require('./email-auth');
  if (!emailAuth.validateSession(token).valid) {
    throw new Error('无效的会话令牌');
  }

  const handshake = pendingHandshakes.get(handshakeId);
  if (!handshake) {
    throw new Error('握手ID无效或已过期');
  }

  if (Date.now() - handshake.createdAt > HANDSHAKE_EXPIRE_MS) {
    pendingHandshakes.delete(handshakeId);
    throw new Error('握手已过期');
  }

  if (handshake.targetId !== targetId) {
    throw new Error('目标Agent不匹配');
  }

  if (handshake.tempKey !== targetKey) {
    throw new Error('临时密钥验证失败');
  }

  pendingHandshakes.delete(handshakeId);

  const connectionId = generateConnectionId();
  const now = Date.now();

  const connection = {
    initiatorId: handshake.initiatorId,
    targetId: handshake.targetId,
    lastHeartbeat: now,
    modelVersion: null
  };

  activeConnections.set(connectionId, connection);
  saveConnectionState(connectionId, connection);

  return { connectionId };
}

// ─── 新增: 模型版本更新 ───
function updateModelVersion(agentId, modelKey) {
  for (const [connId, conn] of activeConnections) {
    if (conn.initiatorId === agentId || conn.targetId === agentId) {
      conn.modelVersion = modelKey;
      saveConnectionState(connId, conn);
    }
  }
  return { success: true };
}

// ─── 原有心跳和断开逻辑保持不变 ───
function heartbeat(connectionId, token) {
  const emailAuth = require('./email-auth');
  if (!emailAuth.validateSession(token).valid) {
    throw new Error('无效的会话令牌');
  }

  const connection = activeConnections.get(connectionId);
  if (!connection) {
    throw new Error('连接不存在');
  }

  connection.lastHeartbeat = Date.now();
  saveConnectionState(connectionId, connection);

  return { success: true };
}

function disconnect(connectionId, token) {
  const emailAuth = require('./email-auth');
  if (!emailAuth.validateSession(token).valid) {
    throw new Error('无效的会话令牌');
  }

  if (!activeConnections.has(connectionId)) {
    throw new Error('连接不存在');
  }

  activeConnections.delete(connectionId);
  
  // 异步删除COS中的状态
  cos.deleteObject({
    Bucket: CONNECTION_BUCKET,
    Region: CONNECTION_REGION,
    Key: `connections/${connectionId}.json`
  }).catch(console.error);

  return { success: true };
}

// ─── 新增: 自动恢复连接 ───
async function restoreConnections() {
  try {
    const { Contents } = await cos.listObjects({
      Bucket: CONNECTION_BUCKET,
      Region: CONNECTION_REGION,
      Prefix: 'connections/'
    });

    if (!Contents) return;

    for (const item of Contents) {
      const connectionId = item.Key.split('/')[1].replace('.json', '');
      const state = await loadConnectionState(connectionId);
      
      if (state && Date.now() - state.lastHeartbeat < HEARTBEAT_TIMEOUT) {
        activeConnections.set(connectionId, state);
      } else {
        // 清理过期的连接状态
        await cos.deleteObject({
          Bucket: CONNECTION_BUCKET,
          Region: CONNECTION_REGION,
          Key: item.Key
        }).catch(console.error);
      }
    }
  } catch (err) {
    console.error('[Agent Handshake] 连接恢复失败:', err);
  }
}

// 启动时自动恢复连接
restoreConnections();

// 定时检查并恢复连接
setInterval(restoreConnections, HEARTBEAT_INTERVAL * 2);

module.exports = {
  initiateHandshake,
  completeHandshake,
  heartbeat,
  disconnect,
  updateModelVersion
};