fix(ios-runner): keep snapshot-plan target state on main and bound the query sweep by its slice

.github/workflows/ios.yml

@@ -256,6 +256,10 @@ jobs:
             -only-testing:AgentDeviceRunnerUITests/RunnerTests/testActionNamesAreCappedPerElementAndReported \
             -only-testing:AgentDeviceRunnerUITests/RunnerTests/testHungCustomActionReadIsContainedAndRecovers \
             -only-testing:AgentDeviceRunnerUITests/RunnerTests/testAbandonedTreeCaptureSkipsQuerySweepAndHonorsWarmupExemption \
+            -only-testing:AgentDeviceRunnerUITests/RunnerTests/testNonInteractiveQuerySweepStopsAtTheSliceItsCallerWaitsFor \
+            -only-testing:AgentDeviceRunnerUITests/RunnerTests/testQuerySweepThatEndsOnItsSliceDeadlinePenalizesChannelAndReachesPrivateAX \
+            -only-testing:AgentDeviceRunnerUITests/RunnerTests/testCoordinateTapRoutingPenalizesTheIdentityMainSettledOnWhileTheWriteWasPending \
+            -only-testing:AgentDeviceRunnerUITests/RunnerTests/testBlockingModalSnapshotLeavesWarmupExemptionForTheFirstCapturePlan \
             -only-testing:AgentDeviceRunnerUITests/RunnerTests/testPrivateAXPinnedRegularDepthReachesAcquisitionAndPresentation 2>&1 | tee /tmp/agent-device-runner-regressions.log
           node --input-type=module -e '
             import { readFileSync } from "node:fs";

apple/runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+AXSnapshotFallback.swift

