@anusornneal/chat-relay 0.7.2 → 0.8.0

package/agent/local-agent.mjs

@@ -1,8 +1,12 @@
 import WebSocket from "ws";
-import { DesktopManager } from "./desktop-manager.mjs";
+import { DesktopManager, createDesktopPlatformAdapter } from "./desktop-manager.mjs";
 import { FileManager } from "./file-manager.mjs";
 import { ProcessManager } from "./process-manager.mjs";
 import { TerminalManager } from "./terminal-manager.mjs";
+import { CapabilityScheduler } from "./capability-scheduler.mjs";
+import { AgentConnectionState } from "./connection-state.mjs";
+import { buildAgentHello } from "./protocol.mjs";
+import { AgentLifecycle, AGENT_RESTART_EXIT_CODE } from "./lifecycle.mjs";
 
 const relayUrl = process.env.RELAY_URL;
 const agentToken = process.env.AGENT_TOKEN;
@@ -11,16 +15,59 @@ const agentName = process.env.AGENT_NAME || agentId;
 const reconnectMs = Number(process.env.RECONNECT_MS ?? 2000);
 const terminalEnabled = process.env.TERMINAL_ENABLED === "1";
 const desktopEnabled = process.env.DESKTOP_ENABLED === "1";
+const desktopAdapter = createDesktopPlatformAdapter({ platform: process.platform });
+const agentHello = buildAgentHello({
+  platform: desktopAdapter.platform,
+  terminalEnabled,
+  desktopEnabled: desktopEnabled && desktopAdapter.supported,
+});
 const MAX_RESPONSE_BYTES = 60 * 1024;
 const terminals = new TerminalManager({
   batchConcurrency: process.env.TERMINAL_BATCH_CONCURRENCY,
   maxQueuedJobs: process.env.TERMINAL_BATCH_MAX_QUEUED,
 });
+const scheduler = new CapabilityScheduler({
+  maxQueued: process.env.AGENT_MAX_QUEUED,
+  queueTimeoutMs: process.env.AGENT_QUEUE_TIMEOUT_MS,
+  fileConcurrency: process.env.AGENT_FILE_CONCURRENCY,
+  processConcurrency: process.env.AGENT_PROCESS_CONCURRENCY,
+  terminalExecConcurrency: process.env.AGENT_TERMINAL_EXEC_CONCURRENCY,
+  terminalControlConcurrency: process.env.AGENT_TERMINAL_CONTROL_CONCURRENCY,
+  desktopReadConcurrency: process.env.AGENT_DESKTOP_READ_CONCURRENCY,
+});
+const connectionState = new AgentConnectionState({
+  baseReconnectMs: reconnectMs,
+  maxReconnectMs: process.env.RECONNECT_MAX_MS,
+  heartbeatMs: process.env.AGENT_HEARTBEAT_MS,
+});
 const files = new FileManager(process.env.ALLOWED_ROOTS);
 const processes = new ProcessManager();
-const desktop = new DesktopManager({ enabled: desktopEnabled });
+const desktop = new DesktopManager({
+  enabled: desktopEnabled,
+  adapter: desktopAdapter,
+  controlMaxQueued: process.env.DESKTOP_CONTROL_MAX_QUEUED,
+  controlQueueTimeoutMs: process.env.AGENT_QUEUE_TIMEOUT_MS,
+});
 const recentCalls = [];
 let reauthorizationRequired = false;
