@testdriverai/agent 7.9.103-canary → 7.9.104-canary

package/mcp-server/dist/server.mjs

@@ -266,6 +266,102 @@ function requireActiveSession() {
     sessionManager.refreshSession(session.sessionId);
     return { valid: true };
 }
+const DEFAULT_HEARTBEAT_MS = 3000;
+function makeProgressReporter(extra) {
+    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) => {
+        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) {
+        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, tool) {
+    if (!signal) {
+        // Never-resolving promise with a no-op cleanup — Promise.race ignores it.
+        return { promise: new Promise(() => { }), cleanup: () => { } };
+    }
+    let onAbort = () => { };
+    const promise = new Promise((_, 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(signal, tool, work) {
+    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, tool) {
+    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
@@ -411,8 +507,9 @@ Debug mode (connect to existing sandbox):
 - Use this to interactively debug failed tests without re-running from scratch`,
     inputSchema: SessionStartInputSchema,
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
-}, async (params) => {
+}, async (params, extra) => {
     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
     const { os: resolvedOs, warning: osWarning } = resolveOs(params.os);
@@ -495,10 +592,16 @@ 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();
             logger.info("session_start: Connected to existing sandbox", { instanceId: instance?.instanceId });
@@ -507,7 +610,8 @@ Debug mode (connect to existing sandbox):
             setSessionContext(newSession.sessionId, instance?.instanceId);
             // 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));
             let screenshotResourceUri;
             if (screenshotBase64) {
                 screenshotResourceUri = storeImage(screenshotBase64, "screenshot");
@@ -536,11 +640,19 @@ 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();
         logger.info("session_start: Connected to sandbox", { instanceId: instance?.instanceId });
@@ -550,52 +662,61 @@ You are now connected to the sandbox in its current state. Use find, click, type
         // Get provision-specific options
         const provisionOptions = getProvisionOptions(params);
         let provisionCmd = "";
-        // Provision based on type
-        switch (params.type) {
-            case "chrome": {
-                const chromeOpts = provisionOptions;
-                logger.info("session_start: Provisioning Chrome", { url: chromeOpts.url });
-                await sdk.provision.chrome(chromeOpts);
-                provisionCmd = "provision.chrome";
-                logger.debug("session_start: Chrome provisioned");
-                break;
-            }
-            case "chromeExtension": {
-                const extOpts = provisionOptions;
-                logger.info("session_start: Provisioning Chrome Extension", { extensionPath: extOpts.extensionPath, extensionId: extOpts.extensionId });
-                await sdk.provision.chromeExtension(extOpts);
-                provisionCmd = "provision.chromeExtension";
-                logger.debug("session_start: Chrome Extension provisioned");
-                break;
-            }
-            case "vscode": {
-                const vscodeOpts = provisionOptions;
-                logger.info("session_start: Provisioning VS Code", { workspace: vscodeOpts.workspace });
-                await sdk.provision.vscode(vscodeOpts);
-                provisionCmd = "provision.vscode";
-                logger.debug("session_start: VS Code provisioned");
-                break;
-            }
-            case "installer": {
-                const installerOpts = provisionOptions;
-                logger.info("session_start: Provisioning installer", { url: installerOpts.url });
-                await sdk.provision.installer(installerOpts);
-                provisionCmd = "provision.installer";
-                logger.debug("session_start: Installer provisioned");
-                break;
-            }
-            case "electron": {
-                const electronOpts = provisionOptions;
-                logger.info("session_start: Provisioning Electron", { appPath: electronOpts.appPath });
-                await sdk.provision.electron(electronOpts);
-                provisionCmd = "provision.electron";
-                logger.debug("session_start: Electron app provisioned");
-                break;
+        // 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;
+                    logger.info("session_start: Provisioning Chrome", { url: chromeOpts.url });
+                    await raceAbort(extra.signal, "session_start", sdk.provision.chrome(chromeOpts));
+                    provisionCmd = "provision.chrome";
+                    logger.debug("session_start: Chrome provisioned");
+                    break;
+                }
+                case "chromeExtension": {
+                    const extOpts = provisionOptions;
+                    logger.info("session_start: Provisioning Chrome Extension", { extensionPath: extOpts.extensionPath, extensionId: extOpts.extensionId });
+                    await raceAbort(extra.signal, "session_start", sdk.provision.chromeExtension(extOpts));
+                    provisionCmd = "provision.chromeExtension";
+                    logger.debug("session_start: Chrome Extension provisioned");
+                    break;
+                }
+                case "vscode": {
+                    const vscodeOpts = provisionOptions;
+                    logger.info("session_start: Provisioning VS Code", { workspace: vscodeOpts.workspace });
+                    await raceAbort(extra.signal, "session_start", sdk.provision.vscode(vscodeOpts));
+                    provisionCmd = "provision.vscode";
+                    logger.debug("session_start: VS Code provisioned");
+                    break;
+                }
+                case "installer": {
+                    const installerOpts = provisionOptions;
+                    logger.info("session_start: Provisioning installer", { url: installerOpts.url });
+                    await raceAbort(extra.signal, "session_start", sdk.provision.installer(installerOpts));
+                    provisionCmd = "provision.installer";
+                    logger.debug("session_start: Installer provisioned");
+                    break;
+                }
+                case "electron": {
+                    const electronOpts = provisionOptions;
+                    logger.info("session_start: Provisioning Electron", { appPath: electronOpts.appPath });
+                    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));
         let screenshotResourceUri;
         if (screenshotBase64) {
             screenshotResourceUri = storeImage(screenshotBase64, "screenshot");
@@ -632,6 +753,19 @@ 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, { tags: { tool: "session_start" }, extra: { params } });
         throw error;
@@ -693,8 +827,9 @@ registerAppTool(server, "find", {
         timeout: z.number().optional().describe("Timeout in ms for polling"),
     }),
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
-}, async (params) => {
+}, async (params, extra) => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("find: Starting", { description: params.description, timeout: params.timeout });
     const sessionCheck = requireActiveSession();
     if (!sessionCheck.valid) {
@@ -703,7 +838,14 @@ registerAppTool(server, "find", {
     }
     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;
+        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();
         // Store element ref for later use (stores actual Element instance)
@@ -784,6 +926,9 @@ registerAppTool(server, "find", {
         }, generatedCode);
     }
     catch (error) {
+        const cancelled = cancelledResultOrNull(error, "find");
+        if (cancelled)
+            return cancelled;
         logger.error("find: Failed", { error: String(error), description: params.description });
         captureException(error, { tags: { tool: "find" }, extra: { description: params.description } });
         throw error;
@@ -798,8 +943,9 @@ registerAppTool(server, "findall", {
         timeout: z.number().optional().describe("Timeout in ms for polling"),
     }),
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
-}, async (params) => {
+}, async (params, extra) => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("findall: Starting", { description: params.description, timeout: params.timeout });
     const sessionCheck = requireActiveSession();
     if (!sessionCheck.valid) {
@@ -808,7 +954,14 @@ registerAppTool(server, "findall", {
     }
     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;
+        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
         const refs = [];
@@ -893,6 +1046,9 @@ registerAppTool(server, "findall", {
         }, generatedCode);
     }
     catch (error) {
+        const cancelled = cancelledResultOrNull(error, "findall");
+        if (cancelled)
+            return cancelled;
         logger.error("findall: Failed", { error: String(error), description: params.description });
         captureException(error, { tags: { tool: "findall" }, extra: { description: params.description } });
         throw error;
@@ -1089,8 +1245,9 @@ registerAppTool(server, "find_and_click", {
         action: z.enum(["click", "double-click", "right-click"]).default("click"),
     }),
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
-}, async (params) => {
+}, async (params, extra) => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("find_and_click: Starting", { description: params.description, action: params.action });
     const sessionCheck = requireActiveSession();
     if (!sessionCheck.valid) {
@@ -1099,7 +1256,14 @@ registerAppTool(server, "find_and_click", {
     }
     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;
+        try {
+            element = await raceAbort(extra.signal, "find_and_click", sdk.find(params.description));
+        }
+        finally {
+            stopHeartbeat();
+        }
         const found = element.found();
         if (!found) {
             logger.warn("find_and_click: Element not found", { description: params.description });
@@ -1207,6 +1371,9 @@ registerAppTool(server, "find_and_click", {
         }, 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, { tags: { tool: "find_and_click" }, extra: { description: params.description, action: params.action } });
         throw error;
@@ -1364,8 +1531,9 @@ You can optionally provide a reference image URI to compare against a previous s
         referenceImageUri: z.string().optional().describe("Optional screenshot resource URI (e.g., 'screenshot://testdriver/screenshot/screenshot-1') to compare against instead of the automatically captured 'before' screenshot. Use a screenshotResourceUri from a previous action."),
     }),
     _meta: { ui: { resourceUri: RESOURCE_URI, expanded: true } },
-}, async (params) => {
+}, async (params, extra) => {
     const startTime = Date.now();
+    const progress = makeProgressReporter(extra);
     logger.info("check: Starting", { task: params.task, hasReferenceImageUri: !!params.referenceImageUri });
     const sessionCheck = requireActiveSession();
     if (!sessionCheck.valid) {
@@ -1375,6 +1543,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
         let beforeScreenshot;
@@ -1425,12 +1594,19 @@ 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;
+        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;
         // Store screenshot for resource serving
         let screenshotResourceUri;
@@ -1460,6 +1636,9 @@ 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, { tags: { tool: "check" }, extra: { task: params.task } });
         throw error;