From 05a2bbeeeba1aee8441ab0d818cafb4c01240030 Mon Sep 17 00:00:00 2001
From: Ryan Vogel <89211796+R44VC0RP@users.noreply.github.com>
Date: Wed, 29 Jul 2026 19:46:41 -0400
Subject: [PATCH] feat(cli): add ax command to dump the accessibility tree

---
 packages/serve-sim/README.md                  |   1 +
 .../serve-sim/src/__tests__/ax-cli.test.ts    | 142 ++++++++++++++++++
 packages/serve-sim/src/ax-cli.ts              |  35 +++++
 packages/serve-sim/src/ax.ts                  |   4 +-
 packages/serve-sim/src/index.ts               |  24 +++
 5 files changed, 204 insertions(+), 2 deletions(-)
 create mode 100644 packages/serve-sim/src/__tests__/ax-cli.test.ts
 create mode 100644 packages/serve-sim/src/ax-cli.ts

diff --git a/packages/serve-sim/README.md b/packages/serve-sim/README.md
index 2cb76d28..16c5cc60 100644
--- a/packages/serve-sim/README.md
+++ b/packages/serve-sim/README.md
@@ -53,6 +53,7 @@ serve-sim ca-debug <option> <on|off> [-d udid]
                                       (blended|copies|misaligned|offscreen|slow-animations)
 serve-sim memory-warning [-d udid]    Simulate a memory warning
 serve-sim event-log [-d udid]         Show recent simulator events
+serve-sim ax [-d udid]                Dump the accessibility tree as JSON
 
 serve-sim camera <bundle-id> [-d udid] [source-options]
                                       Inject a synthetic camera feed and (re)launch the app
