fix(computer-use): preserve window identity and snapshots after overlay failure

native/cu-helper/Sources/cu-helper/AXAction.swift

@@ -1255,7 +1255,7 @@ public enum AXAction {
         var value: CFTypeRef?
         if AXUIElementCopyAttributeValue(app, kAXWindowsAttribute as CFString, &value) == .success,
            let windows = value as? [AXUIElement] {
-            for window in windows where windowNumber(of: window) == windowID {
+            for window in windows where WindowGeometry.axWindowID(of: window) == windowID {
                 if actionNames(window).contains(axRaise) {
                     _ = AXUIElementPerformAction(window, axRaise as CFString)
                 }
@@ -1267,18 +1267,6 @@ public enum AXAction {
         return !WindowGeometry.isFullyCovered(windowID: windowID)
     }
 
-    /// The CGWindowID behind an AX window element, via the private-but-stable
-    /// `_AXUIElementGetWindow`. Returns 0 when it cannot be read, which never
-    /// matches a real window and so simply skips that element.
-    private static func windowNumber(of element: AXUIElement) -> CGWindowID {
-        typealias GetWindow = @convention(c) (AXUIElement, UnsafeMutablePointer<CGWindowID>) -> AXError
-        guard let handle = dlopen(nil, RTLD_LAZY),
-              let symbol = dlsym(handle, "_AXUIElementGetWindow") else { return 0 }
-        let getWindow = unsafeBitCast(symbol, to: GetWindow.self)
-        var id: CGWindowID = 0
-        return getWindow(element, &id) == .success ? id : 0
-    }
-
     private static func activateWindow(_ element: AXUIElement) -> Bool {
         var activated = false
         if actionNames(element).contains(axRaise),

native/cu-helper/Sources/cu-helper/AXTree.swift

@@ -690,10 +690,10 @@ public enum AXTree {
         }
     }
 
-    /// Map one AX window to a Window Server id using public information only.
-    /// PID is pre-filtered by `cgWindowCandidates`; the AX/CG frames must match,
-    /// and normalized titles must either agree on both sides or be absent on both
-    /// sides. Every accepted branch requires exactly one candidate.
+    /// Associate the AX root with its actual WindowServer ID, validated against
+    /// this PID's current candidates. Only when that API is unavailable do we
+    /// require the older frame/title evidence. Chrome's AX title includes app
+    /// and profile names that its CG title omits; title equality is not identity.
     private static func mappedWindowID(
         _ window: AXUIElement,
         candidates: [CGWindowCandidate]
@@ -711,7 +711,8 @@ public enum AXTree {
         }
         return SnapshotWindowIdentityEvidence.mappedWindowID(
             axTitle: axTitle,
-            candidates: evidence
+            candidates: evidence,
+            nativeWindowID: WindowGeometry.axWindowID(of: window)
         )
     }
 

native/cu-helper/Sources/cu-helper/ElementFingerprint.swift

@@ -138,9 +138,9 @@ struct SnapshotPathStep: Sendable, Hashable {
     }
 }
 
-/// Pure fail-closed rules shared by snapshot-time AX↔CG window mapping. Frames
-/// are checked by AXTree; this type owns the evidence rules that are easy to
-/// accidentally weaken: titles must be symmetric, and root IDs must be 1:1.
+/// Pure fail-closed rules shared by snapshot-time AX↔CG window mapping. Direct
+/// IDs must belong to the target process's candidates; fallback title evidence
+/// must be symmetric, and root IDs must be 1:1.
 enum SnapshotWindowIdentityEvidence {
     struct Candidate: Sendable, Equatable {
         let id: UInt32
@@ -152,15 +152,23 @@ enum SnapshotWindowIdentityEvidence {
         normalizedTitle(axTitle) == normalizedTitle(cgTitle)
     }
 
-    /// Select a public WindowServer identity without trusting AX window order.
-    /// Exact frame + bilateral title evidence remains authoritative. Stage
+    /// Select a WindowServer identity without trusting AX window order. Prefer
+    /// the native ID, with exact frame + bilateral title as a fallback. Stage
     /// Manager can expose only thumbnail bounds for background windows; when no
     /// candidate matches the AX frame at all, a unique bilateral non-empty title
-    /// match is the bounded fallback. Any ambiguity still fails closed.
+    /// match is the last fallback. Any ambiguity still fails closed.
     static func mappedWindowID(
         axTitle: String?,
-        candidates: [Candidate]
+        candidates: [Candidate],
+        nativeWindowID: UInt32? = nil
     ) -> UInt32? {
+        if let nativeWindowID, nativeWindowID != 0 {
+            // A direct ID is stronger than a mutable title or frame. Never
+            // fall back to another window when that ID is missing from the
+            // target PID's live candidates (closed, wrong process, or layer).
+            let matches = candidates.filter { $0.id == nativeWindowID }
+            return matches.count == 1 ? nativeWindowID : nil
+        }
         let frameMatches = candidates.filter(\.frameMatches)
         let framedEvidence = frameMatches.filter {
             titlesMatch(axTitle: axTitle, cgTitle: $0.title)

native/cu-helper/Sources/cu-helper/WindowGeometry.swift

@@ -1,4 +1,6 @@
+import ApplicationServices
 import CoreGraphics
+import Darwin
 import Foundation
 
 /// Which on-screen window owns a global point, and where that window sits.
@@ -11,6 +13,28 @@ import Foundation
 /// Window *names* would require Screen Recording; these three do not, so this
 /// works before any capture grant exists.
 enum WindowGeometry {
+    private typealias GetAXWindow = @convention(c) (
+        AXUIElement, UnsafeMutablePointer<CGWindowID>
+    ) -> AXError
+
+    private static let getAXWindow: GetAXWindow? = {
+        guard let handle = dlopen(nil, RTLD_LAZY),
+              let symbol = dlsym(handle, "_AXUIElementGetWindow") else { return nil }
+        return unsafeBitCast(symbol, to: GetAXWindow.self)
+    }()
+
+    /// Read the window's identity directly. AX and WindowServer titles are
+    /// presentation strings and need not agree (Chrome decorates its AX title).
+    /// Callers must still validate this ID against the target process's live
+    /// windows. An unavailable SPI falls back to the existing evidence rules.
+    static func axWindowID(of element: AXUIElement) -> CGWindowID? {
+        var id: CGWindowID = 0
+        guard let getAXWindow,
+              getAXWindow(element, &id) == .success,
+              id != kCGNullWindowID else { return nil }
+        return id
+    }
+
     struct Window: Equatable, Sendable {
         let id: CGWindowID
         let bounds: CGRect

native/cu-helper/Tests/CuHelperTests/AXTreePublicationIntegrationTests.swift

@@ -12,16 +12,34 @@ final class AXTreePublicationIntegrationTests: XCTestCase {
     private static let fixtureReadyPath = "CC_HAHA_AX_PUBLICATION_READY"
     private static let fixtureStopPath = "CC_HAHA_AX_PUBLICATION_STOP"
     private static let fixtureTitle = "CC_HAHA_AX_PUBLICATION_TITLE"
+    private static let fixtureMismatchedTitle = "CC_HAHA_AX_PUBLICATION_MISMATCHED_TITLE"
 
     func testPublishedControlBelowDuplicateAncestorClicksImmediatelyAndRejectsOldGeneration() async throws {
+        try await verifyPublishedControl(mismatchedWindowTitle: false)
+    }
+
+    func testChromeStyleAXTitleCanDifferFromWindowServerTitleWithoutDisablingActions() async throws {
+        try await verifyPublishedControl(mismatchedWindowTitle: true)
+    }
+
+    private func verifyPublishedControl(mismatchedWindowTitle: Bool) async throws {
         if ProcessInfo.processInfo.environment[Self.fixtureFlag] == "1" {
-            try await runFixtureProcess()
-            return
+            try await runFixtureProcess(
+                mismatchedWindowTitle: ProcessInfo.processInfo.environment[Self.fixtureMismatchedTitle] == "1"
+            )
+            // This process is a disposable UI fixture, not another suite run.
+            exit(0)
         }
         try XCTSkipUnless(
             AXIsProcessTrusted(),
             "Live AX publication requires Accessibility permission for the test runner"
         )
+        if mismatchedWindowTitle {
+            try XCTSkipUnless(
+                Capture.hasScreenRecordingPermission(),
+                "Coordinate publication requires Screen Recording permission for the test runner"
+            )
+        }
 
         let root = FileManager.default.temporaryDirectory
             .appendingPathComponent("cc-haha-ax-publication-\(UUID().uuidString)")
@@ -43,6 +61,7 @@ final class AXTreePublicationIntegrationTests: XCTestCase {
             Self.fixtureReadyPath: ready.path,
             Self.fixtureStopPath: stop.path,
             Self.fixtureTitle: title,
+            Self.fixtureMismatchedTitle: mismatchedWindowTitle ? "1" : "0",
         ]) { _, fixture in fixture }
         configuration.activates = false
         configuration.createsNewApplicationInstance = true
@@ -60,6 +79,8 @@ final class AXTreePublicationIntegrationTests: XCTestCase {
         defer { AXTree.invalidate(pid: pid) }
 
         let state = try await AXTree.appState(pid: pid, disableDiff: true)
+        let windowID = try XCTUnwrap(AXTree.snapshotEvidence(pid: pid)?.keyWindowID)
+        XCTAssertEqual(AXTree.currentKeyWindowID(pid: pid), windowID)
         let (handle, line) = try publishedHandle(label: "Bold", state: state)
         XCTAssertEqual(
             AXTree.record(pid: pid, index: handle.index)?.role,
@@ -89,6 +110,34 @@ final class AXTreePublicationIntegrationTests: XCTestCase {
         let (_, clickedLine) = try publishedHandle(label: "Bold", state: clickedState)
         XCTAssertTrue(clickedLine.contains("Value: 1"), clickedLine)
 
+        if mismatchedWindowTitle {
+            // Exercise the actual state → screenshot → coordinate action path,
+            // using only this disposable fixture window. Previously the image
+            // was returned while its coordinate transform was never recorded.
+            let captured = try await router.handle(cmd: "get_app_state", payload: .object([
+                "pid": .int(Int(pid)), "disableDiff": .bool(true),
+            ]))
+            let shot = try XCTUnwrap(captured["screenshot"])
+            XCTAssertEqual(shot["windowID"]?.asInt, Int(windowID))
+            let frame = try XCTUnwrap(AXTree.record(pid: pid, index: handle.index)?.frameGlobal)
+            let x = (frame.x + frame.w / 2 - (try XCTUnwrap(shot["originX"]?.asDouble)))
+                * Double(try XCTUnwrap(shot["width"]?.asInt))
+                / (try XCTUnwrap(shot["pointWidth"]?.asDouble))
+            let y = (frame.y + frame.h / 2 - (try XCTUnwrap(shot["originY"]?.asDouble)))
+                * Double(try XCTUnwrap(shot["height"]?.asInt))
+                / (try XCTUnwrap(shot["pointHeight"]?.asDouble))
+            _ = try await router.handle(cmd: "click", payload: .object([
+                "pid": .int(Int(pid)), "x": .double(x), "y": .double(y),
+            ]))
+            let coordinateState = try await AXTree.appState(pid: pid, disableDiff: true)
+            let (coordinateHandle, coordinateLine) = try publishedHandle(label: "Bold", state: coordinateState)
+            XCTAssertTrue(coordinateLine.contains("Value: 0"), coordinateLine)
+            // Restore the checked state for the stale-generation assertion.
+            _ = try await router.handle(cmd: "click", payload: .object([
+                "pid": .int(Int(pid)), "index": .string(coordinateHandle.rawValue),
+            ]))
+        }
+
         AXTree.invalidate(pid: pid)
         let nextState = try await AXTree.appState(pid: pid, disableDiff: true)
         let (nextHandle, _) = try publishedHandle(label: "Bold", state: nextState)
@@ -113,7 +162,7 @@ final class AXTreePublicationIntegrationTests: XCTestCase {
         try await waitUntil { process.isTerminated }
     }
 
-    private func runFixtureProcess() async throws {
+    private func runFixtureProcess(mismatchedWindowTitle: Bool) async throws {
         let environment = ProcessInfo.processInfo.environment
         let readyPath = try XCTUnwrap(environment[Self.fixtureReadyPath])
         let stopPath = try XCTUnwrap(environment[Self.fixtureStopPath])
@@ -129,6 +178,12 @@ final class AXTreePublicationIntegrationTests: XCTestCase {
             defer: false
         )
         window.title = title
+        if mismatchedWindowTitle {
+            // Chrome exposes a decorated AX title while WindowServer uses the
+            // page title (and may add an audio indicator). Both refer to the
+            // same window. Reproduce that mismatch without a real browser.
+            window.setAccessibilityTitle("\(title) - Google Chrome - Fixture")
+        }
         // TextEdit exposes formatting and alignment segments as two sibling
         // AXGroups whose own fingerprints are identical. Their description-only
         // children are the first semantic evidence that distinguishes the paths.

native/cu-helper/Tests/CuHelperTests/ElementFingerprintTests.swift

@@ -246,6 +246,57 @@ final class ElementFingerprintTests: XCTestCase {
         )
     }
 
+    func testNativeWindowIDMatchesChromeDespiteDecoratedAXAndAudioCGTitles() {
+        let candidates = [
+            SnapshotWindowIdentityEvidence.Candidate(id: 17673, frameMatches: false, title: "Window"),
+            SnapshotWindowIdentityEvidence.Candidate(id: 16290, frameMatches: true, title: "“Townscaper”🔊"),
+        ]
+        XCTAssertNil(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Townscaper - Google Chrome - Mi", candidates: candidates
+        ))
+        XCTAssertEqual(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Townscaper - Google Chrome - Mi", candidates: candidates,
+            nativeWindowID: 16290
+        ), 16290)
+    }
+
+    func testNativeWindowIDSurvivesMissingTitleAndMovedFrame() {
+        XCTAssertEqual(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Document", candidates: [
+                .init(id: 11, frameMatches: false, title: nil),
+            ], nativeWindowID: 11
+        ), 11)
+    }
+
+    func testNativeWindowIDOutsideTargetCandidatesCannotFallBackToLookalike() {
+        XCTAssertNil(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Document", candidates: [
+                .init(id: 11, frameMatches: true, title: "Document"),
+            ], nativeWindowID: 12
+        ))
+    }
+
+    func testAmbiguousNativeWindowIDIsRejected() {
+        XCTAssertNil(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Document", candidates: [
+                .init(id: 11, frameMatches: true, title: "Document"),
+                .init(id: 11, frameMatches: false, title: "Other"),
+            ], nativeWindowID: 11
+        ))
+    }
+
+    func testUnavailableNativeWindowIDPreservesStrictTitleFallback() {
+        let candidates = [SnapshotWindowIdentityEvidence.Candidate(
+            id: 11, frameMatches: true, title: "Document"
+        )]
+        XCTAssertEqual(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Document", candidates: candidates, nativeWindowID: 0
+        ), 11)
+        XCTAssertNil(SnapshotWindowIdentityEvidence.mappedWindowID(
+            axTitle: "Other", candidates: candidates, nativeWindowID: 0
+        ))
+    }
+
     func testStageManagerFallsBackToUniqueBilateralTitleWhenNoFrameMatches() {
         let candidates = [
             SnapshotWindowIdentityEvidence.Candidate(

src/utils/computerUse/cuHelperDaemon.snapshot.test.ts

@@ -0,0 +1,102 @@
+import { afterEach, describe, expect, test } from 'bun:test'
+import { EventEmitter } from 'node:events'
+import {
+  __resetDaemonClientForTests,
+  __setDaemonSocketForTests,
+  callDaemon,
+  overlayHide,
+} from './cuHelperDaemon.js'
+import { __resetHelperBridgeState, callHelper } from './helperBridge.js'
+
+type Request = {
+  id: string
+  cmd: string
+  turnId: string
+  payload: Record<string, unknown>
+}
+
+/** Models the native contract: snapshots belong to a turn and turn_end clears them. */
+class SnapshotSocket extends EventEmitter {
+  requests: Request[] = []
+  destroyed = false
+  private snapshotTurn: string | undefined
+
+  write(data: string): boolean {
+    const request = JSON.parse(data) as Request
+    this.requests.push(request)
+    queueMicrotask(() => {
+      let result: unknown = true
+      let error: string | undefined
+      switch (request.cmd) {
+        case 'get_app_state':
+          this.snapshotTurn = request.turnId
+          result = { axText: 'g1:0 button Search', screenshot: { base64: 'fixture' } }
+          break
+        case 'overlay_show':
+          error = 'target_not_running'
+          break
+        case 'turn_end':
+          this.snapshotTurn = undefined
+          break
+        case 'click':
+          if (this.snapshotTurn !== request.turnId) {
+            error = request.payload.index
+              ? 'No element snapshot exists in the active turn'
+              : 'No screenshot snapshot exists for this target'
+          }
+          break
+        default:
+          error = `Unexpected fixture command: ${request.cmd}`
+      }
+      this.emit('data', Buffer.from(`${JSON.stringify({
+        id: request.id,
+        ok: !error,
+        ...(error ? { error: { message: error } } : { result }),
+      })}\n`))
+    })
+    return true
+  }
+
+  destroy(): this {
+    this.destroyed = true
+    return this
+  }
+}
+
+afterEach(() => {
+  __resetDaemonClientForTests()
+  __resetHelperBridgeState()
+})
+
+describe('cu-helper snapshot lifetime after visual feedback failure', () => {
+  test.each([
+    { label: 'coordinates', target: { x: 20, y: 30 } },
+    { label: 'element handle', target: { index: 'g1:0' } },
+  ])('a failed overlay preserves the snapshot for $label', async ({ target }) => {
+    const socket = new SnapshotSocket()
+    __setDaemonSocketForTests(socket as never)
+
+    // Use the real bridge: it starts overlay_show without awaiting it, then reads state.
+    await callHelper('get_app_state', { app: 'Fixture App' }, {
+      platform: 'darwin',
+      cuHelperAvailable: () => true,
+    })
+    // Allow the fire-and-forget overlay reconciliation to finish before the next tool.
+    await new Promise<void>(resolve => setImmediate(resolve))
+
+    await expect(callDaemon('click', { app: 'Fixture App', ...target })).resolves.toBe(true)
+    expect(socket.requests.map(request => request.cmd)).toEqual([
+      'overlay_show', 'get_app_state', 'click',
+    ])
+    const snapshot = socket.requests.find(request => request.cmd === 'get_app_state')!
+    const click = socket.requests.find(request => request.cmd === 'click')!
+    expect(click.turnId).toBe(snapshot.turnId)
+
+    // Explicit host cleanup must still release the read snapshot despite the failed overlay.
+    await overlayHide()
+    expect(socket.requests.at(-1)).toMatchObject({
+      cmd: 'turn_end',
+      turnId: snapshot.turnId,
+    })
+  })
+})

src/utils/computerUse/cuHelperDaemon.test.ts

@@ -443,6 +443,8 @@ describe('cu-helper overlay reconciliation', () => {
     await show
 
     expect(isOverlayShown()).toBe(false)
+    await new Promise<void>(resolve => setImmediate(resolve))
+    expect(socket.writes).toHaveLength(1)
     const hide = overlayHide()
     await waitForWriteCount(socket, 2)
     expect(JSON.parse(socket.writes[1]!)).toMatchObject({ cmd: 'turn_end' })
@@ -450,6 +452,27 @@ describe('cu-helper overlay reconciliation', () => {
     await hide
   })
 
+  test('cleanup requested during a failed show waits until turn_end completes', async () => {
+    const socket = new FakeSocket()
+    __setDaemonSocketForTests(socket as never)
+
+    const show = overlayShow({ app: 'TextEdit' })
+    await waitForWriteCount(socket, 1)
+    let cleanupCompleted = false
+    const hide = overlayHide().then(() => { cleanupCompleted = true })
+    reply(socket, 0, { ok: false, error: { message: 'target_not_running' } })
+    await waitForWriteCount(socket, 2)
+    expect(JSON.parse(socket.writes[1]!)).toMatchObject({ cmd: 'turn_end' })
+    await new Promise<void>(resolve => setImmediate(resolve))
+    expect(cleanupCompleted).toBe(false)
+
+    reply(socket, 1, { ok: true, result: true })
+    await Promise.all([show, hide])
+    expect(cleanupCompleted).toBe(true)
+    expect(isOverlayShown()).toBe(false)
+    expect(socket.writes).toHaveLength(2)
+  })
+
   test('cleanup ends a read-only turn even when no overlay was shown', async () => {
     const socket = new FakeSocket()
     __setDaemonSocketForTests(socket as never)

src/utils/computerUse/cuHelperDaemon.ts

@@ -112,6 +112,7 @@ let daemonStartCount = 0
 
 let overlayDesiredVisible = false
 let overlayActualVisible = false
+let overlayCleanupRequested = false
 let overlayDesiredPayload: Record<string, unknown> = {}
 let overlayDesiredKey = '{}'
 let overlayActualKey: string | undefined
@@ -391,6 +392,7 @@ function resetState(reason: string, expectedGeneration?: number): void {
 
   overlayDesiredVisible = false
   overlayActualVisible = false
+  overlayCleanupRequested = false
   overlayActualKey = undefined
   overlayRevision++
   activeTurnId = undefined
@@ -697,7 +699,12 @@ function needsOverlayReconciliation(): boolean {
   // sleep assertion) still needs an explicit turn_end at host cleanup. The
   // keyed SCStream consumer deliberately survives that boundary and retires on
   // target/config changes or daemon teardown.
-  if (!overlayDesiredVisible) return overlayActualVisible || activeTurnId !== undefined
+  // A failed visual-feedback request also leaves the overlay hidden, but must
+  // not release the active turn's snapshots. Only explicit host cleanup owns
+  // that lifetime boundary.
+  if (!overlayDesiredVisible) {
+    return overlayCleanupRequested && (overlayActualVisible || activeTurnId !== undefined)
+  }
   return !overlayActualVisible || overlayActualKey !== overlayDesiredKey
 }
 
@@ -725,7 +732,7 @@ async function reconcileOverlay(): Promise<void> {
         overlayActualVisible = false
         overlayActualKey = undefined
         logForDebugging(`cu-helper overlay_show failed: ${String(err)}`, { level: 'debug' })
-        return
+        continue
       }
       continue
     }
@@ -734,6 +741,7 @@ async function reconcileOverlay(): Promise<void> {
     // pending show completes this branch sees the already-owned daemon and
     // serially ends the turn. This also covers read-only turns whose overlay
     // was never visible.
+    overlayCleanupRequested = false
     if (!statePromise || activeDaemonGeneration === undefined) {
       overlayActualVisible = false
       overlayActualKey = undefined
@@ -770,6 +778,7 @@ export function overlayShow(
   target: Record<string, unknown> = {},
 ): Promise<void> {
   overlayDesiredVisible = true
+  overlayCleanupRequested = false
   overlayDesiredPayload = { ...target }
   overlayDesiredKey = JSON.stringify(overlayDesiredPayload)
   overlayRevision++
@@ -779,6 +788,7 @@ export function overlayShow(
 /** Hide the overlay, serialized after any pending show (best-effort). */
 export function overlayHide(): Promise<void> {
   overlayDesiredVisible = false
+  overlayCleanupRequested = true
   overlayRevision++
   return scheduleOverlayReconciliation()
 }
@@ -817,6 +827,7 @@ export function __resetDaemonClientForTests(): void {
   activeDaemonGeneration = undefined
   overlayDesiredVisible = false
   overlayActualVisible = false
+  overlayCleanupRequested = false
   overlayDesiredPayload = {}
   overlayDesiredKey = '{}'
   overlayActualKey = undefined