@testdriverai/runner 7.10.0-canary.0 → 7.10.1-canary

package/lib/ably-service.js

@@ -40,16 +40,16 @@ function getLocalVersion() {
 
 /**
  * Upload a screenshot to S3 via presigned URL from the API
- * 
+ *
  * @param {string} apiRoot - API base URL
  * @param {string} apiKey - Team API key
  * @param {string} sandboxId - Sandbox ID
- * @param {string} base64 - Base64-encoded image data
+ * @param {string|Buffer} payload - PNG bytes as a Buffer (preferred) or base64 string
  * @param {string} label - Label for the screenshot (e.g., 'screenshot', 'before', 'after')
  * @returns {Promise<{ s3Key: string, fileName: string } | null>}
  */
-async function uploadToS3(apiRoot, apiKey, sandboxId, base64, label = 'screenshot') {
-  if (!base64 || !apiRoot || !apiKey) return null;
+async function uploadToS3(apiRoot, apiKey, sandboxId, payload, label = 'screenshot') {
+  if (!payload || !apiRoot || !apiKey) return null;
 
   try {
     const fileName = `${label}-${Date.now()}.png`;
@@ -68,8 +68,10 @@ async function uploadToS3(apiRoot, apiKey, sandboxId, base64, label = 'screensho
       return null;
     }
 
-    // Upload to S3 via presigned URL
-    const pngBuffer = Buffer.from(base64, 'base64');
+    // Upload to S3 via presigned URL. Accept either a Buffer (no decode) or
+    // a base64 string (legacy callers) — Buffer is preferred to avoid the
+    // wasteful PNG → base64 → PNG round-trip.
+    const pngBuffer = Buffer.isBuffer(payload) ? payload : Buffer.from(payload, 'base64');
     const url = new URL(uploadUrlResponse.uploadUrl);
     const transport = url.protocol === 'https:' ? https : http;
 
@@ -136,6 +138,41 @@ async function httpPost(apiRoot, path, body, extraHeaders = {}) {
   });
 }
 
+// ─── EC2 Metadata Helpers ────────────────────────────────────────────────────
+
+/**
+ * Fetch the public IPv4 address from the EC2 Instance Metadata Service (IMDSv1).
+ * Returns null if not running on EC2, instance has no public IP, or request times out.
+ * Non-blocking: uses a 2-second timeout so it never delays runner startup.
+ *
+ * @returns {Promise<string|null>}
+ */
+function getEc2PublicIp() {
+  return new Promise((resolve) => {
+    const req = http.request(
+      {
+        hostname: '169.254.169.254',
+        port: 80,
+        path: '/latest/meta-data/public-ipv4',
+        method: 'GET',
+        timeout: 2000,
+      },
+      (res) => {
+        if (res.statusCode !== 200) {
+          resolve(null);
+          return;
+        }
+        let data = '';
+        res.on('data', (chunk) => { data += chunk; });
+        res.on('end', () => resolve(data.trim() || null));
+      }
+    );
+    req.on('error', () => resolve(null));
+    req.on('timeout', () => { req.destroy(); resolve(null); });
+    req.end();
+  });
+}
+
 // ─── Ably Service Class ──────────────────────────────────────────────────────
 
 class AblyService extends EventEmitter {
@@ -164,6 +201,7 @@ class AblyService extends EventEmitter {
     this._sessionChannel = null;
     this._connected = false;
     this._statsInterval = null;
+    this._closed = false; // set to true in close() to suppress post-close events
   }
 
   /**
@@ -185,12 +223,12 @@ class AblyService extends EventEmitter {
         const text = typeof msg === 'string' ? msg : String(msg);
         const isError = /error/i.test(text) || /\[Level 1\]/i.test(text);
         if (isError) {
-          console.error(`[Ably runner] ${text}`);
+          console.error('[Ably runner] Ably SDK error', { sandboxId: self._sandboxId, raw: text });
           Sentry.withScope((scope) => {
             scope.setTag('ably.client', 'runner');
             scope.setTag('sandbox.id', self._sandboxId);
             scope.setContext('ably_log', { raw: text });
-            Sentry.captureMessage(`[Ably runner] ${text}`, 'error');
+            Sentry.captureMessage('Ably SDK error (runner)', 'error');
           });
         }
         Sentry.addBreadcrumb({
@@ -367,7 +405,10 @@ class AblyService extends EventEmitter {
       if (current === 'disconnected') {
         this._connected = false;
         this.emit('log', `Realtime connection: ${previous} → ${current}${reasonMsg ? ' — ' + reasonMsg : ''}${retryIn ? ' (retryIn=' + retryIn + 'ms)' : ''}`);
-        this.emit('log', 'Ably disconnected — will auto-reconnect');
+        // Suppress auto-reconnect message when we've already initiated a clean close
+        if (!this._closed) {
+          this.emit('log', 'Ably disconnected — will auto-reconnect');
+        }
       } else if (current === 'connected' && previous !== 'initialized') {
         if (!this._connected) {
           this._connected = true;
@@ -390,8 +431,8 @@ class AblyService extends EventEmitter {
         this.emit('log', `Realtime connection: ${previous} → ${current}${reasonMsg ? ' — ' + reasonMsg : ''}`);
       }
 
-      // Capture exceptions for bad states
-      if (current === 'failed' || current === 'suspended' || (current === 'disconnected' && reason)) {
+      // Capture exceptions for bad states — skip if we're in a clean close sequence
+      if (!this._closed && (current === 'failed' || current === 'suspended' || (current === 'disconnected' && reason))) {
         Sentry.withScope((scope) => {
           scope.setTag('ably.client', 'runner');
           scope.setTag('ably.state', current);
@@ -477,12 +518,22 @@ class AblyService extends EventEmitter {
     this.emit('log', 'Listening for commands on Ably');
 
     // Signal readiness to SDK — commands sent before this would be lost
+
+    // Best-effort: fetch the public IP from EC2 instance metadata so the SDK can
+    // construct the VNC preview URL without relying solely on API-side IP lookup.
+    // Times out after 2s and resolves null if not on EC2 or instance has no public IP.
+    const ec2PublicIp = await getEc2PublicIp();
+    if (ec2PublicIp) {
+      this.emit('log', `EC2 public IP: ${ec2PublicIp}`);
+    }
+
     const readyPayload = {
       type: 'runner.ready',
       os: process.platform === 'win32' ? 'windows' : 'linux',
       sandboxId: this._sandboxId,
       runnerVersion: getLocalVersion() || 'unknown',
       timestamp: Date.now(),
+      ip: ec2PublicIp || undefined,
     };
     if (this._updateInfo) {
       readyPayload.update = {
@@ -579,6 +630,10 @@ class AblyService extends EventEmitter {
    * so that missed commands are actually executed.
    */
   async _recoverFromDiscontinuity() {
+    // Don't attempt recovery after an intentional close — subscriptions are
+    // nulled out in close() and the channel is being torn down
+    if (this._closed) return;
+
     const subs = [
       { name: 'command', sub: this._commandSubscription, handler: this._onCommandMsg },
       { name: 'control', sub: this._controlSubscription, handler: this._onControlMsg },
@@ -668,6 +723,9 @@ class AblyService extends EventEmitter {
    * Disconnect from Ably and clean up.
    */
   async close() {
+    if (this._closed) return; // Prevent double-close
+    this._closed = true;
+
     this.emit('log', 'Closing realtime service...');
 
     this._stopReadySignal();
@@ -677,6 +735,11 @@ class AblyService extends EventEmitter {
       this._statsInterval = null;
     }
 
+    // Null out subscription refs so _recoverFromDiscontinuity won't replay
+    // messages on a channel that is about to be torn down
+    this._commandSubscription = null;
+    this._controlSubscription = null;
+
     try {
       if (this._sessionChannel) this._sessionChannel.detach();
     } catch { }

package/lib/automation.js

@@ -52,6 +52,26 @@ const PY_IMPORT = IS_LINUX
   ? "import os; os.environ.setdefault('DISPLAY', ':0'); import pyautogui, sys; pyautogui.FAILSAFE = False; "
   : 'import pyautogui, sys; pyautogui.FAILSAFE = False; ';
 
+/**
+ * Read width/height from a PNG buffer's IHDR chunk without decoding the image.
+ * PNG layout: 8-byte signature, then the IHDR chunk whose 4-byte width and
+ * height live at byte offsets 16 and 20 (big-endian). Returns {} if the buffer
+ * doesn't look like a PNG so callers degrade to the resize fallback.
+ * @param {Buffer} buffer
+ * @returns {{ width?: number, height?: number }}
+ */
+function readPngDimensions(buffer) {
+  const PNG_SIGNATURE = '89504e470d0a1a0a';
+  if (!buffer || buffer.length < 24 ||
+      buffer.slice(0, 8).toString('hex') !== PNG_SIGNATURE) {
+    return {};
+  }
+  return {
+    width: buffer.readUInt32BE(16),
+    height: buffer.readUInt32BE(20),
+  };
+}
+
 /**
  * Run a pyautogui Python script via subprocess.
  * @param {string} script — Python code (pyautogui + sys are already imported via PY_IMPORT prefix)
@@ -150,6 +170,10 @@ class ShellSession {
       // Clear stderr buffer
       this._stderrLines = [];
 
+      // Capture the process reference now — the exit handler may set
+      // this._process to null before the Promise resolves/times out.
+      const proc = this._process;
+
       return await new Promise((resolve, reject) => {
         const timeoutMs = timeout * 1000;
         const timer = setTimeout(() => {
@@ -208,11 +232,11 @@ class ShellSession {
 
         const cleanup = () => {
           clearTimeout(timer);
-          this._process.stdout.removeListener('data', onData);
+          proc.stdout.removeListener('data', onData);
         };
 
-        this._process.stdout.on('data', onData);
-        this._process.stdin.write(fullCommand);
+        proc.stdout.on('data', onData);
+        proc.stdin.write(fullCommand);
       });
     } catch (err) {
       return {
@@ -587,20 +611,26 @@ class Automation extends EventEmitter {
         // For extract/remember: capture screenshot, upload to S3, return s3Key
         // SDK will then call API with the s3Key
         // For system.screenshot (normalized to 'screenshot'): same flow (return s3Key instead of base64)
-        
-        // Capture screenshot locally
+
+        // Capture screenshot locally — buffer variant avoids an unnecessary
+        // round-trip through base64 (PNG buffer → base64 → PNG buffer) that
+        // doubles peak memory and adds ~10–30 ms of pure encoding cost on
+        // large screenshots.
         console.log('[automation] Step 1: Capturing screenshot...');
-        const screenshot = await this._captureScreenshot();
-        console.log(`[automation] Step 2: Screenshot captured, size: ${screenshot.length} bytes (base64)`);
-        const buffer = Buffer.from(screenshot, 'base64');
-        console.log(`[automation] Step 3: Buffer created, size: ${buffer.length} bytes`);
+        const buffer = await this._captureScreenshotBuffer();
+        console.log(`[automation] Step 2: Screenshot captured, size: ${buffer.length} bytes`);
+
+        // Read the PNG dimensions from the IHDR header (bytes 16–24) without
+        // decoding. The SDK uses these to decide whether it can pass the s3Key
+        // straight to the API (no resize needed) or must download + resize.
+        const dimensions = readPngDimensions(buffer);
 
         // Upload screenshot to S3 and return key
-        console.log('[automation] Step 4: Uploading to S3...');
+        console.log('[automation] Step 3: Uploading to S3...');
         const s3Key = await this._uploadToS3(buffer, this._sandboxId, 'image/png');
-        console.log(`[automation] Step 5: Upload complete, s3Key: ${s3Key}`);
-        
-        return { s3Key };
+        console.log(`[automation] Step 4: Upload complete, s3Key: ${s3Key}`);
+
+        return { s3Key, ...dimensions };
       }
 
       case 'ping':
@@ -664,7 +694,15 @@ class Automation extends EventEmitter {
 
   // ── Screenshot (highest quality PNG, via pyautogui → temp file → sharp) ──
 
-  async _captureScreenshot() {
+  /**
+   * Capture a screenshot and return the raw PNG buffer.
+   *
+   * Prefer this when the next consumer of the bytes is binary (e.g. an HTTP
+   * PUT to S3): it avoids the wasteful PNG → base64 → PNG round-trip that
+   * `_captureScreenshot()` does for backwards compatibility with callers
+   * that still expect a base64 string.
+   */
+  async _captureScreenshotBuffer() {
     const sharp = require('sharp');
     const maxAttempts = 3;
 
@@ -702,11 +740,15 @@ class Automation extends EventEmitter {
           if (maxPixel <= 1) {
             console.warn(`[automation] Screenshot attempt ${attempt}/${maxAttempts}: image is all black (max pixel=${maxPixel})`);
             if (attempt < maxAttempts) {
-              // Try to heal: poke the display to trigger a redraw
+              // Try to heal: poke the display to trigger a redraw.
+              // NOTE: do NOT send the `super` key here — on Linux desktops it's
+              // bound to the Activities overview / app launcher, which steals
+              // focus from (and visually hides) the active window. That made
+              // Chrome appear to "go missing" right after a flaky black capture.
+              // The xset display-wake calls are what actually recover the screen.
               try {
                 await runPyAutoGUI(
                   "import subprocess; " +
-                  "subprocess.run(['xdotool', 'key', '--clearmodifiers', 'super'], timeout=5); " +
                   "subprocess.run(['xset', 's', 'off'], timeout=5); " +
                   "subprocess.run(['xset', 's', 'noblank'], timeout=5); " +
                   "subprocess.run(['xset', '-dpms'], timeout=5)",
@@ -722,8 +764,7 @@ class Automation extends EventEmitter {
           }
         }
 
-        const buffer = await image.png({ compressionLevel: 0 }).toBuffer();
-        return buffer.toString('base64');
+        return await image.png({ compressionLevel: 0 }).toBuffer();
       } finally {
         // Clean up temp file
         try { fs.unlinkSync(tmpFile); } catch {}
@@ -731,6 +772,14 @@ class Automation extends EventEmitter {
     }
   }
 
+  async _captureScreenshot() {
+    // Kept for backwards compatibility with callers (e.g. Ably channel
+    // messages) that still want a base64 string. New code paths should
+    // prefer `_captureScreenshotBuffer()` to avoid the round-trip.
+    const buffer = await this._captureScreenshotBuffer();
+    return buffer.toString('base64');
+  }
+
   // ── Focus window (platform-specific) ───────────────────────────────
 
   async _focusWindow(data) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@testdriverai/runner",
-  "version": "7.10.0-canary.0",
+  "version": "7.10.1-canary",
   "description": "TestDriver Runner - Ably-based remote automation agent with Node.js automation",
   "main": "index.js",
   "bin": {

package/sandbox-agent.js

@@ -33,6 +33,7 @@
  */
 require('dotenv').config();
 const Sentry = require('@sentry/node');
+const crypto = require('crypto');
 const fs = require('fs');
 const os = require('os');
 const path = require('path');
@@ -181,6 +182,7 @@ function log(msg) {
 // ─── Main ────────────────────────────────────────────────────────────────────
 
 async function main() {
+  let rootSpan = null;
   log('TestDriver Sandbox Agent starting...');
   log(`Runner version: ${getLocalVersion() || 'unknown'}`);
 
@@ -207,13 +209,45 @@ async function main() {
         // Sample all runner-initiated traces (shouldn't happen often)
         return 1.0;
       },
-      environment: config.sentryEnvironment || process.env.TD_ENV || 'production',
+      environment: config.sentryChannel || process.env.TD_CHANNEL || 'stable',
       release: version,
       serverName: os.hostname(),
     });
     if (config.sentryChannel || process.env.TD_CHANNEL) {
       Sentry.setTag('channel', config.sentryChannel || process.env.TD_CHANNEL);
     }
+
+    // Create a session-level root span for distributed tracing.
+    // Uses deterministic traceId = MD5(sessionId) so all components
+    // (SDK, API, Runner, Worker) share the same trace.
+    const sessionId = config.sessionId;
+    if (sessionId) {
+      const traceId = crypto.createHash('md5').update(sessionId).digest('hex');
+      const spanId = crypto.randomBytes(8).toString('hex');
+      const sentryTraceHeader = `${traceId}-${spanId}-1`;
+      const baggageHeader = `sentry-trace_id=${traceId},sentry-sampled=true`;
+
+      Sentry.continueTrace(
+        { sentryTrace: sentryTraceHeader, baggage: baggageHeader },
+        () => {
+          rootSpan = Sentry.startInactiveSpan({
+            name: 'runner.session',
+            op: 'session',
+            forceTransaction: true,
+            attributes: {
+              'sandbox.id': config.sandboxId,
+              'session.id': sessionId,
+            },
+          });
+        },
+      );
+
+      Sentry.setTag('session', sessionId);
+      Sentry.setTag('trace_id', traceId);
+      log(`Distributed tracing: session=${sessionId} traceId=${traceId}`);
+    }
+
+    Sentry.setTag('sandbox', config.sandboxId);
     log('Sentry initialized');
   }
   log(`Sandbox ID: ${config.sandboxId}`);
@@ -251,6 +285,9 @@ async function main() {
 
   const shutdown = async () => {
     log('Shutting down...');
+    if (rootSpan) {
+      try { rootSpan.end(); } catch (e) { /* ignore */ }
+    }
     await ablyService.close();
     automation.cleanup();
     await Sentry.close(2000);

package/scripts-desktop/start-desktop.sh

@@ -168,9 +168,12 @@ else
 fi
 
 # ─── Force a screen redraw ────────────────────────────────────────────────────
-# Send a synthetic keypress so the display is definitely "active" and any
-# screen blanking timers are reset.
-xdotool key --clearmodifiers super 2>/dev/null || true
+# Wake the display and reset any blanking timers WITHOUT synthetic input.
+# Do NOT send the `super` key here — on Linux desktops it's bound to the
+# Activities overview / app launcher, which steals focus from (and visually
+# hides) the active window. Same reasoning as the black-screenshot recovery
+# in lib/automation.js.
+xset s reset 2>/dev/null || true
 sleep 1
 
 echo "[start-desktop] Desktop environment ready"