diff --git a/packages/serve-sim/src/__tests__/ax-cli.test.ts b/packages/serve-sim/src/__tests__/ax-cli.test.ts
new file mode 100644
index 00000000..ce064e58
--- /dev/null
+++ b/packages/serve-sim/src/__tests__/ax-cli.test.ts
@@ -0,0 +1,142 @@
+import { afterAll, beforeAll, describe, expect, test } from "bun:test";
+import { createServer, type Server } from "http";
+import type { AddressInfo } from "net";
+import { normalizeAxTree, type RawAxeNode } from "../ax";
+import { axUrlFromStreamUrl, fetchAxSnapshot } from "../ax-cli";
+
+function rawNode(overrides: Partial<RawAxeNode> = {}): RawAxeNode {
+  return {
+    AXUniqueId: null,
+    AXLabel: null,
+    AXValue: null,
+    enabled: true,
+    frame: { x: 0, y: 0, width: 100, height: 50 },
+    role_description: "button",
+    type: "Button",
+    children: [],
+    ...overrides,
+  };
+}
+
+const SCREEN = { x: 0, y: 0, width: 393, height: 852 };
+
+describe("axUrlFromStreamUrl", () => {
+  test("maps the stream URL to the sibling /ax endpoint", () => {
+    expect(axUrlFromStreamUrl("http://127.0.0.1:3200/helper/UDID-1/stream.mjpeg"))
+      .toBe("http://127.0.0.1:3200/helper/UDID-1/ax");
+  });
+
+  test("only rewrites the trailing path segment", () => {
+    expect(axUrlFromStreamUrl("http://127.0.0.1:3200/stream.mjpeg/helper/UDID/stream.mjpeg"))
+      .toBe("http://127.0.0.1:3200/stream.mjpeg/helper/UDID/ax");
+  });
+});
+
+describe("normalizeAxTree", () => {
+  test("flattens the tree into elements with role, label, value, enabled state, and frame", () => {
+    const roots: RawAxeNode[] = [
+      rawNode({
+        frame: SCREEN,
+        role_description: "application",
+        type: "Application",
+        children: [
+          rawNode({
+            AXUniqueId: "login-button",
+            AXLabel: "Log in",
+            AXValue: "",
+            role_description: "button",
+            type: "Button",
+            frame: { x: 20, y: 700, width: 353, height: 44 },
+          }),
+          rawNode({
+            AXLabel: "Username",
+            AXValue: "bacon",
+            role_description: "text field",
+            type: "TextField",
+            enabled: false,
+            frame: { x: 20, y: 200, width: 353, height: 44 },
+          }),
+        ],
+      }),
+    ];
+
+    const snapshot = normalizeAxTree(roots);
+    expect(snapshot.screen).toEqual({ width: SCREEN.width, height: SCREEN.height });
+    // The screen-sized root is dropped; only real elements remain.
+    expect(snapshot.elements).toHaveLength(2);
+    expect(snapshot.elements[0]).toEqual({
+      id: "login-button",
+      path: "0.0",
+      label: "Log in",
+      value: "",
+      role: "button",
+      type: "Button",
+      enabled: true,
+      frame: { x: 20, y: 700, width: 353, height: 44 },
+    });
+    expect(snapshot.elements[1]).toMatchObject({
+      id: "0.1", // falls back to the tree path when AXUniqueId is null
+      label: "Username",
+      value: "bacon",
+      role: "text field",
+      enabled: false,
+    });
+  });
+
+  test("caps the element count so pathological trees stay bounded", () => {
+    const children = Array.from({ length: 600 }, (_, i) =>
+      rawNode({ AXLabel: `Row ${i}`, frame: { x: 0, y: i, width: 100, height: 1 } }));
+    const snapshot = normalizeAxTree([rawNode({ frame: SCREEN, children })]);
+    expect(snapshot.elements.length).toBeLessThanOrEqual(500);
+  });
+});
+
+describe("fetchAxSnapshot", () => {
+  let server: Server;
+  let streamUrl: string;
+  let status = 200;
+  let body: unknown = [];
+
+  beforeAll(async () => {
+    server = createServer((req, res) => {
+      if (req.url !== "/helper/UDID-TEST/ax") {
+        res.writeHead(404);
+        res.end();
+        return;
+      }
+      res.writeHead(status, { "Content-Type": "application/json" });
+      res.end(JSON.stringify(body));
+    });
+    await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+    const { port } = server.address() as AddressInfo;
+    streamUrl = `http://127.0.0.1:${port}/helper/UDID-TEST/stream.mjpeg`;
+  });
+
+  afterAll(() => {
+    server.close();
+  });
+
+  test("fetches the raw tree and returns the normalized snapshot", async () => {
+    status = 200;
+    body = [
+      rawNode({
+        frame: SCREEN,
+        children: [rawNode({ AXUniqueId: "ok", AXLabel: "OK", frame: { x: 10, y: 10, width: 80, height: 40 } })],
+      }),
+    ];
+
+    const snapshot = await fetchAxSnapshot(streamUrl);
+    expect(snapshot.screen).toEqual({ width: SCREEN.width, height: SCREEN.height });
+    expect(snapshot.elements).toHaveLength(1);
+    expect(snapshot.elements[0]).toMatchObject({ id: "ok", label: "OK", role: "button", enabled: true });
+  });
+
+  test("surfaces the helper's message when AX is unavailable (503)", async () => {
+    status = 503;
+    body = { error: "ax_unavailable", message: "Accessibility unavailable on this simulator." };
+
+    await expect(fetchAxSnapshot(streamUrl)).rejects.toThrow(
+      "Accessibility unavailable on this simulator.",
+    );
+  });
+});
diff --git a/packages/serve-sim/src/ax-cli.ts b/packages/serve-sim/src/ax-cli.ts
new file mode 100644
index 00000000..065627b2
--- /dev/null
+++ b/packages/serve-sim/src/ax-cli.ts
@@ -0,0 +1,35 @@
+import { normalizeAxTree } from "./ax";
+import type { RawAxeNode } from "./ax";
+import type { AxSnapshot } from "./ax-shared";
+
+/**
+ * Map a device's stream URL (`…/helper/<udid>/stream.mjpeg`) to its sibling
+ * one-shot accessibility endpoint (`…/helper/<udid>/ax`). The state file only
+ * records the stream URL, so the `ax` CLI command derives the endpoint the
+ * same way the accessibility-endpoint test does.
+ */
+export function axUrlFromStreamUrl(streamUrl: string): string {
+  return streamUrl.replace(/\/stream\.mjpeg$/, "/ax");
+}
+
+/**
+ * Fetch the raw axe-shaped tree from a running serve-sim server and normalize
+ * it into the flat {@link AxSnapshot} shape the web UI consumes (roles,
+ * labels, values, enabled state, frames).
+ *
+ * Throws with the helper's message when the endpoint reports AX unavailable
+ * (503 while the simulator's accessibility framework warms up).
+ */
+export async function fetchAxSnapshot(streamUrl: string): Promise<AxSnapshot> {
+  const res = await fetch(axUrlFromStreamUrl(streamUrl));
+  if (!res.ok) {
+    let message = `HTTP ${res.status}`;
+    try {
+      const body = await res.json() as { message?: string };
+      if (typeof body?.message === "string" && body.message) message = body.message;
+    } catch {}
+    throw new Error(message);
+  }
+  const raw = await res.json() as RawAxeNode[];
+  return normalizeAxTree(raw);
+}
diff --git a/packages/serve-sim/src/ax.ts b/packages/serve-sim/src/ax.ts
index 01e10aa8..92072cde 100644
--- a/packages/serve-sim/src/ax.ts
+++ b/packages/serve-sim/src/ax.ts
@@ -9,7 +9,7 @@ const POLL_INTERVAL_MS = 500;
 const MAX_POLL_INTERVAL_MS = 2000;
 const UNAVAILABLE_RETRY_INTERVAL_MS = 15_000;
 