+let activeSocket = null;
+let reconnectTimer = null;
+let heartbeatTimer = null;
+let stopping = false;
+
+function currentWorkSummary() {
+  const terminal = terminals.observability();
+  const terminalExec = scheduler.snapshot().terminalExec || {};
+  return {
+    activeSessions: terminal.sessions.filter((item) => item.status === "running").length,
+    activeBatchJobs: terminal.activeExecJobs,
+    queuedBatchJobs: terminal.queuedJobs,
+    activeTerminalExecs: terminalExec.active,
+    queuedTerminalExecs: terminalExec.queued,
+  };
+}
+
+const lifecycle = new AgentLifecycle({ workSummary: currentWorkSummary });
 
 if (!relayUrl || !agentToken) {
   console.error("RELAY_URL and AGENT_TOKEN are required");
@@ -30,8 +77,8 @@ if (!relayUrl || !agentToken) {
 const wsUrl = relayUrl.replace(/^http/, "ws").replace(/\/$/, "") +
   `/agent?agentId=${encodeURIComponent(agentId)}`;
 
-function serializeResponse(requestId, payload) {
-  const response = JSON.stringify({ requestId, payload });
+function serializeResponse(requestId, payload, meta) {
+  const response = JSON.stringify({ requestId, payload, ...(meta ? { meta } : {}) });
   if (Buffer.byteLength(response, "utf8") <= MAX_RESPONSE_BYTES) return response;
   return JSON.stringify({
     requestId,
@@ -45,8 +92,11 @@ function serializeResponse(requestId, payload) {
 
 async function handlePayload(payload) {
   if (!payload || typeof payload !== "object") return payload;
+  const action = String(payload.action ?? "");
+  const lifecycleGate = lifecycle.guard(action);
+  if (lifecycleGate) return lifecycleGate;
 
-  switch (payload.action) {
+  switch (action) {
     case "ping":
       return { ok: true, action: "pong", at: new Date().toISOString(), agentId, agentName };
 
@@ -56,12 +106,37 @@ async function handlePayload(payload) {
         agentName,
         terminalEnabled,
         terminalBatch: terminals.getBatchConfig(),
+        concurrency: scheduler.snapshot(),
         desktop: desktop.getConfig(),
         allowedRoots: files.getRoots(),
-        reconnectMs,
+        reconnectMs: connectionState.baseReconnectMs,
+        connection: { ...connectionState.snapshot(), processId: process.pid },
+        lifecycle: lifecycle.snapshot(),
+        terminalContinuity: {
+          processId: process.pid,
+          activeSessions: terminals.list().filter((session) => session.status === "running").length,
+        },
+        protocol: {
+          protocolVersion: agentHello.protocolVersion,
+          agentVersion: agentHello.agentVersion,
+          platform: agentHello.platform,
+          arch: agentHello.arch,
+          capabilities: agentHello.capabilities,
+        },
       });
     case "agent.recentCalls":
       return recentCalls.slice(-(Math.min(Math.max(Number(payload.limit) || 50, 1), 100)));
+    case "agent.lifecycle.status":
+      return { ok: true, lifecycle: lifecycle.snapshot() };
+    case "agent.lifecycle.drain":
+      return lifecycle.drain();
+    case "agent.lifecycle.resume":
+      return lifecycle.resume();
+    case "agent.lifecycle.restart": {
+      const result = lifecycle.requestRestart();
+      if (result.ok) setTimeout(restartAgentProcess, 150);
+      return result;
+    }
 
     case "fs.stat":
       return files.stat(payload.path);
@@ -71,6 +146,10 @@ async function handlePayload(payload) {
       return files.read(payload.path, payload.offset, payload.length, payload.maxBytes);
     case "fs.readMany":
       return files.readMany(payload.paths, payload.maxTotalBytes);
+    case "fs.batch":
+      return files.batch(payload.operations, payload.maxTotalBytes);
+    case "fs.artifact":
+      return files.exportArtifact(payload.path, payload.maxBytes);
     case "fs.write":
       return files.write(payload.path, payload.content, payload.mode);
     case "fs.edit":
@@ -93,6 +172,14 @@ async function handlePayload(payload) {
 
     case "desktop.screenshot":
       return desktop.screenshot(payload);
+    case "desktop.clipboard.read":
+      return desktop.clipboardRead(payload);
+    case "desktop.clipboard.write":
+      return desktop.clipboardWrite(payload);
+    case "desktop.window.list":
+      return desktop.listWindows(payload);
+    case "desktop.window.focus":
+      return desktop.focusWindow(payload);
     case "desktop.mouse.click":
       return desktop.mouseClick(payload);
     case "desktop.keyboard.input":
@@ -142,27 +229,120 @@ async function handlePayload(payload) {
   }
 }
 
-function requireReauthorization() {
+function clearReconnectTimer() {
+  if (!reconnectTimer) return;
+  clearTimeout(reconnectTimer);
+  reconnectTimer = null;
+}
+
+function stopHeartbeat() {
+  if (!heartbeatTimer) return;
+  clearInterval(heartbeatTimer);
+  heartbeatTimer = null;
+}
+
+function sendHeartbeat(socket) {
+  if (socket !== activeSocket || socket.readyState !== WebSocket.OPEN) return;
+  const now = Date.now();
+  try {
+    socket.send(JSON.stringify({
+      control: "agent_heartbeat",
+      at: now,
+      agentId,
+      processId: process.pid,
+      heartbeatMs: connectionState.heartbeatMs,
+      lifecycle: lifecycle.snapshot(),
+      health: {
+        reconnectCount: connectionState.reconnectCount,
+        reconnectAttempt: connectionState.reconnectAttempt,
+        lastDisconnectReason: connectionState.lastDisconnectReason,
+        queues: scheduler.snapshot(),
+      },
+    }));
+    connectionState.markHeartbeat(now);
+  } catch {}
+}
+
+function startHeartbeat(socket) {
+  stopHeartbeat();
+  sendHeartbeat(socket);
+  heartbeatTimer = setInterval(() => sendHeartbeat(socket), connectionState.heartbeatMs);
+}
+
+function sendSocketResponse(socket, requestId, payload, meta) {
+  if (socket !== activeSocket || socket.readyState !== WebSocket.OPEN) return;
+  try { socket.send(serializeResponse(requestId, payload, meta)); } catch {}
+}
+
+function scheduleReconnect(code, reason) {
+  if (stopping || reauthorizationRequired) return;
+  connectionState.markDisconnected(code, reason);
+  const delay = connectionState.nextDelay();
+  connectionState.scheduleReconnect(delay);
+  clearReconnectTimer();
+  console.log(`Reconnect scheduled in ${delay}ms`);
+  reconnectTimer = setTimeout(() => {
+    reconnectTimer = null;
+    connect();
+  }, delay);
+}
+
+
+function requireProtocolUpdate(details = {}) {
+  if (stopping) return;
+  stopping = true;
+  connectionState.markStopping("protocol_incompatible");
+  clearReconnectTimer();
+  stopHeartbeat();
+  const socket = activeSocket;
+  activeSocket = null;
+  try { socket?.terminate(); } catch {}
+  desktop.close();
+  const expected = details.expectedProtocolVersion ?? "current";
+  const received = details.receivedProtocolVersion ?? agentHello.protocolVersion;
+  console.error(`Agent protocol is incompatible (local=${received}, relay=${expected}). Update Chat Relay, then reconnect.`);
+  setImmediate(() => process.exit(3));
+}
+
+function requireReauthorization(reason = "credential_revoked") {
   if (reauthorizationRequired) return;
   reauthorizationRequired = true;
+  connectionState.markReauthorization(reason);
+  clearReconnectTimer();
+  stopHeartbeat();
+  const socket = activeSocket;
+  activeSocket = null;
+  try { socket?.terminate(); } catch {}
   desktop.close();
-  process.exitCode = 2;
   console.error("Agent credential was revoked or rejected.");
   console.error('Recovery: run "chat-relay login --force", then "chat-relay remote".');
+  setImmediate(() => process.exit(2));
 }
 
 function connect() {
-  if (reauthorizationRequired) return;
+  if (stopping || reauthorizationRequired) return;
+  clearReconnectTimer();
+  connectionState.markConnecting();
   console.log(`Connecting to ${wsUrl}`);
   console.log(`Terminal access: ${terminalEnabled ? "enabled" : "disabled"}`);
-  console.log("Desktop access: " + (desktopEnabled ? "enabled" : "disabled"));
+  const desktopConfig = desktop.getConfig();
+  console.log("Desktop access: " + (!desktopEnabled ? "disabled" : desktopConfig.supported ? "enabled" : `unsupported (${desktopConfig.platform})`));
+
   const socket = new WebSocket(wsUrl, {
     headers: { Authorization: `Bearer ${agentToken}` },
   });
+  activeSocket = socket;
 
-  socket.on("open", () => console.log("Agent connected"));
+  socket.on("open", () => {
+    if (socket !== activeSocket) return;
+    socket.send(JSON.stringify(agentHello));
+    connectionState.markConnected();
+    startHeartbeat(socket);
+    console.log("Agent connected");
+  });
 
   socket.on("message", async (raw) => {
+    if (socket !== activeSocket) return;
     let message;
     try {
       message = JSON.parse(raw.toString());
@@ -171,8 +351,11 @@ function connect() {
     }
 
     if (message?.control === "credential_revoked") {
-      requireReauthorization();
-      try { socket.close(4001, "credential_revoked"); } catch {}
+      requireReauthorization("credential_revoked");
+      return;
+    }
+    if (message?.control === "protocol_incompatible") {
+      requireProtocolUpdate(message);
       return;
     }
 
@@ -182,61 +365,111 @@ function connect() {
     const action = message.payload && typeof message.payload === "object"
       ? String(message.payload.action ?? "unknown")
       : "unknown";
+    let scheduleMeta = { lane: "unknown", queueWaitMs: 0, queueDepthAtStart: 0, activeAtStart: 0 };
+    let handlerDurationMs = 0;
     try {
-      const result = await handlePayload(message.payload);
+      const result = await scheduler.run(action, async (meta = {}) => {
+        scheduleMeta = { ...scheduleMeta, ...meta };
+        const handlerStartedAt = Date.now();
+        try {
+          return await handlePayload(message.payload);
+        } finally {
+          handlerDurationMs = Math.max(0, Date.now() - handlerStartedAt);
+        }
+      });
       recentCalls.push({
         at: new Date().toISOString(),
         action,
         durationMs: Date.now() - startedAt,
+        queueWaitMs: scheduleMeta.queueWaitMs,
+        handlerDurationMs,
+        lane: scheduleMeta.lane,
         ok: !(result && typeof result === "object" && result.ok === false),
       });
       if (recentCalls.length > 100) recentCalls.shift();
-      socket.send(serializeResponse(message.requestId, result));
+      sendSocketResponse(socket, message.requestId, result, {
+        agentQueueWaitMs: scheduleMeta.queueWaitMs,
+        agentHandlerMs: handlerDurationMs,
+        lane: scheduleMeta.lane,
+      });
     } catch (error) {
       recentCalls.push({
         at: new Date().toISOString(),
         action,
         durationMs: Date.now() - startedAt,
+        queueWaitMs: scheduleMeta.queueWaitMs,
+        handlerDurationMs,
+        lane: scheduleMeta.lane,
         ok: false,
         error: error instanceof Error ? error.message : "agent_error",
       });
       if (recentCalls.length > 100) recentCalls.shift();
-      socket.send(serializeResponse(message.requestId, {
+      sendSocketResponse(socket, message.requestId, {
         ok: false,
         error: error instanceof Error ? error.message : "agent_error",
-      }));
+      });
     }
   });
 
   socket.on("unexpected-response", (_request, response) => {
     if (response.statusCode === 401 || response.statusCode === 403) {
       response.resume();
-      requireReauthorization();
-      socket.terminate();
+      requireReauthorization("credential_rejected");
     }
   });
 
   socket.on("close", (code, reason) => {
+    if (socket !== activeSocket) return;
+    activeSocket = null;
+    stopHeartbeat();
     const reasonText = reason.toString();
     console.log("Agent disconnected (" + code + ") " + reasonText);
     if (code === 4001 || reasonText === "credential_revoked") {
-      requireReauthorization();
+      requireReauthorization("credential_revoked");
+      return;
+    }
+    if (code === 4002 || reasonText === "protocol_incompatible") {
+      requireProtocolUpdate({ receivedProtocolVersion: agentHello.protocolVersion });
       return;
     }
-    if (!reauthorizationRequired) setTimeout(connect, reconnectMs);
+    scheduleReconnect(code, reasonText || "socket_closed");
   });
 
   socket.on("error", (error) => {
-    if (!reauthorizationRequired) console.error("WebSocket error:", error.message);
+    if (socket === activeSocket && !stopping && !reauthorizationRequired) {
+      console.error("WebSocket error:", error.message);
+    }
   });
 }
 
-function shutdown() {
+function restartAgentProcess() {
+  if (stopping) return;
+  stopping = true;
+  connectionState.markStopping("restart_requested");
+  clearReconnectTimer();
+  stopHeartbeat();
+  const socket = activeSocket;
+  activeSocket = null;
+  try { socket?.close(1012, "restart_requested"); } catch {}
+  desktop.close();
+  setTimeout(() => process.exit(AGENT_RESTART_EXIT_CODE), 50);
+}
+
+function shutdown(reason) {
+  if (stopping) return;
+  stopping = true;
+  connectionState.markStopping(reason);
+  clearReconnectTimer();
+  stopHeartbeat();
+  const socket = activeSocket;
+  activeSocket = null;
+  try { socket?.close(1000, "client_shutdown"); } catch {}
   desktop.close();
+  setTimeout(() => process.exit(0), 50);
 }
 
-process.once("SIGINT", shutdown);
-process.once("SIGTERM", shutdown);
-process.once("exit", shutdown);
+process.once("SIGINT", () => shutdown("SIGINT"));
+process.once("SIGTERM", () => shutdown("SIGTERM"));
+process.once("exit", () => desktop.close());
 
 connect();

package/agent/protocol.mjs

@@ -0,0 +1,67 @@
+import fs from "node:fs";
+
+export const AGENT_PROTOCOL_VERSION = 1;
+
+let cachedVersion = null;
+
+export function getAgentVersion() {
+  if (cachedVersion) return cachedVersion;
+  try {
+    const packageUrl = new URL("../package.json", import.meta.url);
+    const parsed = JSON.parse(fs.readFileSync(packageUrl, "utf8"));
+    cachedVersion = typeof parsed.version === "string" && parsed.version ? parsed.version : "0.0.0";
+  } catch {
+    cachedVersion = "0.0.0";
+  }
+  return cachedVersion;
+}
+
+export function getAgentCapabilities(options = {}) {
+  const terminalEnabled = options.terminalEnabled === true;
+  const desktopEnabled = options.desktopEnabled === true;
+  const platform = String(options.platform || process.platform);
+
+  const capabilities = [
+    "transport.heartbeat",
+    "filesystem.stat",
+    "filesystem.list",
+    "filesystem.read",
+    "filesystem.readMany",
+    "filesystem.batch",
+    "filesystem.artifact",
+    "filesystem.write",
+    "filesystem.edit",
+    "filesystem.mkdir",
+    "filesystem.move",
+    "filesystem.delete",
+    "filesystem.search",
+    "process.list",
+    "process.kill",
+  ];
+
+  if (terminalEnabled) {
+    capabilities.push("terminal.exec", "terminal.sessions", "terminal.batch");
+  }
+
+  if (desktopEnabled && platform === "win32") {
+    capabilities.push("desktop.screenshot", "desktop.control", "desktop.step");
+  }
+
+  return capabilities;
+}
+
+export function buildAgentHello(options = {}) {
+  const platform = String(options.platform || process.platform);
+  return {
+    control: "agent_hello",
+    protocolVersion: Number(options.protocolVersion ?? AGENT_PROTOCOL_VERSION),
+    agentVersion: String(options.agentVersion || getAgentVersion()),
+    platform,
+    arch: String(options.arch || process.arch),
+    capabilities: getAgentCapabilities({
+      terminalEnabled: options.terminalEnabled === true,
+      desktopEnabled: options.desktopEnabled === true,
+      platform,
+    }),
+  };
+}