260 lines
7.8 KiB
JavaScript
260 lines
7.8 KiB
JavaScript
#!/usr/bin/env node
|
||
// federation-agent/server/registry.js · v1.0
|
||
// 服务器端 · Agent 注册管理中心
|
||
// 挂载到 gatekeeper 上，管理所有联邦 Agent 节点
|
||
// ICE-GL-ZY001 · D129 · 2026-06-10
|
||
|
||
const WebSocket = require("ws");
|
||
const http = require("http");
|
||
|
||
// ===== 配置 =====
|
||
const PORT = process.env.AGENT_PORT || 3912;
|
||
const AUTH_TOKEN = process.env.AGENT_AUTH || "zy_gtw_federation_agent_v1";
|
||
|
||
// ===== 注册表 =====
|
||
class AgentRegistry {
|
||
constructor() {
|
||
this.agents = new Map(); // agentId → { ws, info, status }
|
||
this.personas = new Map(); // personaId → [agentId, ...]
|
||
this.users = new Map(); // userId → [agentId, ...]
|
||
this.commandQueue = new Map(); // agentId → [pending commands]
|
||
}
|
||
|
||
register(ws, info) {
|
||
const { agentId, personaId, userId, hostname, platform, workDir } = info;
|
||
|
||
const agent = {
|
||
ws,
|
||
info: { agentId, personaId, userId, hostname, platform, workDir, version: info.version || "1.0.0" },
|
||
status: "online",
|
||
connectedAt: new Date().toISOString(),
|
||
lastHeartbeat: Date.now(),
|
||
commandCount: 0,
|
||
};
|
||
|
||
this.agents.set(agentId, agent);
|
||
|
||
if (personaId) {
|
||
if (!this.personas.has(personaId)) this.personas.set(personaId, []);
|
||
this.personas.get(personaId).push(agentId);
|
||
}
|
||
if (userId) {
|
||
if (!this.users.has(userId)) this.users.set(userId, []);
|
||
this.users.get(userId).push(agentId);
|
||
}
|
||
|
||
console.log(`[REGISTRY] + ${agentId} (${personaId || "未注册"} @ ${userId || "未知用户"} · ${hostname})`);
|
||
return agent;
|
||
}
|
||
|
||
unregister(agentId) {
|
||
const agent = this.agents.get(agentId);
|
||
if (!agent) return;
|
||
|
||
this.agents.delete(agentId);
|
||
|
||
if (agent.info.personaId) {
|
||
const list = this.personas.get(agent.info.personaId);
|
||
if (list) {
|
||
const idx = list.indexOf(agentId);
|
||
if (idx > -1) list.splice(idx, 1);
|
||
}
|
||
}
|
||
if (agent.info.userId) {
|
||
const list = this.users.get(agent.info.userId);
|
||
if (list) {
|
||
const idx = list.indexOf(agentId);
|
||
if (idx > -1) list.splice(idx, 1);
|
||
}
|
||
}
|
||
|
||
console.log(`[REGISTRY] - ${agentId} (断开)`);
|
||
}
|
||
|
||
sendCommand(agentId, command) {
|
||
const agent = this.agents.get(agentId);
|
||
if (!agent || !agent.ws || agent.ws.readyState !== WebSocket.OPEN) {
|
||
// 加入待发送队列
|
||
if (!this.commandQueue.has(agentId)) this.commandQueue.set(agentId, []);
|
||
this.commandQueue.get(agentId).push(command);
|
||
return { ok: false, error: "agent offline" };
|
||
}
|
||
|
||
agent.ws.send(JSON.stringify({ type: "command", command }));
|
||
agent.commandCount++;
|
||
return { ok: true };
|
||
}
|
||
|
||
flushQueue(agentId) {
|
||
const queue = this.commandQueue.get(agentId);
|
||
if (!queue || queue.length === 0) return 0;
|
||
|
||
let sent = 0;
|
||
while (queue.length > 0) {
|
||
const cmd = queue.shift();
|
||
const result = this.sendCommand(agentId, cmd);
|
||
if (result.ok) sent++;
|
||
else break;
|
||
}
|
||
return sent;
|
||
}
|
||
|
||
broadcast(personaIds, command) {
|
||
const results = [];
|
||
for (const personaId of personaIds) {
|
||
const agentIds = this.personas.get(personaId) || [];
|
||
for (const agentId of agentIds) {
|
||
results.push({ agentId, ...this.sendCommand(agentId, command) });
|
||
}
|
||
}
|
||
return results;
|
||
}
|
||
|
||
list(target) {
|
||
if (target === "all") {
|
||
return Array.from(this.agents.values()).map(a => ({
|
||
agentId: a.info.agentId,
|
||
personaId: a.info.personaId,
|
||
userId: a.info.userId,
|
||
hostname: a.info.hostname,
|
||
status: a.status,
|
||
commandCount: a.commandCount,
|
||
connectedAt: a.connectedAt,
|
||
}));
|
||
}
|
||
// 按用户过滤
|
||
const agentIds = this.users.get(target) || [];
|
||
return agentIds.map(id => {
|
||
const a = this.agents.get(id);
|
||
return a ? {
|
||
agentId: a.info.agentId,
|
||
personaId: a.info.personaId,
|
||
status: a.status,
|
||
commandCount: a.commandCount,
|
||
} : null;
|
||
}).filter(Boolean);
|
||
}
|
||
|
||
healthCheck() {
|
||
const now = Date.now();
|
||
for (const [agentId, agent] of this.agents) {
|
||
if (now - agent.lastHeartbeat > 120000) {
|
||
agent.status = "timeout";
|
||
console.log(`[REGISTRY] ⚠ ${agentId} 心跳超时`);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// ===== WebSocket 服务器 =====
|
||
const registry = new AgentRegistry();
|
||
const server = http.createServer((req, res) => {
|
||
// REST API 端点
|
||
if (req.url === "/api/agents" && req.method === "GET") {
|
||
const target = new URL(req.url, "http://localhost").searchParams.get("user") || "all";
|
||
res.writeHead(200, { "Content-Type": "application/json" });
|
||
res.end(JSON.stringify(registry.list(target)));
|
||
return;
|
||
}
|
||
if (req.url === "/api/command" && req.method === "POST") {
|
||
let body = "";
|
||
req.on("data", chunk => body += chunk);
|
||
req.on("end", () => {
|
||
try {
|
||
const { agentId, command } = JSON.parse(body);
|
||
const result = registry.sendCommand(agentId, command);
|
||
res.writeHead(result.ok ? 200 : 404, { "Content-Type": "application/json" });
|
||
res.end(JSON.stringify(result));
|
||
} catch (e) {
|
||
res.writeHead(400, { "Content-Type": "application/json" });
|
||
res.end(JSON.stringify({ ok: false, error: e.message }));
|
||
}
|
||
});
|
||
return;
|
||
}
|
||
res.writeHead(404);
|
||
res.end("Not Found");
|
||
});
|
||
|
||
const wss = new WebSocket.Server({ server, path: "/agent" });
|
||
|
||
wss.on("connection", (ws, req) => {
|
||
const url = new URL(req.url, "http://localhost");
|
||
const agentId = url.searchParams.get("agentId") || "unknown";
|
||
const personaId = url.searchParams.get("persona") || "";
|
||
const userId = url.searchParams.get("user") || "";
|
||
|
||
console.log(`[WS] 新连接: ${agentId}`);
|
||
|
||
// 先不发 welcome，等 register 消息
|
||
ws.agentId = agentId;
|
||
|
||
ws.on("message", (data) => {
|
||
try {
|
||
const msg = JSON.parse(data.toString());
|
||
|
||
switch (msg.type) {
|
||
case "register": {
|
||
const agent = registry.register(ws, msg);
|
||
ws.send(JSON.stringify({
|
||
type: "welcome",
|
||
message: `光湖联邦托管 · 注册成功 · ${agent.info.personaId || agentId}`,
|
||
personaId: agent.info.personaId,
|
||
timestamp: Date.now(),
|
||
}));
|
||
// 发送队列中的命令
|
||
const flushed = registry.flushQueue(agentId);
|
||
if (flushed > 0) console.log(`[REGISTRY] → ${agentId} 补发 ${flushed} 条待处理指令`);
|
||
break;
|
||
}
|
||
case "result": {
|
||
console.log(`[RESULT] ${agentId}: ${msg.command} → ${msg.ok ? "OK" : "FAIL"}`);
|
||
if (msg.stderr) console.log(` stderr: ${msg.stderr.substring(0, 200)}`);
|
||
break;
|
||
}
|
||
case "heartbeat": {
|
||
const agent = registry.agents.get(agentId);
|
||
if (agent) agent.lastHeartbeat = Date.now();
|
||
break;
|
||
}
|
||
case "health_report": {
|
||
console.log(`[HEALTH] ${agentId}: uptime=${(msg.uptime / 60).toFixed(1)}min · mem=${(msg.memory.heapUsed / 1024 / 1024).toFixed(1)}MB`);
|
||
break;
|
||
}
|
||
case "confirm_request": {
|
||
console.log(`[CONFIRM] ${agentId}: 请求确认危险操作: ${msg.operation}`);
|
||
break;
|
||
}
|
||
default:
|
||
console.log(`[MSG] ${agentId}: ${msg.type}`);
|
||
}
|
||
} catch (e) {
|
||
console.error(`[ERROR] ${agentId}: ${e.message}`);
|
||
}
|
||
});
|
||
|
||
ws.on("close", () => {
|
||
registry.unregister(agentId);
|
||
});
|
||
|
||
ws.on("error", (err) => {
|
||
console.error(`[WS ERROR] ${agentId}: ${err.message}`);
|
||
});
|
||
});
|
||
|
||
// 健康检查定时器
|
||
setInterval(() => registry.healthCheck(), 30000);
|
||
|
||
server.listen(PORT, () => {
|
||
console.log(`[REGISTRY] 光湖联邦 Agent 注册中心 · 端口 ${PORT}`);
|
||
console.log(`[REGISTRY] WebSocket: ws://0.0.0.0:${PORT}/agent`);
|
||
console.log(`[REGISTRY] REST API: http://0.0.0.0:${PORT}/api/agents`);
|
||
});
|
||
|
||
process.on("SIGINT", () => {
|
||
console.log("[REGISTRY] 关闭中...");
|
||
wss.close();
|
||
server.close();
|
||
process.exit(0);
|
||
});
|