-interface RawAxeNode {
+export interface RawAxeNode {
   AXUniqueId: string | null;
   AXLabel: string | null;
   AXValue: string | null;
@@ -38,7 +38,7 @@ function sameRect(a: AxRect, b: AxRect) {
   );
 }
 
-function normalizeAxTree(roots: RawAxeNode[]): AxSnapshot {
+export function normalizeAxTree(roots: RawAxeNode[]): AxSnapshot {
   const screen = chooseScreenFrame(roots);
   const elements: AxElement[] = [];
 
diff --git a/packages/serve-sim/src/index.ts b/packages/serve-sim/src/index.ts
index 24d72356..e8f0089f 100755
--- a/packages/serve-sim/src/index.ts
+++ b/packages/serve-sim/src/index.ts
@@ -15,6 +15,7 @@ import { setUiOption, uiSettings } from "./ui-settings";
 import { debugCli, debugHelper, debugState } from "./debug";
 import type { EventLogEntry } from "./event-log";
 import { formatEventLogLine } from "./event-log-format";
+import { fetchAxSnapshot } from "./ax-cli";
 import {
   parsePreviewPanes,
   parseSimulatorTheme,
@@ -718,6 +719,23 @@ function deviceLabelsForEvents(events: EventLogEntry[]): Map<string, string> {
   return labels;
 }
 
+async function axDump(deviceArg?: string) {
+  const udid = deviceArg ? resolveDevice(deviceArg) : undefined;
+  const state = readState(udid);
+  if (!state) {
+    console.error("No serve-sim server running. Run `serve-sim` first.");
+    process.exit(1);
+  }
+
+  try {
+    const snapshot = await fetchAxSnapshot(state.streamUrl);
+    console.log(JSON.stringify(snapshot));
+  } catch (err) {
+    console.error(`Failed to read accessibility tree: ${err instanceof Error ? err.message : String(err)}`);
+    process.exit(1);
+  }
+}
+
 async function gesture(jsonStr: string, deviceArg?: string) {
   const state = readState(deviceArg);
   if (!state) {
@@ -1867,6 +1885,12 @@ program
   .option("-n, --limit <count>", "Maximum number of events")
   .action((opts) => eventLog(opts.device, { json: opts.json, limit: opts.limit }));
 
+program
+  .command("ax")
+  .description("Dump the current accessibility tree as JSON")
+  .option(...deviceOpt)
+  .action((opts) => axDump(opts.device));
+
 // `camera` and `permissions` keep their own dedicated argument parsers (the
 // camera verb has nested sub-verbs and source flags; permissions has a
 // unit-tested parser module). Register them as passthrough commands so they