@testdriverai/agent 7.9.103-canary → 7.9.104-test

package/mcp-server/src/server.ts

@@ -12,8 +12,9 @@ process.env.TD_DEBUG = "true";
 import { registerAppResource, registerAppTool, RESOURCE_MIME_TYPE } from "@modelcontextprotocol/ext-apps/server";
 import { McpServer, ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
 import type { Variables } from "@modelcontextprotocol/sdk/shared/uriTemplate.js";
-import type { CallToolResult, ReadResourceResult } from "@modelcontextprotocol/sdk/types.js";
+import type { CallToolResult, ReadResourceResult, ServerNotification, ServerRequest } from "@modelcontextprotocol/sdk/types.js";
 import * as Sentry from "@sentry/node";
 import * as fs from "fs";
 import * as os from "os";
@@ -333,11 +334,147 @@ function requireActiveSession(): { valid: true } | { valid: false; error: CallTo
   return { valid: true };
 }
 
+// =============================================================================
+// Progress reporting (MCP `notifications/progress`)
+// =============================================================================
+
+/** The `extra` argument every tool callback receives from the MCP SDK. */
+type ToolExtra = RequestHandlerExtra<ServerRequest, ServerNotification>;
+
+/**
+ * Reports progress for a long-running tool back to the client.
+ *
+ * Per the MCP spec, progress is only sent when the caller included a
+ * `progressToken` in the request's `_meta`. When no token is present this is a
+ * no-op, so callers can report unconditionally without branching.
+ *
+ * Each `report()` call increments an internal counter (indeterminate progress —
+ * we rarely know a real total ahead of time) and forwards a human-readable
+ * `message`. A `heartbeat()` helper keeps the client's idle timeout alive while
+ * a single long SDK call is in flight (e.g. provisioning a sandbox, polling a
+ * find), which is the situation that was tripping the `session_start` timeout.
+ */
+interface ProgressReporter {
+  /** Emit one progress step with an optional human-readable message. */
+  report(message?: string): void;
+  /**
+   * Emit a progress tick every `intervalMs` until the returned stop function is
+   * called. Use this around a single long `await` so the client keeps receiving
+   * activity. Always pair with the returned `stop()` in a `finally`.
+   */
+  heartbeat(message: string, intervalMs?: number): () => void;
+}
+
+const DEFAULT_HEARTBEAT_MS = 3000;
+
+function makeProgressReporter(extra: ToolExtra): ProgressReporter {
+  const progressToken = extra?._meta?.progressToken;
+
+  // No token → client did not opt into progress. Return a no-op reporter.
+  if (progressToken === undefined || progressToken === null) {
+    return {
+      report: () => {},
+      heartbeat: () => () => {},
+    };
+  }
+
+  let progress = 0;
+
+  const send = (message?: string) => {
+    progress += 1;
+    // Fire-and-forget: a failed notification must never break the tool call.
+    void extra
+      .sendNotification({
+        method: "notifications/progress",
+        params: { progressToken, progress, message },
+      })
+      .catch((err) => {
+        logger.debug("progress: sendNotification failed", { error: String(err) });
+      });
+  };
+
+  return {
+    report: (message) => send(message),
+    heartbeat: (message, intervalMs = DEFAULT_HEARTBEAT_MS) => {
+      send(message);
+      const timer = setInterval(() => send(message), intervalMs);
+      // Don't let the heartbeat keep the event loop alive on its own.
+      timer.unref?.();
+      return () => clearInterval(timer);
+    },
+  };
+}
+
+// =============================================================================
+// Cancellation (MCP `notifications/cancelled` → `extra.signal`)
+// =============================================================================
+
+/** Thrown when a tool call is aborted by the client. */
+class ToolAbortError extends Error {
+  constructor(tool: string) {
+    super(`${tool} was cancelled by the client`);
+    this.name = "ToolAbortError";
+  }
+}
+
+/** Reject as soon as `signal` aborts. Used to race against long SDK calls. */
+function rejectOnAbort(signal: AbortSignal | undefined, tool: string): { promise: Promise<never>; cleanup: () => void } {
+  if (!signal) {
+    // Never-resolving promise with a no-op cleanup — Promise.race ignores it.
+    return { promise: new Promise<never>(() => {}), cleanup: () => {} };
+  }
+  let onAbort: () => void = () => {};
+  const promise = new Promise<never>((_, reject) => {
+    onAbort = () => reject(new ToolAbortError(tool));
+    if (signal.aborted) {
+      onAbort();
+    } else {
+      signal.addEventListener("abort", onAbort, { once: true });
+    }
+  });
+  return { promise, cleanup: () => signal.removeEventListener("abort", onAbort) };
+}
+
+/**
+ * Run a long-running SDK call, but settle as soon as the client aborts.
+ *
+ * The wrapped SDK methods are not themselves signal-aware, so on abort the
+ * underlying work keeps running to completion in the background — but the tool
+ * call returns promptly with a `ToolAbortError` instead of blocking the client.
+ * Callers that hold cleanable resources (e.g. `session_start`) should catch
+ * `ToolAbortError` and tear them down.
+ */
+async function raceAbort<T>(signal: AbortSignal | undefined, tool: string, work: Promise<T>): Promise<T> {
+  if (signal?.aborted) {
+    throw new ToolAbortError(tool);
+  }
+  const { promise, cleanup } = rejectOnAbort(signal, tool);
+  try {
+    return await Promise.race([work, promise]);
+  } finally {
+    cleanup();
+  }
+}
+
+/**
+ * If `error` is a client cancellation, return a "cancelled" tool result so the
+ * caller can `return` it; otherwise return null so normal error handling (log +
+ * Sentry + rethrow) proceeds. Keeps abort out of error reporting — a user
+ * cancelling is not a failure.
+ */
+function cancelledResultOrNull(error: unknown, tool: string): CallToolResult | null {
+  if (error instanceof ToolAbortError) {
+    logger.info(`${tool}: Cancelled by client`);
+    return createToolResult(false, `${tool} was cancelled.`, { action: tool, cancelled: true });
+  }
+  return null;
+}
+
 /**
  * Create tool result with structured content for MCP App
  * Images: imageUrl (data URL) goes to structuredContent for UI to display
  * The croppedImage from find() is small (~10KB) so it's acceptable as data URL
- * 
+ *
  * If generatedCode is provided, it's appended to the text response with instructions
  * for the agent to write it to the test file.
  */
@@ -524,8 +661,9 @@ Debug mode (connect to existing sandbox):
     inputSchema: SessionStartInputSchema as any,
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
   },
