394 lines
14 KiB
JavaScript
394 lines
14 KiB
JavaScript
import crypto from "node:crypto";
|
|
import { createServer } from "node:http";
|
|
import express from "express";
|
|
import { WebSocket, WebSocketServer } from "ws";
|
|
|
|
const DEFAULT_WS_PATH = "/api/ai-workspace/hub";
|
|
const MAX_EVENTS_PER_AGENT = 5000;
|
|
const MAX_REQUEST_META_PER_AGENT = 1000;
|
|
|
|
const config = readConfig();
|
|
const app = express();
|
|
const httpServer = createServer(app);
|
|
const wss = new WebSocketServer({ noServer: true });
|
|
const agentsByCode = new Map();
|
|
let eventSeq = 0;
|
|
|
|
app.disable("x-powered-by");
|
|
app.use(express.json({ limit: "1mb" }));
|
|
|
|
app.get("/healthz", (_req, res) => {
|
|
res.json({
|
|
ok: true,
|
|
service: "nodedc-ai-workspace-hub",
|
|
wsPath: config.wsPath,
|
|
agentsOnline: Array.from(agentsByCode.values()).filter(isAgentOnline).length,
|
|
internalApiConfigured: config.internalAccessTokens.length > 0,
|
|
});
|
|
});
|
|
|
|
app.get("/api/ai-workspace/hub/v1/agents/:pairingCode", requireInternalApi, (req, res) => {
|
|
const agent = agentsByCode.get(cleanPairingCode(req.params.pairingCode));
|
|
res.json({ ok: true, agent: publicAgent(agent) });
|
|
});
|
|
|
|
app.get("/api/ai-workspace/hub/v1/agents/:pairingCode/events", requireInternalApi, (req, res) => {
|
|
const payload = readAgentEvents(req.params.pairingCode, {
|
|
since: req.query.since,
|
|
limit: req.query.limit,
|
|
});
|
|
res.json({ ok: true, ...payload });
|
|
});
|
|
|
|
app.post("/api/ai-workspace/hub/v1/agents/:pairingCode/request", requireInternalApi, asyncRoute(async (req, res) => {
|
|
const command = cleanString(req.body?.command, 100);
|
|
const timeoutMs = sanitizeTimeoutMs(req.body?.timeoutMs, 30000);
|
|
const payload = await requestAgent(req.params.pairingCode, command, req.body?.payload || {}, timeoutMs, {
|
|
quiet: Boolean(req.body?.quiet),
|
|
});
|
|
res.json({ ok: true, payload });
|
|
}));
|
|
|
|
app.post("/api/ai-workspace/hub/v1/agents/:pairingCode/dispatch", requireInternalApi, (req, res) => {
|
|
const command = cleanString(req.body?.command, 100);
|
|
const timeoutMs = sanitizeTimeoutMs(req.body?.timeoutMs, 30000);
|
|
const dispatch = dispatchAgent(req.params.pairingCode, command, req.body?.payload || {}, timeoutMs, {
|
|
quiet: Boolean(req.body?.quiet),
|
|
});
|
|
res.json({ ok: true, requestId: dispatch.requestId });
|
|
});
|
|
|
|
app.use((error, _req, res, _next) => {
|
|
const status = Number(error?.status || 500);
|
|
res.status(status >= 400 && status < 600 ? status : 500).json({
|
|
ok: false,
|
|
error: String(error?.message || error || "internal_error"),
|
|
});
|
|
});
|
|
|
|
httpServer.on("upgrade", (req, socket, head) => {
|
|
let pathname = "";
|
|
try {
|
|
pathname = new URL(req.url || "/", "http://localhost").pathname;
|
|
} catch {}
|
|
if (pathname !== config.wsPath) return;
|
|
wss.handleUpgrade(req, socket, head, (ws) => {
|
|
wss.emit("connection", ws, req);
|
|
});
|
|
});
|
|
|
|
wss.on("connection", attachAgent);
|
|
|
|
const heartbeat = setInterval(() => {
|
|
for (const agent of agentsByCode.values()) {
|
|
try {
|
|
agent.socket.ping();
|
|
} catch {}
|
|
}
|
|
}, 30000);
|
|
|
|
httpServer.listen(config.port, "0.0.0.0", () => {
|
|
console.log(`NODE.DC AI Workspace Hub listening on http://0.0.0.0:${config.port}`);
|
|
console.log(`NODE.DC AI Workspace Hub ws path: ${config.wsPath}`);
|
|
});
|
|
|
|
process.on("SIGTERM", shutdown);
|
|
process.on("SIGINT", shutdown);
|
|
|
|
function attachAgent(socket, req) {
|
|
const url = new URL(req.url || config.wsPath, "http://localhost");
|
|
const pairingCode = cleanPairingCode(url.searchParams.get("pairingCode"));
|
|
if (!pairingCode) {
|
|
socket.close(1008, "pairing_code_required");
|
|
return;
|
|
}
|
|
|
|
const previous = agentsByCode.get(pairingCode);
|
|
if (previous) {
|
|
settleAll(previous, Object.assign(new Error("agent_replaced"), { status: 503 }));
|
|
previous.socket.close(4000, "agent_replaced");
|
|
}
|
|
|
|
const now = new Date().toISOString();
|
|
const agent = {
|
|
socket,
|
|
pairingCode,
|
|
machineName: cleanString(url.searchParams.get("machineName"), 120),
|
|
host: "",
|
|
connectedAt: now,
|
|
lastSeenAt: now,
|
|
pending: new Map(),
|
|
requestMeta: previous?.requestMeta || new Map(),
|
|
events: previous?.events || [],
|
|
};
|
|
agentsByCode.set(pairingCode, agent);
|
|
pushAgentEvent(agent, "", { kind: "connected", message: "Bridge agent connected." });
|
|
|
|
socket.on("message", (raw) => handleAgentMessage(agent, raw));
|
|
socket.on("close", () => {
|
|
pushAgentEvent(agent, "", { kind: "disconnected", message: "Bridge agent disconnected." });
|
|
if (agentsByCode.get(pairingCode) === agent) agentsByCode.delete(pairingCode);
|
|
settleAll(agent, Object.assign(new Error("bridge_agent_offline"), { status: 503 }));
|
|
});
|
|
socket.on("error", () => {});
|
|
|
|
sendJson(socket, {
|
|
type: "ready",
|
|
pairingCode,
|
|
connectedAt: now,
|
|
});
|
|
}
|
|
|
|
function handleAgentMessage(agent, raw) {
|
|
agent.lastSeenAt = new Date().toISOString();
|
|
let message = null;
|
|
try {
|
|
message = JSON.parse(String(raw));
|
|
} catch {
|
|
return;
|
|
}
|
|
|
|
if (message?.type === "hello") {
|
|
agent.machineName = cleanString(message.machineName || agent.machineName, 120);
|
|
agent.host = cleanString(message.host, 120);
|
|
pushAgentEvent(agent, "", {
|
|
kind: "hello",
|
|
message: "Bridge agent hello received.",
|
|
cwd: cleanString(message.cwd, 500),
|
|
});
|
|
return;
|
|
}
|
|
|
|
if (message?.type === "event") {
|
|
const pending = message?.requestId ? agent.pending.get(String(message.requestId || "")) : null;
|
|
if (pending?.quiet) return;
|
|
pushAgentEvent(agent, message.requestId, message.event || message.payload || {});
|
|
return;
|
|
}
|
|
|
|
if (message?.type !== "response") return;
|
|
const requestId = String(message.requestId || "");
|
|
const pending = agent.pending.get(requestId);
|
|
if (!pending) return;
|
|
agent.pending.delete(requestId);
|
|
clearTimeout(pending.timeout);
|
|
if (!pending.quiet) {
|
|
pushAgentEvent(agent, requestId, {
|
|
kind: "response_received",
|
|
message: `Bridge response received after ${elapsedLabel(pending.startedAt)}.`,
|
|
});
|
|
}
|
|
if (message.ok === false) {
|
|
const error = new Error(String(message.error || "bridge_agent_failed"));
|
|
error.status = Number(message.status || 502);
|
|
if (!pending.quiet) pushAgentEvent(agent, requestId, { kind: "error", message: error.message });
|
|
pending.reject(error);
|
|
return;
|
|
}
|
|
if (!pending.quiet) pushAgentEvent(agent, requestId, { kind: "done", message: "Bridge request completed." });
|
|
pending.resolve(message.payload || { ok: true });
|
|
}
|
|
|
|
function requestAgent(pairingCodeRaw, command, payload = {}, timeoutMs = 30000, options = {}) {
|
|
return startAgentRequest(pairingCodeRaw, command, payload, timeoutMs, options).promise;
|
|
}
|
|
|
|
function dispatchAgent(pairingCodeRaw, command, payload = {}, timeoutMs = 30000, options = {}) {
|
|
const { requestId, promise } = startAgentRequest(pairingCodeRaw, command, payload, timeoutMs, options);
|
|
promise.catch(() => {});
|
|
return { requestId };
|
|
}
|
|
|
|
function startAgentRequest(pairingCodeRaw, command, payload = {}, timeoutMs = 30000, options = {}) {
|
|
const pairingCode = cleanPairingCode(pairingCodeRaw);
|
|
const agent = agentsByCode.get(pairingCode);
|
|
if (!agent || !isAgentOnline(agent)) {
|
|
const error = new Error("bridge_agent_offline");
|
|
error.status = 503;
|
|
throw error;
|
|
}
|
|
|
|
const requestId = crypto.randomUUID();
|
|
const quiet = Boolean(options?.quiet);
|
|
agent.requestMeta.set(requestId, requestMetaFromPayload(command, payload));
|
|
trimRequestMeta(agent);
|
|
const promise = new Promise((resolve, reject) => {
|
|
const timeout = setTimeout(() => {
|
|
agent.pending.delete(requestId);
|
|
const error = new Error("bridge_timeout");
|
|
error.status = 504;
|
|
if (!quiet) pushAgentEvent(agent, requestId, { kind: "timeout", command, message: "Bridge request timed out." });
|
|
reject(error);
|
|
}, timeoutMs);
|
|
const startedAt = Date.now();
|
|
agent.pending.set(requestId, { resolve, reject, timeout, startedAt, quiet });
|
|
if (!quiet) pushAgentEvent(agent, requestId, { kind: "engine_send", command, message: `Engine sent request: ${command}` });
|
|
sendJson(agent.socket, {
|
|
type: "request",
|
|
requestId,
|
|
command: String(command || ""),
|
|
payload,
|
|
quiet,
|
|
});
|
|
});
|
|
return { requestId, promise };
|
|
}
|
|
|
|
function readAgentEvents(pairingCodeRaw, options = {}) {
|
|
const pairingCode = cleanPairingCode(pairingCodeRaw);
|
|
const agent = agentsByCode.get(pairingCode);
|
|
const since = Number(options.since || 0);
|
|
const limit = Math.max(0, Math.min(2000, Number(options.limit || 100)));
|
|
const events = agent?.events || [];
|
|
const latestEventId = events.length ? events[events.length - 1].id : eventSeq;
|
|
return {
|
|
events: events.filter((event) => event.id > since).slice(-limit),
|
|
latestEventId,
|
|
online: Boolean(agent && isAgentOnline(agent)),
|
|
};
|
|
}
|
|
|
|
function requestMetaFromPayload(command, payload = {}) {
|
|
const context = payload?.context && typeof payload.context === "object" ? payload.context : {};
|
|
return {
|
|
command: cleanString(command, 80),
|
|
threadId: cleanString(payload?.threadId, 160),
|
|
threadTitle: cleanString(payload?.threadTitle, 240),
|
|
userMessage: cleanString(payload?.message, 12000),
|
|
remoteControlMode: cleanString(context.remoteControlMode, 40),
|
|
};
|
|
}
|
|
|
|
function pushAgentEvent(agent, requestId, rawEvent = {}) {
|
|
if (!agent) return null;
|
|
const now = new Date().toISOString();
|
|
const cleanRequestId = String(requestId || rawEvent.requestId || "");
|
|
const meta = cleanRequestId ? agent.requestMeta?.get(cleanRequestId) : null;
|
|
const event = {
|
|
id: ++eventSeq,
|
|
pairingCode: agent.pairingCode,
|
|
requestId: cleanRequestId,
|
|
threadId: rawEvent.threadId == null ? cleanString(meta?.threadId, 160) : cleanString(rawEvent.threadId, 160),
|
|
threadTitle: rawEvent.threadTitle == null ? cleanString(meta?.threadTitle, 240) : cleanString(rawEvent.threadTitle, 240),
|
|
userMessage: rawEvent.userMessage == null ? cleanString(meta?.userMessage, 12000) : cleanString(rawEvent.userMessage, 12000),
|
|
remoteControlMode: rawEvent.remoteControlMode == null ? cleanString(meta?.remoteControlMode, 40) : cleanString(rawEvent.remoteControlMode, 40),
|
|
kind: cleanString(rawEvent.kind || rawEvent.type || "event", 40),
|
|
stream: rawEvent.stream ? cleanString(rawEvent.stream, 40) : "",
|
|
text: rawEvent.text == null ? "" : cleanString(rawEvent.text, 12000),
|
|
message: rawEvent.message == null ? "" : cleanString(rawEvent.message, 12000),
|
|
command: rawEvent.command == null ? "" : cleanString(rawEvent.command, 1000),
|
|
cwd: rawEvent.cwd == null ? "" : cleanString(rawEvent.cwd, 500),
|
|
currentFile: rawEvent.currentFile == null ? "" : cleanString(rawEvent.currentFile, 1000),
|
|
eventsFile: rawEvent.eventsFile == null ? "" : cleanString(rawEvent.eventsFile, 1000),
|
|
at: rawEvent.at ? cleanString(rawEvent.at, 80) : now,
|
|
};
|
|
agent.events.push(event);
|
|
if (agent.events.length > MAX_EVENTS_PER_AGENT) {
|
|
agent.events.splice(0, agent.events.length - MAX_EVENTS_PER_AGENT);
|
|
}
|
|
return event;
|
|
}
|
|
|
|
function publicAgent(agent) {
|
|
if (!agent) return null;
|
|
return {
|
|
online: isAgentOnline(agent),
|
|
pairingCode: agent.pairingCode,
|
|
machineName: agent.machineName,
|
|
host: agent.host,
|
|
connectedAt: agent.connectedAt,
|
|
lastSeenAt: agent.lastSeenAt,
|
|
};
|
|
}
|
|
|
|
function isAgentOnline(agent) {
|
|
return Boolean(agent?.socket && agent.socket.readyState === WebSocket.OPEN);
|
|
}
|
|
|
|
function settleAll(agent, error) {
|
|
for (const pending of agent.pending.values()) {
|
|
clearTimeout(pending.timeout);
|
|
pending.reject(error);
|
|
}
|
|
agent.pending.clear();
|
|
}
|
|
|
|
function trimRequestMeta(agent) {
|
|
if (!agent?.requestMeta || agent.requestMeta.size <= MAX_REQUEST_META_PER_AGENT) return;
|
|
const excess = agent.requestMeta.size - MAX_REQUEST_META_PER_AGENT;
|
|
Array.from(agent.requestMeta.keys()).slice(0, excess).forEach((key) => agent.requestMeta.delete(key));
|
|
}
|
|
|
|
function sendJson(socket, payload) {
|
|
socket.send(JSON.stringify(payload));
|
|
}
|
|
|
|
function cleanPairingCode(value) {
|
|
return String(value || "")
|
|
.trim()
|
|
.toUpperCase()
|
|
.replace(/[^A-Z0-9]/g, "")
|
|
.slice(0, 32);
|
|
}
|
|
|
|
function cleanString(value, maxLength) {
|
|
return String(value || "").trim().slice(0, maxLength);
|
|
}
|
|
|
|
function sanitizeTimeoutMs(value, fallback) {
|
|
const numeric = Number(value || fallback);
|
|
if (!Number.isFinite(numeric)) return fallback;
|
|
return Math.max(1000, Math.min(12 * 60 * 60 * 1000, numeric));
|
|
}
|
|
|
|
function elapsedLabel(startedAt) {
|
|
const elapsed = Math.max(0, Date.now() - Number(startedAt || Date.now()));
|
|
if (elapsed < 1000) return `${elapsed}ms`;
|
|
return `${(elapsed / 1000).toFixed(1)}s`;
|
|
}
|
|
|
|
function asyncRoute(handler) {
|
|
return (req, res, next) => Promise.resolve(handler(req, res, next)).catch(next);
|
|
}
|
|
|
|
function requireInternalApi(req, res, next) {
|
|
if (!config.internalAccessTokens.length) {
|
|
res.status(503).json({ ok: false, error: "ai_workspace_hub_token_not_configured" });
|
|
return;
|
|
}
|
|
const header = String(req.headers.authorization || "");
|
|
const token = header.startsWith("Bearer ") ? header.slice(7) : "";
|
|
if (!token || !config.internalAccessTokens.some((candidate) => safeEqual(token, candidate))) {
|
|
res.status(401).json({ ok: false, error: "ai_workspace_hub_unauthorized" });
|
|
return;
|
|
}
|
|
next();
|
|
}
|
|
|
|
function safeEqual(left, right) {
|
|
const leftBuffer = Buffer.from(String(left || ""));
|
|
const rightBuffer = Buffer.from(String(right || ""));
|
|
return leftBuffer.length === rightBuffer.length && crypto.timingSafeEqual(leftBuffer, rightBuffer);
|
|
}
|
|
|
|
function readConfig() {
|
|
const internalAccessTokens = Array.from(new Set([
|
|
process.env.AI_WORKSPACE_HUB_TOKEN,
|
|
process.env.NODEDC_INTERNAL_ACCESS_TOKEN,
|
|
].map((value) => cleanString(value, 1000)).filter(Boolean)));
|
|
return {
|
|
port: Number(process.env.PORT || 18081),
|
|
wsPath: process.env.AI_WORKSPACE_HUB_WS_PATH || DEFAULT_WS_PATH,
|
|
internalAccessTokens,
|
|
};
|
|
}
|
|
|
|
function shutdown() {
|
|
clearInterval(heartbeat);
|
|
wss.close();
|
|
httpServer.close(() => {
|
|
process.exit(0);
|
|
});
|
|
setTimeout(() => process.exit(1), 5000).unref();
|
|
}
|