neoagent 3.2.1-beta.1 → 3.2.1-beta.3

package/server/services/cli/executor.js

@@ -7,6 +7,22 @@ const FORCE_KILL_GRACE_MS = 5000;
 const MAX_STDOUT_CHARS = 50000;
 const MAX_STDERR_CHARS = 10000;
 
+function abortedResult(command, cwd, startedAt = Date.now()) {
+  return {
+    exitCode: null,
+    stdout: '',
+    stderr: 'Command aborted before it started.',
+    killed: true,
+    timedOut: false,
+    aborted: true,
+    signal: null,
+    durationMs: Date.now() - startedAt,
+    pid: null,
+    command,
+    cwd,
+  };
+}
+
 function resolveDefaultShell() {
   const candidates = [
     process.env.SHELL,
@@ -63,6 +79,22 @@ function terminateProcess(proc, signal = 'SIGTERM') {
   proc.kill?.(signal) || proc.kill?.();
 }
 
+function terminateWithEscalation(proc, isActive) {
+  terminateProcess(proc, 'SIGTERM');
+  if (proc.__neoagentForceKillTimer) return;
+  proc.__neoagentForceKillTimer = setTimeout(() => {
+    if (isActive()) terminateProcess(proc, 'SIGKILL');
+  }, FORCE_KILL_GRACE_MS);
+  proc.__neoagentForceKillTimer.unref?.();
+}
+
+function clearForceKill(proc) {
+  if (proc?.__neoagentForceKillTimer) {
+    clearTimeout(proc.__neoagentForceKillTimer);
+    proc.__neoagentForceKillTimer = null;
+  }
+}
+
 function shellSupportsPipefail(shellPath) {
   const normalized = String(shellPath || '').trim().toLowerCase();
   return /(?:^|\/)(?:bash|zsh|ksh|mksh|yash)$/.test(normalized);
@@ -111,6 +143,7 @@ class CLIExecutor {
     const cwd = options.cwd || process.env.HOME;
     const timeout = clampTimeout(options.timeout, DEFAULT_TIMEOUT_MS);
     const stdinInput = options.stdinInput;
+    if (options.signal?.aborted) return abortedResult(command, cwd);
 
     return new Promise((resolve) => {
       let stdout = '';
@@ -131,6 +164,13 @@ class CLIExecutor {
       const pid = proc.pid;
       this.activeProcesses.set(pid, proc);
       options.onSpawn?.(pid);
+      const onAbort = () => {
+        killed = true;
+        proc.__neoagentKilled = true;
+        proc.__neoagentKillReason = 'aborted';
+        terminateWithEscalation(proc, () => this.activeProcesses.get(pid) === proc);
+      };
+      options.signal?.addEventListener('abort', onAbort, { once: true });
 
       proc.stdout.on('data', (data) => {
         stdout += data.toString();
@@ -156,14 +196,14 @@ class CLIExecutor {
         timedOut = true;
         proc.__neoagentKilled = true;
         proc.__neoagentKillReason = 'timeout';
-        terminateProcess(proc, 'SIGTERM');
-        setTimeout(() => {
-          if (!proc.killed) terminateProcess(proc, 'SIGKILL');
-        }, FORCE_KILL_GRACE_MS);
+        terminateWithEscalation(proc, () => this.activeProcesses.get(pid) === proc);
       }, timeout);
+      timer.unref?.();
 
       proc.on('close', (code, signal) => {
         clearTimeout(timer);
+        clearForceKill(proc);
+        options.signal?.removeEventListener('abort', onAbort);
         this.activeProcesses.delete(pid);
         const durationMs = Date.now() - startedAt;
 
@@ -173,6 +213,7 @@ class CLIExecutor {
           stderr: truncateOutput(stderr.trim(), MAX_STDERR_CHARS),
           killed: killed || proc.__neoagentKilled === true,
           timedOut: timedOut || proc.__neoagentKillReason === 'timeout',
+          aborted: proc.__neoagentKillReason === 'aborted',
           signal: signal || null,
           durationMs,
           pid,
@@ -183,6 +224,8 @@ class CLIExecutor {
 
       proc.on('error', (err) => {
         clearTimeout(timer);
+        clearForceKill(proc);
+        options.signal?.removeEventListener('abort', onAbort);
         this.activeProcesses.delete(pid);
         resolve({
           exitCode: -1,
@@ -204,6 +247,7 @@ class CLIExecutor {
   async executeInteractive(command, inputs = [], options = {}) {
     const cwd = options.cwd || process.env.HOME;
     const timeout = clampTimeout(options.timeout, DEFAULT_INTERACTIVE_TIMEOUT_MS);
+    if (options.signal?.aborted) return abortedResult(command, cwd);
 
     return new Promise((resolve) => {
       let output = '';
@@ -231,6 +275,13 @@ class CLIExecutor {
       const pid = proc.pid;
       this.activeProcesses.set(pid, proc);
       options.onSpawn?.(pid);
+      const onAbort = () => {
+        killed = true;
+        proc.__neoagentKilled = true;
+        proc.__neoagentKillReason = 'aborted';
+        terminateWithEscalation(proc, () => this.activeProcesses.get(pid) === proc);
+      };
+      options.signal?.addEventListener('abort', onAbort, { once: true });
 
       proc.onData((data) => {
         output += data;
@@ -256,11 +307,14 @@ class CLIExecutor {
         timedOut = true;
         proc.__neoagentKilled = true;
         proc.__neoagentKillReason = 'timeout';
-        proc.kill();
+        terminateWithEscalation(proc, () => this.activeProcesses.get(pid) === proc);
       }, timeout);
+      timer.unref?.();
 
       proc.onExit(({ exitCode, signal }) => {
         clearTimeout(timer);
+        clearForceKill(proc);
+        options.signal?.removeEventListener('abort', onAbort);
         this.activeProcesses.delete(pid);
 
         const cleanOutput = output.replace(/\x1B\[[0-9;]*[a-zA-Z]/g, '').trim();
@@ -270,6 +324,7 @@ class CLIExecutor {
           stderr: '',
           killed: killed || proc.__neoagentKilled === true,
           timedOut: timedOut || proc.__neoagentKillReason === 'timeout',
+          aborted: proc.__neoagentKillReason === 'aborted',
           signal: typeof signal === 'number' ? String(signal) : signal || null,
           durationMs: Date.now() - startedAt,
           pid,
@@ -286,8 +341,7 @@ class CLIExecutor {
     if (proc) {
       proc.__neoagentKilled = true;
       proc.__neoagentKillReason = reason;
-      terminateProcess(proc, 'SIGTERM');
-      this.activeProcesses.delete(pid);
+      terminateWithEscalation(proc, () => this.activeProcesses.get(pid) === proc);
       return true;
     }
     return false;
@@ -301,9 +355,8 @@ class CLIExecutor {
     for (const [pid, proc] of this.activeProcesses) {
       proc.__neoagentKilled = true;
       proc.__neoagentKillReason = reason;
-      terminateProcess(proc, 'SIGTERM');
+      terminateWithEscalation(proc, () => this.activeProcesses.get(pid) === proc);
     }
-    this.activeProcesses.clear();
   }
 }
 

package/server/services/desktop/gateway.js

@@ -1,3 +1,5 @@
+'use strict';
+
 const { WebSocketServer } = require('ws');
 const {
   DESKTOP_COMPANION_WS_PATH,
@@ -15,6 +17,9 @@ const {
 const UPGRADE_RATE_LIMIT_WINDOW_MS = 60 * 1000;
 const UPGRADE_RATE_LIMIT_MAX_ATTEMPTS = 30;
 const UPGRADE_RATE_LIMIT_ENTRY_TTL_MS = 10 * 60 * 1000;
+const UPGRADE_AUTH_TIMEOUT_MS = 5000;
+const HELLO_TIMEOUT_MS = 5000;
+const MAX_HELLO_BYTES = 64 * 1024;
 
 function rejectUpgrade(socket, statusCode, message) {
   try {
@@ -88,6 +93,8 @@ function bindDesktopCompanionGateway(httpServer, app, sessionMiddleware, streamH
   });
   const upgradeAttempts = new Map();
   const upgradeThrottleObserver = createUpgradeThrottleObserver();
+  let closing = false;
+  let closePromise = null;
 
   if (app?.locals) {
     app.locals.getDesktopGatewayRateLimitSnapshot = () => upgradeThrottleObserver.snapshot();
@@ -115,7 +122,7 @@ function bindDesktopCompanionGateway(httpServer, app, sessionMiddleware, streamH
     return true;
   }
 
-  httpServer.on('upgrade', (req, socket, head) => {
+  const handleUpgrade = (req, socket, head) => {
     let url;
     try {
       url = new URL(req.url, 'http://localhost');
@@ -125,6 +132,10 @@ function bindDesktopCompanionGateway(httpServer, app, sessionMiddleware, streamH
     if (url.pathname !== DESKTOP_COMPANION_WS_PATH) {
       return;
     }
+    if (closing) {
+      rejectUpgrade(socket, 503, 'Service Unavailable');
+      return;
+    }
 
     const remoteAddress = remoteAddressFromRequest(req);
     if (!allowUpgradeAttempt(remoteAddress)) {
@@ -133,7 +144,16 @@ function bindDesktopCompanionGateway(httpServer, app, sessionMiddleware, streamH
       return;
     }
 
+    const authTimer = setTimeout(() => {
+      rejectUpgrade(socket, 504, 'Gateway Timeout');
+    }, UPGRADE_AUTH_TIMEOUT_MS);
+    authTimer.unref?.();
     sessionMiddleware(req, {}, (err) => {
+      clearTimeout(authTimer);
+      if (closing || socket.destroyed) {
+        if (!socket.destroyed) rejectUpgrade(socket, 503, 'Service Unavailable');
+        return;
+      }
       if (err) {
         rejectUpgrade(socket, 500, 'Session Error');
         return;
@@ -155,11 +175,18 @@ function bindDesktopCompanionGateway(httpServer, app, sessionMiddleware, streamH
           if (!initialized) {
             try { ws.close(1008, 'Desktop companion hello timed out'); } catch {}
           }
-        }, 5000);
+        }, HELLO_TIMEOUT_MS);
+        helloTimer.unref?.();
+        const clearHelloTimer = () => clearTimeout(helloTimer);
+        ws.once('close', clearHelloTimer);
+        ws.once('error', clearHelloTimer);
 
         ws.once('message', (data) => {
           clearTimeout(helloTimer);
           try {
+            if (Buffer.byteLength(data) > MAX_HELLO_BYTES) {
+              throw Object.assign(new Error('Desktop companion hello is too large.'), { status: 413 });
+            }
             const message = parseDesktopMessage(data);
             if (!isDesktopCompanionHello(message)) {
               throw Object.assign(new Error('Desktop companion hello is required.'), { status: 400 });
@@ -227,11 +254,21 @@ function bindDesktopCompanionGateway(httpServer, app, sessionMiddleware, streamH
         });
       });
     });
-  });
+  };
+  httpServer.on('upgrade', handleUpgrade);
 
   if (app?.locals) {
     app.locals.desktopCompanionGateway = {
-      close: () => new Promise((resolve) => wss.close(() => resolve())),
+      close: () => {
+        if (closePromise) return closePromise;
+        closing = true;
+        httpServer.removeListener('upgrade', handleUpgrade);
+        for (const client of wss.clients) {
+          try { client.terminate(); } catch {}
+        }
+        closePromise = new Promise((resolve) => wss.close(() => resolve()));
+        return closePromise;
+      },
     };
   }
 

package/server/services/desktop/protocol.js

@@ -1,3 +1,5 @@
+'use strict';
+
 const DESKTOP_COMPANION_WS_PATH = '/api/desktop/ws';
 
 const DESKTOP_COMMANDS = Object.freeze({
@@ -17,6 +19,7 @@ const DESKTOP_COMMANDS = Object.freeze({
   GET_TREE: 'getTree',
   PAUSE_CONTROL: 'pauseControl',
   EXECUTE_COMMAND: 'executeCommand',
+  CANCEL_COMMAND: 'cancelCommand',
   PING: 'ping',
   MOUSE_MOVE: 'mouseMove',
 });

package/server/services/desktop/provider.js

@@ -1,28 +1,18 @@
+'use strict';
+
 const fs = require('fs');
 const path = require('path');
 const { DATA_DIR } = require('../../../runtime/paths');
+const { writeBufferAtomic } = require('../../utils/files');
+const { decodeBase64Image } = require('../../utils/image_payload');
 const {
   DESKTOP_COMMANDS,
-  DesktopCompanionSelectionError,
   DesktopCompanionUnavailableError,
 } = require('./protocol');
 
 const SCREENSHOTS_DIR = path.join(DATA_DIR, 'screenshots');
 if (!fs.existsSync(SCREENSHOTS_DIR)) fs.mkdirSync(SCREENSHOTS_DIR, { recursive: true });
 
-function extractBase64Image(value) {
-  const text = String(value || '');
-  if (!text) return null;
-  const match = text.match(/^data:image\/(?:png|jpeg|jpg);base64,(.+)$/i);
-  return match ? match[1] : text;
-}
-
-function guessExtension(result = {}) {
-  const mime = String(result.contentType || result.mimeType || 'image/png').toLowerCase();
-  if (mime.includes('jpeg') || mime.includes('jpg')) return 'jpg';
-  return 'png';
-}
-
 class DesktopProvider {
   constructor(options = {}) {
     this.registry = options.registry;
@@ -36,24 +26,21 @@ class DesktopProvider {
     }
   }
 
-  _writeScreenshotArtifact(base64, result = {}) {
-    const buffer = Buffer.from(base64, 'base64');
-    const extension = guessExtension(result);
-    const contentType = result.contentType || (extension === 'jpg' ? 'image/jpeg' : 'image/png');
+  async _writeScreenshotArtifact(image, result = {}, options = {}) {
     if (this.artifactStore && this.userId != null) {
-      const artifact = this.artifactStore.allocateFile(this.userId, {
+      const artifact = await this.artifactStore.createBufferArtifact(this.userId, {
         kind: 'desktop-screenshot',
         backend: 'desktop-companion',
-        extension,
-        contentType,
+        extension: image.extension,
+        contentType: image.contentType,
         filenameBase: 'desktop-companion-screenshot',
+        content: image.buffer,
+        signal: options.signal,
         metadata: {
           deviceId: result.device?.deviceId || null,
           displayId: result.displayId || result.device?.activeDisplayId || null,
         },
       });
-      fs.writeFileSync(artifact.storagePath, buffer);
-      this.artifactStore.finalizeFile(artifact.artifactId, artifact.storagePath);
       return {
         screenshotPath: artifact.url,
         artifactId: artifact.artifactId,
@@ -62,9 +49,9 @@ class DesktopProvider {
       };
     }
 
-    const filename = `desktop_${Date.now()}_${Math.random().toString(16).slice(2)}.${extension}`;
+    const filename = `desktop_${Date.now()}_${Math.random().toString(16).slice(2)}.${image.extension}`;
     const fullPath = path.join(SCREENSHOTS_DIR, filename);
-    fs.writeFileSync(fullPath, buffer);
+    await writeBufferAtomic(fullPath, image.buffer, { signal: options.signal });
     return {
       screenshotPath: `/screenshots/${filename}`,
       artifactId: null,
@@ -73,13 +60,14 @@ class DesktopProvider {
     };
   }
 
-  _materialize(result) {
+  async _materialize(result, options = {}) {
     if (!result || typeof result !== 'object') return result;
     const raw = result.screenshotDataUrl || result.screenshotData || result.screenshotBase64;
     if (!raw) return result;
-    const base64 = extractBase64Image(raw);
-    if (!base64) return result;
-    const screenshot = this._writeScreenshotArtifact(base64, result);
+    const image = decodeBase64Image(raw, {
+      allowedTypes: ['image/png', 'image/jpeg'],
+    });
+    const screenshot = await this._writeScreenshotArtifact(image, result, options);
     const next = { ...result, ...screenshot };
     delete next.screenshotDataUrl;
     delete next.screenshotData;
@@ -89,16 +77,19 @@ class DesktopProvider {
 
   async _dispatch(command, payload = {}, options = {}) {
     this._assertReady();
-    try {
-      return this._materialize(
-        await this.registry.dispatch(this.userId, payload.deviceId || null, command, payload, options),
-      );
-    } catch (error) {
-      if (error instanceof DesktopCompanionSelectionError || error instanceof DesktopCompanionUnavailableError) {
-        throw error;
-      }
-      throw error;
-    }
+    const safePayload = { ...(payload || {}) };
+    const signal = options.signal || safePayload.signal || null;
+    const timeoutMs = options.timeoutMs ?? safePayload.timeoutMs;
+    delete safePayload.signal;
+    delete safePayload.timeoutMs;
+    const result = await this.registry.dispatch(
+      this.userId,
+      safePayload.deviceId || null,
+      command,
+      safePayload,
+      { ...options, signal, timeoutMs },
+    );
+    return this._materialize(result, { signal });
   }
 
   getStatus() {
@@ -121,9 +112,9 @@ class DesktopProvider {
     return this.registry.revoke(this.userId, deviceId);
   }
 
-  pauseDevice(deviceId, paused = true) {
+  pauseDevice(deviceId, paused = true, options = {}) {
     this._assertReady();
-    return this.registry.pause(this.userId, deviceId, paused);
+    return this.registry.pause(this.userId, deviceId, paused, options);
   }
 
   screenshot(options = {}) {
@@ -137,7 +128,7 @@ class DesktopProvider {
 
   stopStream(options = {}) {
     this._assertReady();
-    return this.registry.stopStream(this.userId, options.deviceId || null);
+    return this.registry.stopStream(this.userId, options.deviceId || null, options);
   }
 
   observe(options = {}) {
@@ -192,6 +183,12 @@ class DesktopProvider {
       stdin_input: options.stdinInput || null,
       pty: options.pty === true,
       inputs: options.inputs || [],
+      signal: options.signal,
+    }, {
+      signal: options.signal,
+      timeoutMs: Number(options.timeout || 0) > 0
+        ? Number(options.timeout) + 10_000
+        : undefined,
     });
   }
 }