-  async (params: SessionStartInput): Promise<CallToolResult> => {
+  async (params: SessionStartInput, extra: ToolExtra): Promise<CallToolResult> => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
 
     // Resolve OS with priority: explicit param > TD_OS env var > "linux" default
     // This mirrors the behavior of the Vitest hooks (hooks.mjs) which also reads TD_OS
@@ -621,10 +759,15 @@ Debug mode (connect to existing sandbox):
       // Handle sandboxId mode - connect to existing sandbox (debug-on-failure mode)
       if (params.sandboxId) {
         logger.info("session_start: Connecting to existing sandbox (debug mode)", { sandboxId: params.sandboxId });
-        await sdk.connect({
-          sandboxId: params.sandboxId,
-          keepAlive: params.keepAlive,
-        });
+        const stopHeartbeat = progress.heartbeat(`Connecting to existing sandbox ${params.sandboxId}...`);
+        try {
+          await raceAbort(extra.signal, "session_start", sdk.connect({
+            sandboxId: params.sandboxId,
+            keepAlive: params.keepAlive,
+          }));
+        } finally {
+          stopHeartbeat();
+        }
         
         // Get sandbox ID
         const instance = sdk.getInstance();
@@ -636,7 +779,8 @@ Debug mode (connect to existing sandbox):
         
         // Capture screenshot of current state
         logger.debug("session_start: Capturing screenshot of existing sandbox");
-        const screenshotBase64 = await sdk.agent.system.captureScreenBase64(1, false, true);
+        progress.report("Capturing screenshot...");
+        const screenshotBase64 = await raceAbort(extra.signal, "session_start", sdk.agent.system.captureScreenBase64(1, false, true)) as string;
         
         let screenshotResourceUri: string | undefined;
         if (screenshotBase64) {
@@ -673,11 +817,20 @@ You are now connected to the sandbox in its current state. Use find, click, type
       } else {
         logger.info("session_start: Connecting to cloud sandbox...");
       }
-      await sdk.connect({
-        reconnect: params.reconnect,
-        keepAlive: params.keepAlive,
-        ip: instanceIp,
-      });
+      {
+        const stopHeartbeat = progress.heartbeat(
+          instanceIp ? `Connecting to self-hosted instance ${instanceIp}...` : "Connecting to cloud sandbox..."
+        );
+        try {
+          await raceAbort(extra.signal, "session_start", sdk.connect({
+            reconnect: params.reconnect,
+            keepAlive: params.keepAlive,
+            ip: instanceIp,
+          }));
+        } finally {
+          stopHeartbeat();
+        }
+      }
 
       // Get sandbox ID
       const instance = sdk.getInstance();
@@ -691,12 +844,16 @@ You are now connected to the sandbox in its current state. Use find, click, type
       const provisionOptions = getProvisionOptions(params);
       let provisionCmd = "";
 
+      // Provisioning can take tens of seconds (downloading installers, booting
+      // apps). Heartbeat so the client's idle timeout keeps resetting.
+      const stopProvisionHeartbeat = progress.heartbeat(`Provisioning ${params.type}...`);
+      try {
       // Provision based on type
       switch (params.type) {
         case "chrome": {
           const chromeOpts = provisionOptions as { url: string; maximized?: boolean; guest?: boolean };
           logger.info("session_start: Provisioning Chrome", { url: chromeOpts.url });
-          await sdk.provision.chrome(chromeOpts);
+          await raceAbort(extra.signal, "session_start", sdk.provision.chrome(chromeOpts));
           provisionCmd = "provision.chrome";
           logger.debug("session_start: Chrome provisioned");
           break;
@@ -705,7 +862,7 @@ You are now connected to the sandbox in its current state. Use find, click, type
         case "chromeExtension": {
           const extOpts = provisionOptions as { extensionPath?: string; extensionId?: string; maximized?: boolean };
           logger.info("session_start: Provisioning Chrome Extension", { extensionPath: extOpts.extensionPath, extensionId: extOpts.extensionId });
-          await sdk.provision.chromeExtension(extOpts);
+          await raceAbort(extra.signal, "session_start", sdk.provision.chromeExtension(extOpts));
           provisionCmd = "provision.chromeExtension";
           logger.debug("session_start: Chrome Extension provisioned");
           break;
@@ -714,7 +871,7 @@ You are now connected to the sandbox in its current state. Use find, click, type
         case "vscode": {
           const vscodeOpts = provisionOptions as { workspace?: string; extensions?: string[] };
           logger.info("session_start: Provisioning VS Code", { workspace: vscodeOpts.workspace });
-          await sdk.provision.vscode(vscodeOpts);
+          await raceAbort(extra.signal, "session_start", sdk.provision.vscode(vscodeOpts));
           provisionCmd = "provision.vscode";
           logger.debug("session_start: VS Code provisioned");
           break;
@@ -723,7 +880,7 @@ You are now connected to the sandbox in its current state. Use find, click, type
         case "installer": {
           const installerOpts = provisionOptions as { url: string; filename?: string; appName?: string; launch?: boolean };
           logger.info("session_start: Provisioning installer", { url: installerOpts.url });
-          await sdk.provision.installer(installerOpts);
+          await raceAbort(extra.signal, "session_start", sdk.provision.installer(installerOpts));
           provisionCmd = "provision.installer";
           logger.debug("session_start: Installer provisioned");
           break;
@@ -732,16 +889,20 @@ You are now connected to the sandbox in its current state. Use find, click, type
         case "electron": {
           const electronOpts = provisionOptions as { appPath: string; args?: string[] };
           logger.info("session_start: Provisioning Electron", { appPath: electronOpts.appPath });
-          await sdk.provision.electron(electronOpts);
+          await raceAbort(extra.signal, "session_start", sdk.provision.electron(electronOpts));
           provisionCmd = "provision.electron";
           logger.debug("session_start: Electron app provisioned");
           break;
         }
       }
+      } finally {
+        stopProvisionHeartbeat();
+      }
 
       // Capture initial screenshot after provisioning
       logger.debug("session_start: Capturing initial screenshot");
-      const screenshotBase64 = await sdk.agent.system.captureScreenBase64(1, false, true);
+      progress.report("Capturing screenshot...");
+      const screenshotBase64 = await raceAbort(extra.signal, "session_start", sdk.agent.system.captureScreenBase64(1, false, true)) as string;
       
       let screenshotResourceUri: string | undefined;
       if (screenshotBase64) {
@@ -787,6 +948,18 @@ IMPORTANT - If creating a new test project, use these EXACT dependencies in pack
         generatedCode
       );
     } catch (error) {
+      // On client cancellation, tear down the half-provisioned session so we
+      // don't leak a connected sandbox. The underlying SDK call may still be
+      // running in the background; best-effort cleanup is all we can do.
+      if (error instanceof ToolAbortError) {
+        logger.info("session_start: Cancelled by client, tearing down session");
+        try {
+          await sdk?.disconnect?.();
+        } catch (cleanupErr) {
+          logger.warn("session_start: Cleanup after cancel failed", { error: String(cleanupErr) });
+        }
+        return createToolResult(false, "Session start was cancelled.", { action: "session_start", cancelled: true });
+      }
       logger.error("session_start: Failed", { error: String(error) });
       captureException(error as Error, { tags: { tool: "session_start" }, extra: { params } });
       throw error;
@@ -874,8 +1047,9 @@ registerAppTool(
     }) as any,
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
   },
-  async (params: { description: string; timeout?: number }): Promise<CallToolResult> => {
+  async (params: { description: string; timeout?: number }, extra: ToolExtra): Promise<CallToolResult> => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("find: Starting", { description: params.description, timeout: params.timeout });
 
     const sessionCheck = requireActiveSession();
@@ -886,7 +1060,17 @@ registerAppTool(
 
     try {
       logger.debug("find: Calling SDK find");
-      const element = await sdk.find(params.description, params.timeout ? { timeout: params.timeout } : undefined);
+      const stopHeartbeat = progress.heartbeat(`Looking for "${params.description}"...`);
+      let element: any;
+      try {
+        element = await raceAbort(
+          extra.signal,
+          "find",
+          sdk.find(params.description, params.timeout ? { timeout: params.timeout } : undefined)
+        );
+      } finally {
+        stopHeartbeat();
+      }
       const found = element.found();
       const coords = element.getCoordinates();
 
@@ -975,6 +1159,8 @@ registerAppTool(
         generatedCode
       );
     } catch (error) {
+      const cancelled = cancelledResultOrNull(error, "find");
+      if (cancelled) return cancelled;
       logger.error("find: Failed", { error: String(error), description: params.description });
       captureException(error as Error, { tags: { tool: "find" }, extra: { description: params.description } });
       throw error;
@@ -995,8 +1181,9 @@ registerAppTool(
     }) as any,
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
   },
-  async (params: { description: string; timeout?: number }): Promise<CallToolResult> => {
+  async (params: { description: string; timeout?: number }, extra: ToolExtra): Promise<CallToolResult> => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("findall: Starting", { description: params.description, timeout: params.timeout });
 
     const sessionCheck = requireActiveSession();
@@ -1007,7 +1194,17 @@ registerAppTool(
 
     try {
       logger.debug("findall: Calling SDK findAll");
-      const elements = await sdk.findAll(params.description, params.timeout ? { timeout: params.timeout } : undefined);
+      const stopHeartbeat = progress.heartbeat(`Looking for all "${params.description}"...`);
+      let elements: any;
+      try {
+        elements = await raceAbort(
+          extra.signal,
+          "findall",
+          sdk.findAll(params.description, params.timeout ? { timeout: params.timeout } : undefined)
+        );
+      } finally {
+        stopHeartbeat();
+      }
       const count = elements.length;
 
       // Store element refs for later use
@@ -1104,6 +1301,8 @@ registerAppTool(
         generatedCode
       );
     } catch (error) {
+      const cancelled = cancelledResultOrNull(error, "findall");
+      if (cancelled) return cancelled;
       logger.error("findall: Failed", { error: String(error), description: params.description });
       captureException(error as Error, { tags: { tool: "findall" }, extra: { description: params.description } });
       throw error;
@@ -1373,8 +1572,9 @@ registerAppTool(
     }) as any,
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
   },
-  async (params: { description: string; action: "click" | "double-click" | "right-click" }): Promise<CallToolResult> => {
+  async (params: { description: string; action: "click" | "double-click" | "right-click" }, extra: ToolExtra): Promise<CallToolResult> => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("find_and_click: Starting", { description: params.description, action: params.action });
 
     const sessionCheck = requireActiveSession();
@@ -1385,7 +1585,13 @@ registerAppTool(
 
     try {
       logger.debug("find_and_click: Finding element");
-      const element = await sdk.find(params.description);
+      const stopHeartbeat = progress.heartbeat(`Looking for "${params.description}"...`);
+      let element: any;
+      try {
+        element = await raceAbort(extra.signal, "find_and_click", sdk.find(params.description));
+      } finally {
+        stopHeartbeat();
+      }
       const found = element.found();
 
       if (!found) {
@@ -1511,6 +1717,8 @@ registerAppTool(
         generatedCode
       );
     } catch (error) {
+      const cancelled = cancelledResultOrNull(error, "find_and_click");
+      if (cancelled) return cancelled;
       logger.error("find_and_click: Failed", { error: String(error), description: params.description });
       captureException(error as Error, { tags: { tool: "find_and_click" }, extra: { description: params.description, action: params.action } });
       throw error;
@@ -1730,8 +1938,9 @@ You can optionally provide a reference image URI to compare against a previous s
     }) as any,
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
   },
-  async (params: { task: string; referenceImageUri?: string }): Promise<CallToolResult> => {
+  async (params: { task: string; referenceImageUri?: string }, extra: ToolExtra): Promise<CallToolResult> => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("check: Starting", { task: params.task, hasReferenceImageUri: !!params.referenceImageUri });
 
     const sessionCheck = requireActiveSession();
@@ -1743,6 +1952,7 @@ You can optionally provide a reference image URI to compare against a previous s
     try {
       // Capture current screenshot
       logger.debug("check: Capturing current screenshot");
+      progress.report("Capturing screenshot...");
       const currentScreenshot = await sdk.agent.system.captureScreenBase64(1, false, true);
       
       // Use provided reference image URI, last screenshot as "before" state, or current if no previous screenshot
@@ -1798,12 +2008,18 @@ You can optionally provide a reference image URI to compare against a previous s
         beforeScreenshotPreview: beforeScreenshot?.substring(0, 50),
         currentScreenshotPreview: currentScreenshot?.substring(0, 50)
       });
-      const response = await sdk.agent.sdk.req("check", {
-        tasks: [params.task],
-        images: [beforeScreenshot, currentScreenshot],
-        mousePosition,
-        activeWindow,
-      });
+      const stopHeartbeat = progress.heartbeat(`Checking: "${params.task}"...`);
+      let response: any;
+      try {
+        response = await raceAbort(extra.signal, "check", sdk.agent.sdk.req("check", {
+          tasks: [params.task],
+          images: [beforeScreenshot, currentScreenshot],
+          mousePosition,
+          activeWindow,
+        }));
+      } finally {
+        stopHeartbeat();
+      }
 
       const aiResponse = response.data;
       
@@ -1843,6 +2059,8 @@ You can optionally provide a reference image URI to compare against a previous s
         }
       );
     } catch (error) {
+      const cancelled = cancelledResultOrNull(error, "check");
+      if (cancelled) return cancelled;
       logger.error("check: Failed", { error: String(error), task: params.task });
       captureException(error as Error, { tags: { tool: "check" }, extra: { task: params.task } });
       throw error;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@testdriverai/agent",
-  "version": "7.9.103-canary",
+  "version": "7.9.104-test",
   "description": "Next generation autonomous AI agent for end-to-end testing of web & desktop",
   "main": "sdk.js",
   "types": "sdk.d.ts",

package/sdk.js

@@ -472,10 +472,12 @@ class Element {
     this.sdk.emitter.emit(events.log.log, findingMessage);
 
     try {
-      const screenshot = await this.system.captureScreenBase64();
+      // Returns { imageKey } (fast S3-key path, no local round-trip) or
+      // { image } (base64 fallback). See system.captureScreenImage.
+      const imagePayload = await this.system.captureScreenImage();
       // Only store screenshot in DEBUG mode to prevent memory leaks
-      if (debugMode) {
-        this._screenshot = screenshot;
+      if (debugMode && imagePayload.image) {
+        this._screenshot = imagePayload.image;
       }
 
       // Handle options - can be a number (cacheThreshold) or object with cacheKey/cacheThreshold/cache
@@ -566,7 +568,7 @@ class Element {
       response = await this.sdk.apiClient.req("find", {
         session: this.sdk.getSessionId(),
         element: description,
-        image: screenshot,
+        ...imagePayload,
         threshold: threshold,
         elementSimilarity: elementSimilarity,
         cacheKey: cacheKey,
@@ -2405,7 +2407,8 @@ CAPTCHA_SOLVER_EOF`,
     const { events } = require("./agent/events.js");
 
     try {
-      const screenshot = await this.system.captureScreenBase64();
+      // { imageKey } (fast S3-key path) or { image } (base64 fallback).
+      const imagePayload = await this.system.captureScreenImage();
 
       // Handle options - can be a number (cacheThreshold) or object with cacheKey/cacheThreshold/cache
       let cacheKey = null;
@@ -2474,7 +2477,7 @@ CAPTCHA_SOLVER_EOF`,
         {
           session: this.getSessionId(),
           element: description,
-          image: screenshot,
+          ...imagePayload,
           threshold: threshold,
           cacheKey: cacheKey,
           os: this.os,
@@ -2513,11 +2516,12 @@ CAPTCHA_SOLVER_EOF`,
             elementData,
           );
 
-          // Only store screenshot in DEBUG mode
+          // Only store screenshot in DEBUG mode (and only when we have the
+          // bytes locally — the fast imageKey path doesn't download them)
           const debugMode =
             process.env.VERBOSE || process.env.TD_DEBUG;
-          if (debugMode) {
-            element._screenshot = screenshot;
+          if (debugMode && imagePayload.image) {
+            element._screenshot = imagePayload.image;
           }
 
           return element;
@@ -3001,11 +3005,9 @@ CAPTCHA_SOLVER_EOF`,
     const { events } = require("./agent/events.js");
     this.emitter.emit(events.log.log, "🔍 Running OmniParser screen analysis...");
 
-    const screenshot = await this.system.captureScreenBase64();
-
     const response = await this.apiClient.req("parse", {
       session: this.getSessionId(),
-      image: screenshot,
+      ...(await this.system.captureScreenImage()),
     });
 
     if (response.error) {