fix(ui-automation): Tighten post-action snapshot handling

Skip post-action snapshot capture for safe same-screen batches, add long-press capture parity, and cover runtime next-step foreground ranking rules directly.

Co-Authored-By: Codex <noreply@openai.com>

Author: cameroncooke

src/mcp/tools/ui-automation/__tests__/batch.test.ts

@@ -203,16 +203,22 @@ describe('Batch UI Automation Tool', () => {
         createNode({ type: 'Switch', role: 'AXSwitch', AXValue: '0' }),
         createNode({ type: 'Switch', role: 'AXSwitch', AXValue: 'off' }),
       ]);
+      const { calls, executor } = createTrackingExecutor();
 
-      const result = await runBatch({
-        simulatorId,
-        steps: [
-          { action: 'tap', elementRef: 'e1' },
-          { action: 'tap', elementRef: 'e2' },
-        ],
-      });
+      const result = await runBatch(
+        {
+          simulatorId,
+          steps: [
+            { action: 'tap', elementRef: 'e1' },
+            { action: 'tap', elementRef: 'e2' },
+          ],
+        },
+        executor,
+      );
 
       expect(result.didError).toBe(false);
+      expect(result.capture).toBeUndefined();
+      expect(calls.some((call) => call.command[1] === 'describe-ui')).toBe(false);
       expect(getRuntimeSnapshot(simulatorId)).not.toBeNull();
     });
 

src/mcp/tools/ui-automation/__tests__/long_press.test.ts

