fix(computer-use-mcp): harden desktop runtime cleanup (#1746)

---------

Co-authored-by-agent: Antigravity <antigravity@gemini.com>

services/computer-use-mcp/src/browser-dom/cdp-bridge.test.ts

@@ -1,7 +1,36 @@
-import { describe, expect, it } from 'vitest'
+import { EventEmitter } from 'node:events'
+
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import { WebSocket } from 'ws'
 
 import { CdpBridge } from '../browser-dom/cdp-bridge'
 
+afterEach(() => {
+  vi.useRealTimers()
+})
+
+function attachHeartbeatSocket(bridge: CdpBridge, socket: EventEmitter & {
+  readyState: number
+  ping: () => void
+  terminate: () => void
+  close: () => void
+}) {
+  const internals = bridge as unknown as {
+    awaitingHeartbeatPong: boolean
+    consecutiveHeartbeatFailures: number
+    socket: typeof socket
+    status: { connected: boolean }
+    startHeartbeat: () => void
+  }
+  internals.socket = socket
+  internals.status.connected = true
+  socket.on('pong', () => {
+    internals.awaitingHeartbeatPong = false
+    internals.consecutiveHeartbeatFailures = 0
+  })
+  internals.startHeartbeat()
+}
+
 describe('cdpBridge', () => {
   it('creates with correct initial status', () => {
     const bridge = new CdpBridge({
@@ -117,4 +146,114 @@ describe('cdpBridge', () => {
     await bridge.close()
     expect(bridge.getStatus().connected).toBe(false)
   })
+
+  it('keeps the CDP bridge alive when heartbeat pongs arrive', () => {
+    vi.useFakeTimers()
+    const bridge = new CdpBridge({
+      cdpUrl: 'http://localhost:9222',
+      requestTimeoutMs: 10_000,
+      heartbeatIntervalMs: 10,
+      heartbeatFailureLimit: 3,
+    })
+    const socket = Object.assign(new EventEmitter(), {
+      readyState: WebSocket.OPEN,
+      ping: vi.fn(() => socket.emit('pong')),
+      terminate: vi.fn(),
+      close: vi.fn(),
+    })
+
+    attachHeartbeatSocket(bridge, socket)
+
+    vi.advanceTimersByTime(40)
+
+    expect(socket.ping).toHaveBeenCalled()
+    expect(socket.terminate).not.toHaveBeenCalled()
+    expect(bridge.getStatus().connected).toBe(true)
+  })
+
+  it('tears down the CDP bridge after consecutive missed heartbeat pongs', () => {
+    vi.useFakeTimers()
+    const bridge = new CdpBridge({
+      cdpUrl: 'http://localhost:9222',
+      requestTimeoutMs: 10_000,
+      heartbeatIntervalMs: 10,
+      heartbeatFailureLimit: 3,
+    })
+    const socket = Object.assign(new EventEmitter(), {
+      readyState: WebSocket.OPEN,
+      ping: vi.fn(),
+      terminate: vi.fn(),
+      close: vi.fn(),
+    })
+
+    attachHeartbeatSocket(bridge, socket)
+
+    vi.advanceTimersByTime(30)
+
+    expect(socket.ping).toHaveBeenCalledTimes(3)
+    expect(socket.terminate).not.toHaveBeenCalled()
+    expect(bridge.getStatus().connected).toBe(true)
+
+    vi.advanceTimersByTime(10)
+
+    expect(socket.terminate).toHaveBeenCalledTimes(1)
+    expect(bridge.getStatus().connected).toBe(false)
+    expect(bridge.getStatus().lastError).toBe('CDP heartbeat failed after 3 consecutive missed pongs')
+  })
+
+  it('does not tear down the CDP bridge before the first heartbeat ping can miss', () => {
+    vi.useFakeTimers()
+    const bridge = new CdpBridge({
+      cdpUrl: 'http://localhost:9222',
+      requestTimeoutMs: 10_000,
+      heartbeatIntervalMs: 10,
+      heartbeatFailureLimit: 1,
+    })
+    const socket = Object.assign(new EventEmitter(), {
+      readyState: WebSocket.OPEN,
+      ping: vi.fn(),
+      terminate: vi.fn(),
+      close: vi.fn(),
+    })
+
+    attachHeartbeatSocket(bridge, socket)
+
+    vi.advanceTimersByTime(10)
+
+    expect(socket.ping).toHaveBeenCalledTimes(1)
+    expect(socket.terminate).not.toHaveBeenCalled()
+    expect(bridge.getStatus().connected).toBe(true)
+
+    vi.advanceTimersByTime(10)
+
+    expect(socket.ping).toHaveBeenCalledTimes(1)
+    expect(socket.terminate).toHaveBeenCalledTimes(1)
+    expect(bridge.getStatus().connected).toBe(false)
+  })
+
+  it('preserves heartbeat ping errors in the bridge status', () => {
+    vi.useFakeTimers()
+    const bridge = new CdpBridge({
+      cdpUrl: 'http://localhost:9222',
+      requestTimeoutMs: 10_000,
+      heartbeatIntervalMs: 10,
+      heartbeatFailureLimit: 3,
+    })
+    const socket = Object.assign(new EventEmitter(), {
+      readyState: WebSocket.OPEN,
+      ping: vi.fn(() => {
+        throw new Error('CDP ping write failed')
+      }),
+      terminate: vi.fn(),
+      close: vi.fn(),
+    })
+
+    attachHeartbeatSocket(bridge, socket)
+
+    vi.advanceTimersByTime(10)
+
+    expect(socket.terminate).toHaveBeenCalledTimes(1)
+    expect(bridge.getStatus().connected).toBe(false)
+    expect(bridge.getStatus().lastError).toBe('CDP ping write failed')
+  })
 })

services/computer-use-mcp/src/browser-dom/cdp-bridge.ts

@@ -19,6 +19,18 @@ export interface CdpBridgeConfig {
   cdpUrl: string
   /** Request timeout in milliseconds */
   requestTimeoutMs: number
+  /**
+   * WebSocket heartbeat interval in milliseconds.
+   *
+   * @default 5000
+   */
+  heartbeatIntervalMs?: number
+  /**
+   * Consecutive missed heartbeat pongs before tearing down the bridge.
+   *
+   * @default 3
+   */
+  heartbeatFailureLimit?: number
 }
 
 export interface CdpBridgeStatus {
@@ -75,6 +87,9 @@ interface CdpTargetInfo {
 
 export class CdpBridge {
   private socket?: WebSocket
+  private heartbeatTimer?: NodeJS.Timeout
+  private awaitingHeartbeatPong = false
+  private consecutiveHeartbeatFailures = 0
   private nextId = 1
   private pending = new Map<number, PendingCdpRequest>()
   private status: CdpBridgeStatus
@@ -117,16 +132,19 @@ export class CdpBridge {
    */
   async connectToTarget(target: CdpTargetInfo): Promise<void> {
     if (this.socket) {
-      this.socket.close()
-      this.socket = undefined
+      await this.close()
     }
 
     const wsUrl = target.webSocketDebuggerUrl!
 
     await new Promise<void>((resolve, reject) => {
       const socket = new WebSocket(wsUrl)
+      let connectionSettled = false
 
       socket.on('open', () => {
+        if (connectionSettled)
+          return
+        connectionSettled = true
         this.socket = socket
         this.status.connected = true
         this.status.pageTitle = target.title
@@ -136,18 +154,36 @@ export class CdpBridge {
       })
 
       socket.on('message', (data) => {
+        if (this.socket !== socket)
+          return
         this.handleMessage(data)
       })
 
+      socket.on('pong', () => {
+        if (this.socket !== socket)
+          return
+        this.consecutiveHeartbeatFailures = 0
+        this.awaitingHeartbeatPong = false
+      })
+
       socket.on('close', () => {
-        this.socket = undefined
-        this.status.connected = false
+        if (this.socket !== socket)
+          return
+        this.handleSocketClosed('CDP WebSocket closed')
       })
 
       socket.on('error', (error) => {
-        this.status.lastError = error instanceof Error ? error.message : String(error)
+        if (this.socket && this.socket !== socket)
+          return
+
+        const message = error instanceof Error ? error.message : String(error)
+        if (!this.socket && connectionSettled)
+          return
+
+        this.status.lastError = message
         if (!this.socket) {
-          reject(new Error(`CDP WebSocket connection failed: ${this.status.lastError}`))
+          connectionSettled = true
+          reject(new Error(`CDP WebSocket connection failed: ${message}`))
         }
       })
     })
@@ -156,17 +192,17 @@ export class CdpBridge {
     await this.send('Accessibility.enable', {})
     await this.send('DOM.enable', {})
     await this.send('Runtime.enable', {})
+    this.startHeartbeat()
   }
 
   /**
    * Close the CDP connection.
    */
   async close(): Promise<void> {
-    for (const [id, pending] of this.pending.entries()) {
-      clearTimeout(pending.timeoutId)
-      pending.reject(new Error(`CDP bridge closed before completing request ${id}`))
-    }
-    this.pending.clear()
+    this.clearHeartbeat()
+    this.rejectPendingRequests('CDP bridge closed')
+    this.awaitingHeartbeatPong = false
+    this.consecutiveHeartbeatFailures = 0
 
     if (this.socket) {
       this.socket.close()
@@ -387,4 +423,84 @@ export class CdpBridge {
       pending.resolve(data.result)
     }
   }
+
+  private startHeartbeat() {
+    this.clearHeartbeat()
+    this.awaitingHeartbeatPong = false
+    this.consecutiveHeartbeatFailures = 0
+
+    const intervalMs = this.config.heartbeatIntervalMs ?? 5_000
+    const failureLimit = this.config.heartbeatFailureLimit ?? 3
+
+    if (intervalMs <= 0 || failureLimit <= 0) {
+      return
+    }
+
+    this.heartbeatTimer = setInterval(() => {
+      const socket = this.socket
+      if (!socket || socket.readyState !== WebSocket.OPEN) {
+        this.handleSocketClosed('CDP heartbeat found closed WebSocket')
+        return
+      }
+
+      if (this.awaitingHeartbeatPong) {
+        this.consecutiveHeartbeatFailures += 1
+        if (this.consecutiveHeartbeatFailures >= failureLimit) {
+          this.teardownAfterHeartbeatFailure()
+          return
+        }
+      }
+
+      try {
+        socket.ping()
+        this.awaitingHeartbeatPong = true
+      }
+      catch (error) {
+        const message = error instanceof Error ? error.message : String(error)
+        this.teardownAfterHeartbeatFailure(message)
+      }
+    }, intervalMs)
+    this.heartbeatTimer.unref?.()
+  }
+
+  private clearHeartbeat() {
+    if (!this.heartbeatTimer) {
+      return
+    }
+
+    clearInterval(this.heartbeatTimer)
+    this.heartbeatTimer = undefined
+  }
+
+  private teardownAfterHeartbeatFailure(reason?: string) {
+    const failureLimit = this.config.heartbeatFailureLimit ?? 3
+    this.status.lastError = reason ?? `CDP heartbeat failed after ${failureLimit} consecutive missed pongs`
+    this.clearHeartbeat()
+    this.rejectPendingRequests(this.status.lastError)
+
+    const socket = this.socket
+    this.socket = undefined
+    this.status.connected = false
+    this.awaitingHeartbeatPong = false
+    this.consecutiveHeartbeatFailures = 0
+    socket?.terminate()
+  }
+
+  private handleSocketClosed(reason: string) {
+    this.status.lastError = reason
+    this.clearHeartbeat()
+    this.rejectPendingRequests(reason)
+    this.socket = undefined
+    this.status.connected = false
+    this.awaitingHeartbeatPong = false
+    this.consecutiveHeartbeatFailures = 0
+  }
+
+  private rejectPendingRequests(reason: string) {
+    for (const [id, pending] of this.pending.entries()) {
+      clearTimeout(pending.timeoutId)
+      pending.reject(new Error(`${reason} before completing request ${id}`))
+    }
+    this.pending.clear()
+  }
 }

services/computer-use-mcp/src/executors/macos-local.test.ts

@@ -0,0 +1,23 @@
+import { describe, expect, it } from 'vitest'
+
+import { moveAndClickScript } from './macos-local'
+
+describe('moveAndClickScript', () => {
+  it('saves and restores the real macOS cursor around CGEvent clicks', () => {
+    const source = moveAndClickScript()
+
+    const saveCursorIndex = source.indexOf('let originalCursorLocation = CGEvent(source: nil)?.location')
+    const restoreCursorIndex = source.indexOf('CGWarpMouseCursorPosition(savedCursorLocation)')
+    const moveTraceIndex = source.indexOf('for point in trace')
+    const clickIndex = source.indexOf('down.post(tap: .cghidEventTap)')
+
+    expect(saveCursorIndex).toBeGreaterThanOrEqual(0)
+    expect(restoreCursorIndex).toBeGreaterThanOrEqual(0)
+    expect(moveTraceIndex).toBeGreaterThanOrEqual(0)
+    expect(clickIndex).toBeGreaterThanOrEqual(0)
+    expect(saveCursorIndex).toBeLessThan(moveTraceIndex)
+    expect(restoreCursorIndex).toBeLessThan(moveTraceIndex)
+    expect(restoreCursorIndex).toBeLessThan(clickIndex)
+    expect(source).toContain('defer {')
+  })
+})

services/computer-use-mcp/src/executors/macos-local.ts

@@ -210,7 +210,21 @@ print(String(data: data, encoding: .utf8)!)
 `
 }
 
-function moveAndClickScript() {
+/**
+ * Creates Swift source that posts a local macOS click while restoring the user's cursor.
+ *
+ * Use when:
+ * - The macOS executor needs to perform a pointer-based click through Quartz.
+ * - Tests need to inspect the generated Swift source without moving the real cursor.
+ *
+ * Expects:
+ * - `COMPUTER_USE_SWIFT_STDIN` contains `pointerTrace`, `button`, and `clickCount`.
+ * - The caller already verified that the host platform is macOS.
+ *
+ * Returns:
+ * - Swift source that saves the real cursor location before CGEvent movement and restores it with `CGWarpMouseCursorPosition`.
+ */
+export function moveAndClickScript() {
   return String.raw`
 import CoreGraphics
 import Foundation
@@ -248,6 +262,16 @@ let buttonRaw = input["button"] as? Int ?? 0
 let clickCount = input["clickCount"] as? Int ?? 1
 let button = mouseButton(buttonRaw)
 
+// CGEvent mouse posts move the real macOS cursor. Save and restore it so the
+// overlay ghost pointer can visualize agent intent without leaving the user's
+// actual cursor at the agent click target.
+let originalCursorLocation = CGEvent(source: nil)?.location
+defer {
+  if let savedCursorLocation = originalCursorLocation {
+    CGWarpMouseCursorPosition(savedCursorLocation)
+  }
+}
+
 for point in trace {
   let x = point["x"] as? Double ?? 0
   let y = point["y"] as? Double ?? 0