@@ -110,14 +110,16 @@ extension RunnerTests {
   /// remembered depth deliberately does NOT refresh the TTL, so expiry re-probes the full
   /// requested depth once per window instead of capping this screen class forever.
   func recordPrivateAXAcceptedDepth(
+    bundleId: String?,
+    processIdentifier: Int?,
     exactDepthRequested: Bool,
     effectiveDepth: Int,
     attemptDepths: [Int]
   ) {
     guard !exactDepthRequested, effectiveDepth != attemptDepths.first else { return }
     rememberPrivateAXAcceptedDepth(
-      bundleId: currentBundleId,
-      processIdentifier: currentAppProcessIdentifier,
+      bundleId: bundleId,
+      processIdentifier: processIdentifier,
       depth: effectiveDepth
     )
   }
@@ -160,11 +162,12 @@ extension RunnerTests {
   }
 
   func privateAXSnapshotAcquisition(
-    app: XCUIApplication,
+    target: SnapshotCaptureTarget,
     hint: CaptureHint,
     deadline: Date = .distantFuture
   ) -> SnapshotAcquisition? {
     #if os(iOS) && targetEnvironment(simulator)
+      let app = target.app
       let requestedDepth = hint.rawTraversalDepth ?? 64
       // An explicit --depth request is honored as asked: no accepted-depth
       // memory, no frontier extension past it.
@@ -173,8 +176,8 @@ extension RunnerTests {
         exactDepthRequested
         ? nil
         : rememberedPrivateAXAcceptedDepth(
-          bundleId: currentBundleId,
-          processIdentifier: currentAppProcessIdentifier
+          bundleId: target.bundleId,
+          processIdentifier: target.processIdentifier
         )
       let attemptDepths = Self.privateAXAttemptDepths(
         requestedDepth: requestedDepth,
@@ -202,6 +205,8 @@ extension RunnerTests {
         return nil
       }
       recordPrivateAXAcceptedDepth(
+        bundleId: target.bundleId,
+        processIdentifier: target.processIdentifier,
         exactDepthRequested: exactDepthRequested,
         effectiveDepth: effectiveDepth,
         attemptDepths: attemptDepths
@@ -212,7 +217,11 @@ extension RunnerTests {
       }
 
       let rootFrame = privateAXRect(root["frame"])
-      let geometry = privateAXSnapshotGeometry(app: app, rootFrame: rootFrame)
+      let geometry = privateAXSnapshotGeometry(
+        app: app,
+        bundleId: target.bundleId,
+        rootFrame: rootFrame
+      )
       let viewport = geometry.viewport
       let nodes = privateAXAcquisition(
         rawRoot: root,
@@ -262,8 +271,8 @@ extension RunnerTests {
   /// grinding on this screen class. Under penalty it reliably burns its full timeout and
   /// falls back anyway (~1s added to every private AX capture on the Bluesky bench feed),
   /// so honor the penalty here the same way capture plans do.
-  func shouldReadPrivateAXViewportViaXCTest() -> Bool {
-    !hasAbandonedMainThreadWork() && !isSnapshotXCTestChannelPenalized(bundleId: currentBundleId)
+  func shouldReadPrivateAXViewportViaXCTest(bundleId: String?) -> Bool {
+    !hasAbandonedMainThreadWork() && !isSnapshotXCTestChannelPenalized(bundleId: bundleId)
   }
 
   /// The geometry this tier may anchor a rotation on. The bridge's own root frame is one more
@@ -272,10 +281,11 @@ extension RunnerTests {
   /// consumers already treat as geometry they cannot measure (#2612).
   private func privateAXSnapshotGeometry(
     app: XCUIApplication,
+    bundleId: String?,
     rootFrame: CGRect
   ) -> (viewport: CGRect, interfaceOrientation: Int) {
     let fallback = rootFrame.isEmpty ? CGRect.infinite : rootFrame
-    guard shouldReadPrivateAXViewportViaXCTest() else {
+    guard shouldReadPrivateAXViewportViaXCTest(bundleId: bundleId) else {
       return (fallback, RunnerInterfaceOrientation.unknown)
     }
     do {

apple/runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+CommandDispatch.swift

@@ -579,8 +579,13 @@ extension RunnerTests {
     }
     #endif
     let probeDeadline = Date().addingTimeInterval(systemModalProbeBudget)
+    // Routing runs on the command queue, so this hands the probe a target rather than a bundle id it
+    // read across the main boundary: the penalty an abandoned probe arms carries the identity main
+    // holds once the probe starts, not one whose write was still queued behind the block that
+    // occupied main (#2781).
     return boundedBlockingSystemAlertSnapshot(
-      deadline: probeDeadline
+      deadline: probeDeadline,
+      penaltyTarget: .mainOwnedTarget
     ) != nil
 #else
     return false

apple/runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+Lifecycle.swift

@@ -202,7 +202,7 @@ extension RunnerTests {
   /// markers. Every site that binds, rebinds, or drops the target runs this.
   func resetTargetBoundState() {
     clearRememberedTextEntryTap()
-    snapshotXCTestPenaltyWarmupExemptionPending = false
+    snapshotXCTestPenaltyWarmupExemption.isPending = false
     lastLoggedFastAppGuardLine = nil
     lastLoggedGesturePolicyLines.removeAll()
   }
@@ -247,7 +247,7 @@ extension RunnerTests {
     resetTargetBoundState()
     clearSnapshotXCTestChannelPenalty(reason: "target_process_changed")
     clearPrivateAXAcceptedDepth(reason: "target_process_changed")
-    snapshotXCTestPenaltyWarmupExemptionPending = true
+    snapshotXCTestPenaltyWarmupExemption.isPending = true
     beginFirstInteractionStabilization()
   }
 

apple/runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+Snapshot.swift

@@ -132,6 +132,20 @@ extension RunnerTests {
   ]
 
   static let flatInteractiveFallbackBudget: TimeInterval = 1.0
+  /// The least slice time a sweep query may start with. XCTest cannot cancel a query, so one that
+  /// starts later outlives the slice its caller waits for and holds the main thread (#2783).
+  static let flatInteractiveQueryBudget: TimeInterval = 0.1
+
+  /// The deadline the query-sweep tier's caller waits for: one slice, clamped to the plan deadline.
+  /// Interactive and non-interactive requests share it, since the caller discards a later result.
+  static func querySweepSliceDeadline(startedAt: Date, planDeadline: Date) -> Date {
+    min(startedAt.addingTimeInterval(flatInteractiveFallbackBudget), planDeadline)
+  }
+
+  /// Whether a sweep query started at `now` can still finish before the slice `deadline`.
+  static func querySweepCanStartQuery(deadline: Date, now: Date) -> Bool {
+    deadline.timeIntervalSince(now) >= flatInteractiveQueryBudget
+  }
 
   /// What one capture may spend reading the keyboard band before it gives up on the fact and lets the
   /// tap guard fall back to the tree rule. The scroll path pays this query per gesture and stays well
@@ -145,14 +159,17 @@ extension RunnerTests {
   // `boundedBlockingSystemAlertSnapshot`'s probe closure (see `systemModalProbeOverrideForTesting`
   // in RunnerTests.swift), so reverting this entry point to bypass the bounded probe fails the
   // regression test.
-  func snapshotFast(app: XCUIApplication, options: PresentationOptions) throws -> DataPayload {
+  func snapshotFast(target: SnapshotCaptureTarget, options: PresentationOptions) throws -> DataPayload {
     let deadline = Date().addingTimeInterval(Self.snapshotPlanBudget)
-    if let blocking = boundedBlockingSystemAlertSnapshot(deadline: deadline) {
+    if let blocking = boundedBlockingSystemAlertSnapshot(
+      deadline: deadline,
+      penaltyTarget: .prepared(bundleId: target.bundleId)
+    ) {
       return blocking
     }
     return try runSnapshotCapturePlan(
       Self.regularVisiblePlan,
-      app: app,
+      target: target,
       options: options,
       terminal: .sparseWithFatalOnAXFailure,
       deadline: deadline
@@ -267,14 +284,17 @@ extension RunnerTests {
   }
 
   // See `snapshotFast` above: the single production entry point, no unit-test overload.
-  func snapshotRaw(app: XCUIApplication, options: PresentationOptions) throws -> DataPayload {
+  func snapshotRaw(target: SnapshotCaptureTarget, options: PresentationOptions) throws -> DataPayload {
     let deadline = Date().addingTimeInterval(Self.snapshotPlanBudget)
-    if let blocking = boundedBlockingSystemAlertSnapshot(deadline: deadline) {
+    if let blocking = boundedBlockingSystemAlertSnapshot(
+      deadline: deadline,
+      penaltyTarget: .prepared(bundleId: target.bundleId)
+    ) {
       return blocking
     }
     return try runSnapshotCapturePlan(
       Self.rawDiagnosticPlan,
-      app: app,
+      target: target,
       options: options,
       terminal: .throwOnAXFailure,
       deadline: deadline
@@ -283,8 +303,15 @@ extension RunnerTests {
 
   /// Runs the pre-plan SpringBoard system-modal probe as a bounded capture tier sharing the plan
   /// deadline, so a slow alert enumeration cannot bypass the snapshot timeout and stall (#1244).
-  func boundedBlockingSystemAlertSnapshot(deadline: Date) -> DataPayload? {
-    boundedBlockingSystemAlertSnapshotBody(deadline: deadline) { probeDeadline in
+  /// An abandoned probe penalizes the XCTest channel for `penaltyTarget`.
+  func boundedBlockingSystemAlertSnapshot(
+    deadline: Date,
+    penaltyTarget: SnapshotProbePenaltyTarget
+  ) -> DataPayload? {
+    boundedBlockingSystemAlertSnapshotBody(
+      deadline: deadline,
+      penaltyTarget: penaltyTarget
+    ) { probeDeadline in
       #if AGENT_DEVICE_RUNNER_UNIT_TESTS
       if let override = self.systemModalProbeOverrideForTesting {
         return override(probeDeadline)
@@ -301,6 +328,7 @@ extension RunnerTests {
   /// production runs and what the unit tests exercise.
   private func boundedBlockingSystemAlertSnapshotBody(
     deadline: Date,
+    penaltyTarget: SnapshotProbePenaltyTarget,
     probe: @escaping (Date) -> DataPayload?
   ) -> DataPayload? {
     #if os(macOS)
@@ -316,6 +344,7 @@ extension RunnerTests {
     }
     let probeDeadline = Date().addingTimeInterval(slice)
     let startedAt = Date()
+    let penaltyIdentity = SnapshotProbePenaltyIdentity(penaltyTarget)
     do {
       return try runMainThreadWork(
         "system_modal_probe",
@@ -329,12 +358,13 @@ extension RunnerTests {
         },
         onAbandoned: {
           self.penalizeSnapshotXCTestChannel(
-            bundleId: self.currentBundleId,
+            bundleId: penaltyIdentity.penalizedBundleId,
             reason: "system_modal_probe_timeout"
           )
         }
       ) {
-        probe(probeDeadline)
+        penaltyIdentity.captureFromMain(bundleId: self.currentBundleId)
+        return probe(probeDeadline)
       }
     } catch {
       NSLog(
@@ -413,37 +443,34 @@ extension RunnerTests {
   func querySweepSnapshotAcquisition(
     app: XCUIApplication,
     hint: CaptureHint,
-    planDeadline: Date = .distantFuture
-  ) -> SnapshotAcquisition {
+    sliceDeadline deadline: Date
+  ) -> (acquisition: SnapshotAcquisition, outcome: SnapshotTierOutcome) {
     var nodes: [RawAXNode] = [
       interactiveRootNode(rect: .zero)
     ]
     if hint.rawTraversalDepth == 0 || hint.regularPresentedDepth == 0 {
-      return SnapshotAcquisition(
-        hint: hint,
-        nodes: nodes,
-        truncated: false,
-        effectiveDepth: nil,
-        viewport: .infinite,
-        interfaceOrientation: RunnerInterfaceOrientation.unknown
+      return (
+        SnapshotAcquisition(
+          hint: hint,
+          nodes: nodes,
+          truncated: false,
+          effectiveDepth: nil,
+          viewport: .infinite,
+          interfaceOrientation: RunnerInterfaceOrientation.unknown
+        ),
+        .completed
       )
     }
 
-    // Bounded by both its own sweep budget and the umbrella capture-plan deadline, so a
-    // chained recovery tier can never push the plan past the main-thread watchdog (#1105).
-    let sweepDeadline = hint.interactiveOnly
-      ? Date().addingTimeInterval(Self.flatInteractiveFallbackBudget)
-      : Date.distantFuture
-    let deadline = min(sweepDeadline, planDeadline)
     let viewport = safeSnapshotViewport(app: app)
     var seen = Set<String>()
     var candidates: [RawAXNode] = []
     let flatElements = flatInteractiveElements(app: app, deadline: deadline)
-    var truncated = flatElements.truncated
+    var outcome = flatElements.outcome
     for element in flatElements.elements {
-      if Date() >= deadline {
+      if !Self.querySweepCanStartQuery(deadline: deadline, now: Date()) {
         NSLog("AGENT_DEVICE_RUNNER_SNAPSHOT_FLAT_FALLBACK_DEADLINE")
-        truncated = true
+        outcome = .deadlineExhausted
         break
       }
       guard let node = flatSnapshotNode(element: element, index: 0, parentIndex: 0) else {
@@ -491,20 +518,25 @@ extension RunnerTests {
         )
       )
     }
-    return SnapshotAcquisition(
-      hint: hint,
-      nodes: nodes,
-      truncated: truncated,
-      effectiveDepth: nil,
-      viewport: viewport,
-      interfaceOrientation: RunnerInterfaceOrientation.unknown
+    return (
+      SnapshotAcquisition(
+        hint: hint,
+        nodes: nodes,
+        truncated: outcome == .deadlineExhausted,
+        effectiveDepth: nil,
+        viewport: viewport,
+        interfaceOrientation: RunnerInterfaceOrientation.unknown
+      ),
+      outcome
     )
   }
 
   func snapshotAccessibilityUnavailable(failure: SnapshotCaptureFailure) -> DataPayload {
     NSLog("AGENT_DEVICE_RUNNER_SNAPSHOT_AX_UNAVAILABLE=%@", failure.message)
-    runnerAccessibilityHealth = .unavailable
-    invalidateCachedTarget(reason: Self.axSnapshotUnavailableReason)
+    applyMainOwnedSnapshotState("ax_unavailable_invalidation") {
+      self.runnerAccessibilityHealth = .unavailable
+      self.invalidateCachedTarget(reason: Self.axSnapshotUnavailableReason)
+    }
     // This is a planned terminal result, so it carries the structured verdict like every other
     // planned snapshot — downstream sparse handling keys off the verdict, not node shapes.
     return sparseTruncatedSnapshotPayload(

apple/runner/AgentDeviceRunner/AgentDeviceRunnerUITests/RunnerTests+SnapshotAcquisition.swift

@@ -464,7 +464,7 @@ extension RunnerTests {
   func flatInteractiveElements(
     app: XCUIApplication,
     deadline: Date
-  ) -> (elements: [XCUIElement], truncated: Bool) {
+  ) -> (elements: [XCUIElement], outcome: SnapshotTierOutcome) {
     let queries: [XCUIElementQuery] = [
       app.buttons,
       app.links,
@@ -487,23 +487,36 @@ extension RunnerTests {
       app.images
     ]
 
-    var elements: [XCUIElement] = []
-    var truncated = false
+    return Self.runFlatInteractiveQueries(queries, deadline: deadline) { query in
+      self.snapshotElementsQuery {
+        query.allElementsBoundByIndex
+      }
+    }
+  }
+
+  /// Runs sweep queries in order until one reports AX unavailable, or until the next one could not
+  /// finish before `deadline` (`querySweepCanStartQuery`). A sweep stopped by the deadline reports
+  /// `.deadlineExhausted` rather than a truncation flag: what it collected is a partial tree, and
+  /// only the caller that owns the tier decides whether that counts as an answer (#2781).
+  static func runFlatInteractiveQueries<Query, Element>(
+    _ queries: [Query],
+    deadline: Date,
+    now: () -> Date = { Date() },
+    run: (Query) -> (elements: [Element], axUnavailable: Bool)
+  ) -> (elements: [Element], outcome: SnapshotTierOutcome) {
+    var elements: [Element] = []
     for query in queries {
-      if Date() >= deadline {
+      if !querySweepCanStartQuery(deadline: deadline, now: now()) {
         NSLog("AGENT_DEVICE_RUNNER_SNAPSHOT_FLAT_FALLBACK_DEADLINE")
-        truncated = true
-        break
-      }
-      let result = snapshotElementsQuery {
-        query.allElementsBoundByIndex
+        return (elements, .deadlineExhausted)
       }
+      let result = run(query)
       elements.append(contentsOf: result.elements)
       if result.axUnavailable {
         break
       }
     }
-    return (elements, truncated)
+    return (elements, .completed)
   }
 
   func snapshotElementsQuery(