@@ -59,8 +59,8 @@ describe('Long Press Plugin', () => {
       expect(result).toMatchObject({
         didError: false,
         action: { type: 'long-press', elementRef: 'e1', durationMs: 1500 },
+        capture: { type: 'runtime-snapshot', simulatorId },
       });
-      expect(calls).toHaveLength(1);
       expect(calls[0]?.command).toEqual([
         '/mocked/axe/path',
         'touch',

src/mcp/tools/ui-automation/__tests__/runtime-next-steps.test.ts

@@ -0,0 +1,161 @@
+import { beforeEach, describe, expect, it } from 'vitest';
+import type { AccessibilityNode } from '../../../../types/domain-results.ts';
+import { createRuntimeSnapshotNextSteps } from '../shared/runtime-next-steps.ts';
+import {
+  __resetRuntimeSnapshotStoreForTests,
+  getRuntimeSnapshot,
+} from '../shared/snapshot-ui-state.ts';
+import { createNode, recordSnapshot, simulatorId } from './ui-action-test-helpers.ts';
+
+function currentRuntimeSnapshot() {
+  const snapshot = getRuntimeSnapshot(simulatorId);
+  expect(snapshot).not.toBeNull();
+  return snapshot!.payload;
+}
+
+function createScrollView(overrides: Partial<AccessibilityNode> = {}): AccessibilityNode {
+  return createNode({
+    type: 'ScrollView',
+    role: 'AXScrollArea',
+    frame: { x: 0, y: 0, width: 390, height: 844 },
+    AXIdentifier: 'scroll-view',
+    ...overrides,
+  });
+}
+
+function nestNode(node: AccessibilityNode, depth: number): AccessibilityNode {
+  let current = node;
+  for (let index = 0; index < depth; index += 1) {
+    current = createNode({
+      type: 'Group',
+      role: 'AXGroup',
+      AXIdentifier: `container.${index}`,
+      frame: current.frame,
+      children: [current],
+    });
+  }
+  return current;
+}
+
+describe('runtime snapshot next steps', () => {
+  beforeEach(() => {
+    __resetRuntimeSnapshotStoreForTests();
+  });
+
+  it('prefers tap and scroll examples from the active foreground container', () => {
+    recordSnapshot([
+      createScrollView({
+        AXIdentifier: 'weather.backgroundList',
+        children: [
+          createNode({
+            AXLabel: 'Background, Details',
+            AXIdentifier: 'weather.backgroundCard',
+            frame: { x: 20, y: 120, width: 350, height: 80 },
+          }),
+        ],
+      }),
+      createScrollView({
+        AXIdentifier: 'weather.settingsSheet',
+        frame: { x: 0, y: 420, width: 390, height: 424 },
+        children: [
+          createNode({ AXLabel: 'Close', frame: { x: 310, y: 430, width: 60, height: 40 } }),
+          createNode({
+            type: 'TextField',
+            role: 'AXTextField',
+            AXLabel: 'Search',
+            frame: { x: 20, y: 480, width: 350, height: 40 },
+          }),
+          createNode({
+            AXLabel: 'London, England',
+            AXIdentifier: 'weather.locationCard',
+            frame: { x: 20, y: 540, width: 350, height: 80 },
+          }),
+        ],
+      }),
+    ]);
+
+    const snapshot = currentRuntimeSnapshot();
+    const foregroundScrollRef = snapshot.elements.find(
+      (element) => element.identifier === 'weather.settingsSheet',
+    )?.ref;
+    const foregroundCardRef = snapshot.elements.find(
+      (element) => element.identifier === 'weather.locationCard',
+    )?.ref;
+
+    const steps = createRuntimeSnapshotNextSteps({
+      simulatorId,
+      runtimeSnapshot: snapshot,
+      includeRefreshAndWait: false,
+    });
+
+    expect(steps).toContainEqual({
+      label: 'Tap an elementRef',
+      tool: 'tap',
+      params: { simulatorId, elementRef: foregroundCardRef },
+    });
+    expect(steps).toContainEqual({
+      label: 'Scroll visible content',
+      tool: 'swipe',
+      params: {
+        simulatorId,
+        withinElementRef: foregroundScrollRef,
+        direction: 'up',
+        distance: 0.5,
+      },
+    });
+  });
+
+  it('uses hierarchy depth only as a foreground-root tie breaker', () => {
+    recordSnapshot([
+      nestNode(
+        createScrollView({
+          AXIdentifier: 'deep.stateControls',
+          frame: { x: 0, y: 0, width: 390, height: 80 },
+          children: [
+            createNode({
+              type: 'Switch',
+              role: 'AXSwitch',
+              AXLabel: 'Nested switch',
+              AXValue: '0',
+            }),
+          ],
+        }),
+        40,
+      ),
+      createScrollView({
+        AXIdentifier: 'shallow.searchPanel',
+        frame: { x: 0, y: 100, width: 390, height: 500 },
+        children: [
+          createNode({
+            type: 'TextField',
+            role: 'AXTextField',
+            AXLabel: 'Search',
+            frame: { x: 20, y: 130, width: 350, height: 40 },
+          }),
+        ],
+      }),
+    ]);
+
+    const snapshot = currentRuntimeSnapshot();
+    const shallowSearchRef = snapshot.elements.find(
+      (element) => element.identifier === 'shallow.searchPanel',
+    )?.ref;
+
+    const steps = createRuntimeSnapshotNextSteps({
+      simulatorId,
+      runtimeSnapshot: snapshot,
+      includeRefreshAndWait: false,
+    });
+
+    expect(steps).toContainEqual({
+      label: 'Scroll visible content',
+      tool: 'swipe',
+      params: {
+        simulatorId,
+        withinElementRef: shallowSearchRef,
+        direction: 'up',
+        distance: 0.5,
+      },
+    });
+  });
+});

src/mcp/tools/ui-automation/batch.ts

@@ -216,6 +216,10 @@ export function createBatchExecutor(
       });
     }
 
+    if (resolvedSteps.preserveSnapshot) {
+      return createUiActionSuccessResult(action, simulatorId, [guard.warningText]);
+    }
+
     const captureResult = await captureRuntimeSnapshotAfterActionSafely({
       simulatorId,
       executor,

src/mcp/tools/ui-automation/long_press.ts

@@ -20,6 +20,7 @@ import {
 import { clearRuntimeSnapshot, resolveElementRef } from './shared/snapshot-ui-state.ts';
 import { getRuntimeElementActivationPoint } from './shared/runtime-snapshot.ts';
 import { executeAxeCommand, defaultAxeHelpers } from './shared/axe-command.ts';
+import { captureRuntimeSnapshotAfterActionSafely } from './shared/post-action-snapshot.ts';
 import type { AxeHelpers } from './shared/axe-command.ts';
 import type { NonStreamingExecutor } from '../../../types/tool-execution.ts';
 import type { UiActionResultDomainResult } from '../../../types/domain-results.ts';
@@ -100,7 +101,6 @@ export function createLongPressExecutor(
       await executeAxeCommand(commandArgs, simulatorId, 'touch', executor, axeHelpers);
       clearRuntimeSnapshot(simulatorId);
       log('info', `${LOG_PREFIX}/${toolName}: Success for ${simulatorId}`);
-      return createUiActionSuccessResult(action, simulatorId, [guard.warningText]);
     } catch (error) {
       if (shouldInvalidateRuntimeSnapshotAfterActionError(error)) {
         clearRuntimeSnapshot(simulatorId);
@@ -118,6 +118,21 @@ export function createLongPressExecutor(
         }),
       });
     }
+
+    const captureResult = await captureRuntimeSnapshotAfterActionSafely({
+      simulatorId,
+      executor,
+      axeHelpers,
+    });
+    return createUiActionSuccessResult(
+      action,
+      simulatorId,
+      [guard.warningText, captureResult.warning],
+      {
+        ...(captureResult.capture ? { capture: captureResult.capture } : {}),
+        ...(captureResult.uiError ? { uiError: captureResult.uiError } : {}),
+      },
+    );
   };
 }
 

src/mcp/tools/ui-automation/shared/runtime-next-steps.ts

@@ -270,8 +270,8 @@ function findActiveForegroundRoot(
       (hasDismissControl ? 100 : 0) +
       (hasTextEntry ? 60 : 0) +
       (hasStateControls ? 30 : 0) +
-      record.metadata.depth +
-      (indexByRef.get(record.publicElement.ref) ?? 0) / 1000;
+      record.metadata.depth / 1000 +
+      (indexByRef.get(record.publicElement.ref) ?? 0) / 1_000_000;
     scoreByRef.set(record.publicElement.ref, score);
     return score;
   }