From d8f25cad0bc2fa793dda8ac1a090d0d732a662ea Mon Sep 17 00:00:00 2001 From: cryptotavares Date: Fri, 23 Oct 2026 22:57:13 +0100 Subject: [PATCH 1/3] feat: cancel timed-out tasks; stable refs that late writes cannot overwrite A request that timed out kept running: it could still click, snapshot and then overwrite the refMap the next request relied on. Refs were renumbered by every observation and discovery call, so refs an agent held could point at other elements. On mobile, a click observation could show the old screen. - RequestQueue passes an AbortSignal to each task and aborts it on timeout. ToolContext gains optional `signal` and `setRefMap`. Interaction tools, run_steps, discovery tools and the observation check the signal between steps (MM_REQUEST_ABORTED, internal). - New RefStore (server-owned, wraps sessionManager.setRefMap): each request takes a generation; writes from an older or aborted request are dropped. Tools that call context.sessionManager.setRefMap go through it too (proxy). ISessionManager is unchanged. - Stable refs: discovery tools and observations reuse the previous ref for the same element (browser key: role + name + testId; mobile: testId, else role + name; elements sharing a key are matched in screen order). ToolContext.setRefMap returns the stabilized { refMap, nodes }, which describe_screen / accessibility_snapshot return. New elements get refs above a session high-water mark, so a removed ref is never reused until launch/cleanup. When a tool set refs, the observation merges into them. - list_testids no longer reads the a11y tree or sets refs (it returns none). - Mobile observation and describe_screen: one device snapshot for both the tree and test IDs (optional IPlatformDriver.captureScreen; ScreenCapture type exported). After a mutating or batch call, re-snapshot every 300 ms while the screen is busy (", busy" label, progress/activity indicator), at most 3 times / 2 s; an unchanged screen gets one retry. - The observation timeout now aborts the observation instead of leaving it running. --- README.md | 15 + SKILL.md | 6 +- src/index.ts | 1 + src/platform/index.ts | 1 + src/platform/mobile-platform-driver.test.ts | 46 + src/platform/mobile-platform-driver.ts | 30 +- src/platform/types.ts | 22 + src/server/create-server.test.ts | 945 +++++++++++++++++++- src/server/create-server.ts | 391 ++++++-- src/server/ref-stability.test.ts | 451 ++++++++++ src/server/ref-stability.ts | 188 ++++ src/server/ref-store.test.ts | 429 +++++++++ src/server/ref-store.ts | 317 +++++++ src/server/request-queue.test.ts | 93 ++ src/server/request-queue.ts | 37 +- src/server/screen-settle.test.ts | 83 ++ src/server/screen-settle.ts | 56 ++ src/tools/batch.test.ts | 30 + src/tools/batch.ts | 10 +- src/tools/discovery-tools.test.ts | 189 +++- src/tools/discovery-tools.ts | 29 +- src/tools/driver-delegation.test.ts | 3 +- src/tools/interaction.test.ts | 41 + src/tools/interaction.ts | 6 + src/tools/types/errors.ts | 2 + src/tools/utils.ts | 44 + src/tools/utils/constants.ts | 23 + src/types/http.ts | 20 + 28 files changed, 3394 insertions(+), 114 deletions(-) create mode 100644 src/server/ref-stability.test.ts create mode 100644 src/server/ref-stability.ts create mode 100644 src/server/ref-store.test.ts create mode 100644 src/server/ref-store.ts create mode 100644 src/server/screen-settle.test.ts create mode 100644 src/server/screen-settle.ts diff --git a/README.md b/README.md index 5697626..a615ccc 100644 --- a/README.md +++ b/README.md @@ -366,9 +366,17 @@ type ToolContext = { get refMap(): Map; // lazy — returns empty map if no session workflowContext: WorkflowContext; toolRegistry: Map>; + driver?: IPlatformDriver; + signal?: AbortSignal; // aborted when the request times out + setRefMap?: ( + map: Map, + nodes?: A11yNodeTrimmed[], + ) => { refMap: Map; nodes?: A11yNodeTrimmed[] }; }; ``` +Tools check `signal?.aborted` between steps and stop. Tools that build refs call `context.setRefMap(map, nodes)` (fall back to `sessionManager.setRefMap` when it is absent) and return the `nodes` it returns: the daemon may renumber refs so elements the agent already holds keep their refs. The daemon's `RefStore` drops a write from a request that timed out or was replaced by a newer one, so a cancelled call cannot overwrite the refs of later calls. In the daemon, `context.sessionManager.setRefMap` goes through the same guard (without renumbering). + The daemon routes `POST /tool/:name` requests through the registry, applies Zod validation on inputs, executes the tool through the request queue, and captures observations (test IDs, a11y snapshot) according to the call's `observe` level (see [Observations](#observations)). **Registered tools:** @@ -443,6 +451,10 @@ mm type e3 "hello" → types into the element mapped to e3 This accessibility-first approach provides reliable element targeting that survives minor UI changes. +Refs stay valid while the element stays on screen. Each observation, `describe_screen` and `accessibility_snapshot` is renumbered against the last tree the agent received. Elements are matched by key (browser: role + name + testId; mobile: testId, else role + name). A key found once in each tree keeps its ref. Same key: matched in screen order (the k-th element with that key keeps the ref of the k-th element with that key in the last tree), so the many unnamed containers of a mobile screen keep their refs and an unchanged screen diffs as unchanged. When an element is inserted in the middle of a group that shares a key, the later elements of that group shift by one ref (each takes its predecessor's ref; the last gets a fresh one); elements with other keys are not affected. Other elements get fresh refs above the highest ref handed out so far, so an element that leaves the screen and comes back gets a fresh ref, and a removed ref is never reused for another element within a session (numbering restarts at `launch` and `cleanup`). Ref numbers can therefore grow within a session: new elements, options and collapsed runs get fresh refs, and do not reuse lower numbers. Options (`role: option`) and elements inside collapsed runs (`… 4 more`) are not stabilized: they always get fresh, contiguous refs so ranges (`e2–e56`) stay valid. If a tool set refs during a call (for example `describe_screen` with `observe`), the observation keeps them and adds its new refs, so every ref in the response resolves. `list_testids` does not change refs. + +On mobile, after a mutating or batch call, the daemon may take the snapshot again after 300 ms: while it shows a busy element (a label ending in `, busy`, an Android `ProgressBar`, or an iOS `ActivityIndicator` / `ProgressIndicator` / `UIActivityIndicatorView`), and at most once when it still equals the previous screen. In total at most 3 extra snapshots and about 2 s. The test IDs come from the same snapshot as the tree. `describe_screen` on mobile also takes one snapshot for both. + ### Session Metadata Consumers can persist launch metadata with `writeSessionMetadata(artifactsDir, metadata)`. It writes `//session.json` (for example `test-artifacts//session.json` with `artifactsDir: 'test-artifacts'`) and returns the file path. @@ -541,6 +553,8 @@ For `run_steps`, the top-level `observe` sets the level of the one observation c Observations contain `testIds` and `a11y` only. They do not include app state; call `get_state` for that. +A call that times out is cancelled: its signal is aborted, it stops at its next step, it collects no observation, and its refs are dropped. + ## CLI Reference The `mm` CLI provides a unified interface for agents and developers. All commands communicate with the daemon over HTTP — the daemon is auto-started on `mm launch` if not already running. @@ -795,6 +809,7 @@ Tool errors are classified into specific error codes for structured handling: | `MM_PORT_IN_USE` | Port already in use | | **System** | | | `MM_UNKNOWN_TOOL` | Unknown tool name | +| `MM_REQUEST_ABORTED` | Internal: a cancelled (timed-out) call stopped; the caller gets the timeout error | | `MM_INTERNAL_ERROR` | Internal server error | | `MM_BATCH_TIMEOUT` | `batchTimeoutMs` deadline exceeded in run_steps | | `MM_CDP_BLOCKED` | CDP method is blocked (destructive to session); on mobile the Hermes blocklist maps here too | diff --git a/SKILL.md b/SKILL.md index 94923e9..ca2b9d9 100644 --- a/SKILL.md +++ b/SKILL.md @@ -18,7 +18,7 @@ mm cleanup --shutdown # 5. Clean up when done **Critical rules:** -- **Always `describe-screen` before interacting.** Refs like `e1`, `e2` are ephemeral — they change after every action. +- **Always `describe-screen` before interacting.** Refs like `e1`, `e2` are short-lived: a ref stays valid while its element stays on screen. New elements get new refs, an element that leaves and comes back gets a new ref, and a removed ref is never reused for another element (until `launch`/`cleanup`). - **Always `describe-screen` after interacting** — OR use inline `observations` from mutating tool responses. Mutating tools (click, type, navigate, etc.) return an `observations` object with fresh `testIds` and `a11y` refs. You can use these refs directly for the next interaction without calling `describe-screen`. Call `describe-screen` when you need the full tree or screenshots, and `get-state` when you need app state (observations carry none). - **One target per command.** Specify exactly ONE of: a11y ref (`e5`), testId, or CSS selector. - **Errors are structured.** Check the `error.code` field to decide recovery strategy (see Error Codes below). @@ -85,7 +85,9 @@ After a mutating action, the response includes fresh screen data: } ``` -You can use the `ref` values from `observations.a11y.nodes` for the next interaction — no `describe-screen` needed. Note that refs in compacted observations may be summary nodes (e.g., `"55 options (refs e2–e56)"`) when there are 3+ options under a combobox or listbox. +You can use the `ref` values from `observations.a11y.nodes` for the next interaction — no `describe-screen` needed. Refs stay valid while the element stays on screen: an unchanged element keeps the ref you already hold (from `describe-screen` or an earlier observation), even when other elements appear before it. A new element gets a new, higher ref; an element whose name changed, or that left the screen and came back, gets a new ref, so ref numbers grow during a session. Elements with the same key (for example unnamed containers on mobile) are matched in screen order: if one is inserted among them, the later ones of that group shift by one ref. Note that refs in compacted observations may be summary nodes (e.g., `"55 options (refs e2–e56)"`) when there are 3+ options under a combobox or listbox. + +**Timeouts cancel the call:** when a call times out (`Tool execution timed out`), it is cancelled. It stops at its next step and cannot overwrite the refs of later calls. The action itself may still have happened, so run `describe-screen` before retrying. **Quick reference:** diff --git a/src/index.ts b/src/index.ts index 3699cbf..50eb544 100644 --- a/src/index.ts +++ b/src/index.ts @@ -10,6 +10,7 @@ export type { TypeActionResult, GetTextActionResult, PlatformScreenshotOptions, + ScreenCapture, } from './platform'; export { PlaywrightPlatformDriver, MobilePlatformDriver } from './platform'; diff --git a/src/platform/index.ts b/src/platform/index.ts index 6ffb045..dd96aed 100644 --- a/src/platform/index.ts +++ b/src/platform/index.ts @@ -7,6 +7,7 @@ export type { PlatformScreenshotOptions, WithinScope, IPlatformDriver, + ScreenCapture, } from './types.js'; export { PlaywrightPlatformDriver } from './playwright-driver.js'; diff --git a/src/platform/mobile-platform-driver.test.ts b/src/platform/mobile-platform-driver.test.ts index 5ca57d9..19f18d6 100644 --- a/src/platform/mobile-platform-driver.test.ts +++ b/src/platform/mobile-platform-driver.test.ts @@ -71,6 +71,52 @@ describe('MobilePlatformDriver', () => { vi.restoreAllMocks(); }); + describe('captureScreen', () => { + it('returns the tree and the test IDs of one snapshot', async () => { + const backend = createMockBackend({ + getWindowSize: vi.fn().mockResolvedValue({ width: 402, height: 874 }), + snapshot: vi + .fn() + .mockResolvedValueOnce({ + platform: 'ios', + hierarchy: [ + makeElement({ identifier: 'login', label: 'Unlock' }), + makeElement({ + identifier: 'below', + type: 'StaticText', + frame: { x: 0, y: 2000, width: 10, height: 10 }, + }), + ], + raw: '[]', + timestamp: Date.now(), + }) + .mockResolvedValueOnce({ + platform: 'ios', + hierarchy: [makeElement({ identifier: 'wallet' })], + raw: '[]', + timestamp: Date.now(), + }), + }); + const driver = new MobilePlatformDriver(backend); + + const screen = await driver.captureScreen(); + const testIds = await screen.getTestIds(); + + expect(backend.snapshot).toHaveBeenCalledTimes(1); + expect(screen.nodes.map((node) => node.testId)).toStrictEqual([ + 'login', + 'below', + ]); + expect(screen.refMap.get('e1')).toBe('identifier:login'); + // Same items, order and visibility as getTestIds on that snapshot. + expect(testIds).toStrictEqual([ + { testId: 'login', tag: 'Button', text: 'Unlock', visible: true }, + { testId: 'below', tag: 'StaticText', text: undefined, visible: false }, + ]); + expect(await screen.getTestIds(1)).toHaveLength(1); + }); + }); + describe('getPlatform', () => { it('returns ios for ios backend', () => { const driver = new MobilePlatformDriver(createMockBackend()); diff --git a/src/platform/mobile-platform-driver.ts b/src/platform/mobile-platform-driver.ts index 7aea944..4eee7aa 100644 --- a/src/platform/mobile-platform-driver.ts +++ b/src/platform/mobile-platform-driver.ts @@ -22,6 +22,7 @@ import type { GetTextActionResult, PlatformScreenshotOptions, PlatformType, + ScreenCapture, WithinScope, } from './types.js'; import type { @@ -208,6 +209,33 @@ export class MobilePlatformDriver implements IPlatformDriver { */ async getTestIds(limit?: number): Promise { const snapshot = await this.#backend.snapshot(); + return this.#testIdsOf(snapshot.hierarchy, limit); + } + + /** + * Takes one device snapshot and returns its accessibility tree and test + * IDs, so both describe the same screen. + * + * @returns The nodes, refMap and a test ID reader for the same snapshot. + */ + async captureScreen(): Promise { + const snapshot = await this.#backend.snapshot(); + return { + ...normalizeSnapshot(snapshot.hierarchy), + getTestIds: async (limit?: number) => + this.#testIdsOf(snapshot.hierarchy, limit), + }; + } + + /** + * @param hierarchy - The UIElement tree of a snapshot. + * @param limit - Maximum number of test IDs to return. + * @returns The test IDs, with visibility against the device viewport. + */ + async #testIdsOf( + hierarchy: UIElement[], + limit?: number, + ): Promise { const items: TestIdItem[] = []; const max = limit ?? OBSERVATION_TESTID_LIMIT; let viewport: { width: number; height: number } | undefined; @@ -217,7 +245,7 @@ export class MobilePlatformDriver implements IPlatformDriver { // Without a viewport we cannot judge visibility; report everything. viewport = undefined; } - collectTestIds(snapshot.hierarchy, items, max, viewport); + collectTestIds(hierarchy, items, max, viewport); return items; } diff --git a/src/platform/types.ts b/src/platform/types.ts index 529e6ed..caf9e28 100644 --- a/src/platform/types.ts +++ b/src/platform/types.ts @@ -48,6 +48,21 @@ export type WithinScope = { value: string; }; +/** + * One snapshot of the screen: the accessibility tree, and the test IDs of the + * same snapshot. + */ +export type ScreenCapture = { + nodes: A11yNodeTrimmed[]; + refMap: Map; + /** + * Lists the test IDs of this snapshot (no new snapshot). + * + * @param limit - Maximum number of test IDs to return. + */ + getTestIds(limit?: number): Promise; +}; + export type IPlatformDriver = { click( targetType: TargetType, @@ -88,6 +103,13 @@ export type IPlatformDriver = { getTestIds(limit?: number): Promise; + /** + * Takes one snapshot and returns both the accessibility tree and its test + * IDs, so they describe the same screen. Optional: without it, callers use + * `getAccessibilityTree` and `getTestIds` (two snapshots). + */ + captureScreen?(): Promise; + screenshot(options: PlatformScreenshotOptions): Promise; getAppState(): Promise; diff --git a/src/server/create-server.test.ts b/src/server/create-server.test.ts index 13e5a59..b950418 100644 --- a/src/server/create-server.test.ts +++ b/src/server/create-server.test.ts @@ -14,7 +14,13 @@ import { import { readDaemonState } from './daemon-state.js'; import pkg from '../../package.json'; import type { PortMap, WorkflowContext } from '../capabilities/context.js'; -import type { DaemonState, ServerConfig, ToolResponse } from '../types/http.js'; +import { toolRegistry } from '../tools/registry.js'; +import type { + DaemonState, + ServerConfig, + ToolFunction, + ToolResponse, +} from '../types/http.js'; const tmpDir = path.join(os.tmpdir(), `mm-create-server-test-${Date.now()}`); @@ -1846,8 +1852,10 @@ describe('observation compaction in HTTP responses', () => { await post('/tool/click', { a11yRef: 'e1', observe: 'full' }); const body = await post('/tool/click', { a11yRef: 'e1' }); + // "Cancel" left the screen and came back: it gets a fresh ref (e4), a + // ref is never reused for an element the agent did not see last. expect(body.observations?.a11y.diff).toStrictEqual({ - added: ['e2'], + added: ['e4'], removed: ['e3'], unchanged: 1, }); @@ -2150,3 +2158,936 @@ describe('createServer observation driver freshness after launch', () => { expect(body.error.code).toBe('MM_TOOL_NOT_SUPPORTED_ON_PLATFORM'); }); }); + +describe('createServer abort, ref generations and ref stability', () => { + type Node = { ref: string; role: string; name: string; path: string[] }; + type ObsBody = { + ok: boolean; + error?: { code: string; message: string }; + observations?: { + a11y: { + nodes: Node[]; + diff?: { added: string[]; removed: string[]; unchanged: number }; + }; + }; + }; + type Tree = { nodes: Node[]; refMap: Map }; + + const sleep = async (ms: number): Promise => + new Promise((resolve) => setTimeout(resolve, ms)); + + const buildTree = (names: string[]): Tree => ({ + nodes: names.map((name, index) => ({ + ref: `e${index + 1}`, + role: 'button', + name, + path: [], + })), + refMap: new Map( + names.map((name, index) => [`e${index + 1}`, `label:${name}`] as const), + ), + }); + + let server: ServerInstance | undefined; + let state: DaemonState; + let mockSM: ReturnType; + let storedRefMap: Map; + let driver: { + getPlatform: ReturnType; + getAppState: ReturnType; + getTestIds: ReturnType; + getAccessibilityTree: ReturnType; + click: ReturnType; + }; + const addedTools: string[] = []; + + const startServer = async ( + options: { requestTimeoutMs?: number; mobile?: boolean } = {}, + ): Promise => { + if (options.mobile !== false) { + mockSM.getPlatformDriver.mockReturnValue(driver); + } + server = createServer( + buildConfig({ + sessionManager: mockSM as unknown as ServerConfig['sessionManager'], + requestTimeoutMs: options.requestTimeoutMs, + }), + ); + state = await server.start(); + }; + + const post = async ( + urlPath: string, + requestBody: unknown = {}, + ): Promise<{ status: number; body: ObsBody }> => { + const res = await httpRequest(`http://127.0.0.1:${state.port}${urlPath}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(requestBody), + }); + return { status: res.status, body: (await res.json()) as ObsBody }; + }; + + const addTool = (name: string, fn: ToolFunction): void => { + toolRegistry.set(name, fn); + addedTools.push(name); + }; + + // A mutating tool without a `timeoutMs` default, so the server's + // requestTimeoutMs is the queue timeout (click defaults to 15 s). + const addTapTool = (): void => { + addTool('test_tap', async () => ({ ok: true, result: {} })); + }; + + beforeEach(async () => { + await fs.mkdir(tmpDir, { recursive: true }); + exitSpy = vi.spyOn(process, 'exit').mockImplementation((() => {}) as never); + + storedRefMap = new Map(); + mockSM = createMockSessionManager(); + mockSM.hasActiveSession.mockReturnValue(true); + mockSM.setRefMap.mockImplementation((map: Map) => { + storedRefMap = map; + }); + mockSM.getRefMap.mockImplementation(() => storedRefMap); + + driver = { + getPlatform: vi.fn().mockReturnValue('android'), + getAppState: vi.fn().mockResolvedValue({ isLoaded: true }), + getTestIds: vi.fn().mockResolvedValue([]), + getAccessibilityTree: vi.fn().mockResolvedValue(buildTree(['Unlock'])), + click: vi.fn().mockResolvedValue({ clicked: true, target: 'x' }), + }; + }); + + afterEach(async () => { + const running = server; + server = undefined; + await running?.stop(); + for (const name of addedTools.splice(0)) { + toolRegistry.delete(name); + } + exitSpy.mockRestore(); + const discovery = await import('../tools/utils/discovery.js'); + vi.mocked(discovery.collectTrimmedA11ySnapshot).mockReset(); + vi.mocked(discovery.collectTrimmedA11ySnapshot).mockResolvedValue({ + nodes: [], + refMap: new Map(), + }); + await fs.rm(tmpDir, { recursive: true, force: true }).catch(() => {}); + }); + + describe('P1-08: timed-out requests', () => { + it('aborts the signal of a timed-out task and ignores its late setRefMap (AC1)', async () => { + await startServer({ requestTimeoutMs: 50 }); + let abortedWhenLate: boolean | undefined; + let finished!: () => void; + const done = new Promise((resolve) => { + finished = resolve; + }); + addTool('test_slow_tool', async (_input, context) => { + await sleep(120); + abortedWhenLate = context.signal?.aborted; + context.setRefMap?.(new Map([['e1', 'late']])); + finished(); + return { ok: true, result: {} }; + }); + + const { status, body } = await post('/tool/test_slow_tool'); + await done; + + expect(status).toBe(500); + expect(body.error?.message).toContain('timed out after 50ms'); + expect(abortedWhenLate).toBe(true); + expect(mockSM.setRefMap).not.toHaveBeenCalled(); + }); + + it('keeps the refMap of the second request after the first times out (AC2)', async () => { + await startServer({ requestTimeoutMs: 50 }); + const firstTree = buildTree(['Old screen']); + const secondTree = buildTree(['New screen']); + // First click: its observation snapshot arrives after the timeout. + driver.getAccessibilityTree + .mockImplementationOnce(async () => { + await sleep(120); + return firstTree; + }) + .mockResolvedValueOnce(secondTree); + addTapTool(); + + const first = await post('/tool/test_tap'); + const second = await post('/tool/test_tap'); + + expect(first.status).toBe(500); + expect(second.status).toBe(200); + expect(second.body.observations?.a11y.nodes).toStrictEqual( + secondTree.nodes, + ); + expect(mockSM.setRefMap).toHaveBeenCalledTimes(1); + expect(storedRefMap).toStrictEqual(secondTree.refMap); + // The aborted observation stopped before its test ID step. + expect(driver.getTestIds).toHaveBeenCalledTimes(1); + }); + + it('stops the browser observation at the next checkpoint after an abort', async () => { + const discovery = await import('../tools/utils/discovery.js'); + vi.mocked(discovery.collectTestIds).mockClear(); + mockSM.getPage.mockReturnValue({ + waitForLoadState: vi.fn(async () => sleep(120)), + waitForFunction: vi.fn(async () => undefined), + } as never); + await startServer({ requestTimeoutMs: 50, mobile: false }); + addTapTool(); + + const { status } = await post('/tool/test_tap'); + await sleep(150); + + expect(status).toBe(500); + expect(vi.mocked(discovery.collectTestIds)).not.toHaveBeenCalled(); + expect(mockSM.setRefMap).not.toHaveBeenCalled(); + }); + + it('cancels a tool that checks the signal between steps', async () => { + await startServer({ requestTimeoutMs: 50 }); + const steps: string[] = []; + addTool('test_steps_tool', async (_input, context) => { + steps.push('one'); + await sleep(100); + if (context.signal?.aborted) { + return { ok: false, error: { code: 'X', message: 'aborted' } }; + } + steps.push('two'); + return { ok: true, result: {} }; + }); + + await post('/tool/test_steps_tool'); + await sleep(80); + + expect(steps).toStrictEqual(['one']); + }); + }); + + describe('P1-09: observations and refs', () => { + it('a non-delivered observation does not change lastObservation (AC1)', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['A', 'B'])) + .mockResolvedValueOnce(buildTree(['A', 'C'])) + .mockResolvedValueOnce(buildTree(['A', 'C'])); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/click', { testId: 'x', observe: 'none' }); + const { body } = await post('/tool/click', { testId: 'x' }); + + // observe none collected nothing; the diff is against the first + // (delivered) observation. + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + expect(body.observations?.a11y.diff).toStrictEqual({ + added: ['e3'], + removed: ['e2'], + unchanged: 1, + }); + }); + + it('keeps the refs the tool set and adds the new observation refs', async () => { + await startServer(); + const toolTree = buildTree(['Send', 'Swap']); + const observedTree = buildTree(['Banner', 'Send', 'Swap']); + driver.getAccessibilityTree + .mockResolvedValueOnce(toolTree) + .mockResolvedValueOnce(observedTree); + + const { body } = await post('/tool/describe_screen', { + observe: 'diff', + }); + + expect(body.ok).toBe(true); + // The observation is stabilized against the tool's tree, so Send/Swap + // keep the refs of the refMap the tool set. + const nodes = body.observations?.a11y.nodes ?? []; + expect(nodes.map((node) => [node.ref, node.name])).toStrictEqual([ + ['e3', 'Banner'], + ['e1', 'Send'], + ['e2', 'Swap'], + ]); + // The new element's ref resolves through the session manager too. + const resolved = mockSM.getRefMap(); + for (const node of nodes) { + expect(resolved.has(node.ref)).toBe(true); + } + expect(resolved.get('e3')).toBe('label:Banner'); + expect(resolved.get('e1')).toBe('label:Send'); + expect(resolved.get('e2')).toBe('label:Swap'); + // No settle: the tool set refs, so it already looked at the screen. + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + }); + + it('lets the run_steps observation replace refs a step set', async () => { + await startServer(); + const stepTree = buildTree(['Send']); + const observedTree = buildTree(['Done']); + driver.getAccessibilityTree + .mockResolvedValueOnce(stepTree) + .mockResolvedValueOnce(observedTree); + + await post('/tool/run_steps', { + steps: [{ tool: 'accessibility_snapshot' }], + }); + + expect(mockSM.setRefMap).toHaveBeenCalledTimes(2); + expect(storedRefMap.get('e2')).toBe('label:Done'); + expect(storedRefMap.has('e1')).toBe(false); + }); + + it('describe_screen then click keeps describe_screen refs for unchanged elements (AC2)', async () => { + const discovery = await import('../tools/utils/discovery.js'); + const described = { + nodes: [ + { ref: 'e1', role: 'button', name: 'Send', path: [] }, + { ref: 'e2', role: 'button', name: 'Swap', path: [] }, + ], + refMap: new Map([ + ['e1', 'role=button[name="Send"]'], + ['e2', 'role=button[name="Swap"]'], + ]), + }; + // A banner appears before the buttons: without stabilization it would + // take e1 and shift Send/Swap to e2/e3. + const afterClick = { + nodes: [ + { ref: 'e1', role: 'alert', name: 'Saved', path: [] }, + { ref: 'e2', role: 'button', name: 'Send', path: [] }, + { ref: 'e3', role: 'button', name: 'Swap', path: [] }, + ], + refMap: new Map([ + ['e1', 'role=alert[name="Saved"]'], + ['e2', 'role=button[name="Send"]'], + ['e3', 'role=button[name="Swap"]'], + ]), + }; + vi.mocked(discovery.collectTrimmedA11ySnapshot) + .mockResolvedValueOnce(described as never) + .mockResolvedValueOnce(afterClick as never); + await startServer({ mobile: false }); + + await post('/tool/describe_screen'); + const { body } = await post('/tool/click', { + a11yRef: 'e2', + observe: 'diff', + }); + + expect(body.observations?.a11y.nodes).toStrictEqual([ + { ref: 'e3', role: 'alert', name: 'Saved', path: [] }, + { ref: 'e1', role: 'button', name: 'Send', path: [] }, + { ref: 'e2', role: 'button', name: 'Swap', path: [] }, + ]); + expect(storedRefMap).toStrictEqual( + new Map([ + ['e3', 'role=alert[name="Saved"]'], + ['e1', 'role=button[name="Send"]'], + ['e2', 'role=button[name="Swap"]'], + ]), + ); + }); + + it('keeps refs stable across two observations and diffs on them', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Send', 'Swap'])) + .mockResolvedValueOnce(buildTree(['Banner', 'Send', 'Swap'])); + + await post('/tool/click', { testId: 'x' }); + const { body } = await post('/tool/click', { testId: 'x' }); + + expect(body.observations?.a11y.diff).toStrictEqual({ + added: ['e3'], + removed: [], + unchanged: 2, + }); + expect(storedRefMap.get('e1')).toBe('label:Send'); + expect(storedRefMap.get('e3')).toBe('label:Banner'); + }); + + it('diffs two mobile observations of the same tree with duplicate keys as unchanged', async () => { + await startServer(); + // Many unnamed containers share the key `role + ""` (Android proof: + // android.view.ViewGroup), plus a status-bar label that changes. + const screen = (signal: string): Tree => { + const nodes: Node[] = [ + { + ref: 'e1', + role: 'android.widget.FrameLayout', + name: signal, + path: [], + }, + ...Array.from( + { length: 30 }, + (_, index): Node => ({ + ref: `e${index + 2}`, + role: 'android.view.ViewGroup', + name: '', + path: ['android.widget.FrameLayout'], + }), + ), + { ref: 'e32', role: 'android.widget.Button', name: 'Send', path: [] }, + ]; + return { + nodes, + refMap: new Map([ + ['e1', `label:${signal}|type:android.widget.FrameLayout`], + ['e32', 'label:Send|type:android.widget.Button'], + ]), + }; + }; + driver.getAccessibilityTree + .mockResolvedValueOnce(screen('T-Mobile, no signal.')) + .mockResolvedValueOnce(screen('T-Mobile, no signal.')) + // Settle retry (unchanged): the status bar changed meanwhile. + .mockResolvedValue(screen('T-Mobile, signal full.')); + + await post('/tool/click', { testId: 'x' }); + const same = await post('/tool/click', { testId: 'x' }); + + expect(same.body.observations?.a11y.diff).toStrictEqual({ + added: ['e33'], + removed: ['e1'], + unchanged: 31, + }); + // Only the status-bar node whose name changed is sent. + expect( + same.body.observations?.a11y.nodes.map((node) => node.name), + ).toStrictEqual(['T-Mobile, signal full.']); + + const again = await post('/tool/click', { testId: 'x' }); + + expect(again.body.observations?.a11y.diff).toStrictEqual({ + added: [], + removed: [], + unchanged: 32, + }); + expect(again.body.observations?.a11y.nodes).toStrictEqual([]); + }); + + it('does not stabilize against refs set outside the RefStore', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Send'])) + .mockResolvedValueOnce(buildTree(['Banner', 'Send'])); + + await post('/tool/click', { testId: 'x' }); + storedRefMap = new Map([['e1', 'external']]); + await post('/tool/click', { testId: 'x' }); + + // No baseline match, and e1 was handed out before: fresh refs only. + expect(storedRefMap).toStrictEqual( + new Map([ + ['e2', 'label:Banner'], + ['e3', 'label:Send'], + ]), + ); + }); + + it('never reuses a removed ref across three observations', async () => { + await startServer(); + const names = Array.from({ length: 10 }, (_, index) => `B${index + 1}`); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(names)) + .mockResolvedValueOnce(buildTree(names.slice(0, 7))) + .mockResolvedValueOnce(buildTree([...names.slice(0, 7), 'New'])); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/click', { testId: 'x' }); + const { body } = await post('/tool/click', { testId: 'x' }); + + expect(body.observations?.a11y.diff?.added).toStrictEqual(['e11']); + expect(storedRefMap.get('e11')).toBe('label:New'); + expect(storedRefMap.has('e8')).toBe(false); + }); + + it('describe_screen refs stay stable against the previous tree, and later observations stay stable against them', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Send', 'Swap'])) + .mockResolvedValueOnce(buildTree(['Banner', 'Send', 'Swap'])) + .mockResolvedValueOnce(buildTree(['Banner', 'Send', 'Swap', 'New'])); + + await post('/tool/click', { testId: 'x' }); + const described = (await post('/tool/describe_screen')) + .body as unknown as { + result: { a11y: { nodes: Node[] } }; + }; + const { body } = await post('/tool/click', { testId: 'x' }); + + expect( + described.result.a11y.nodes.map((node) => [node.ref, node.name]), + ).toStrictEqual([ + ['e3', 'Banner'], + ['e1', 'Send'], + ['e2', 'Swap'], + ]); + expect( + body.observations?.a11y.nodes.map((node) => [node.ref, node.name]), + ).toStrictEqual([ + ['e3', 'Banner'], + ['e1', 'Send'], + ['e2', 'Swap'], + ['e4', 'New'], + ]); + // describe_screen resets lastObservation: no diff on the next call. + expect(body.observations?.a11y.diff).toBeUndefined(); + }); + + it('routes context.sessionManager.setRefMap through the request handle', async () => { + await startServer(); + const target = mockSM as unknown as Record; + target.whoAmI = function (this: unknown): boolean { + return this === mockSM; + }; + Object.defineProperty(mockSM, 'isSelf', { + configurable: true, + get(this: unknown): boolean { + return this === mockSM; + }, + }); + let seen: { bound?: boolean; getter?: boolean; id?: string } = {}; + addTool('test_consumer_tool', async (_input, context) => { + const sessionManager = context.sessionManager as unknown as { + whoAmI(): boolean; + isSelf: boolean; + } & ToolContext['sessionManager']; + const { whoAmI } = sessionManager; + seen = { + bound: whoAmI(), + getter: sessionManager.isSelf, + id: sessionManager.getSessionId(), + }; + context.sessionManager.setRefMap(new Map([['e1', 'consumer']])); + return { ok: true, result: {} }; + }); + + await post('/tool/test_consumer_tool', { observe: 'none' }); + + expect(seen).toStrictEqual({ + bound: true, + getter: true, + id: 'test-session', + }); + expect(mockSM.setRefMap).toHaveBeenCalledTimes(1); + expect(storedRefMap).toStrictEqual(new Map([['e1', 'consumer']])); + }); + + it('binds methods and getters of a session manager class with #private fields', async () => { + class PrivateSessionManager { + readonly #id = 'private-session'; + + #refMap = new Map(); + + get privateId(): string { + return this.#id; + } + + readPrivateId(): string { + return this.#id; + } + + hasActiveSession(): boolean { + return true; + } + + getSessionId(): string { + return this.#id; + } + + getPlatformDriver(): typeof driver { + return driver; + } + + setRefMap(map: Map): void { + this.#refMap = map; + } + + getRefMap(): Map { + return this.#refMap; + } + } + const real = new PrivateSessionManager(); + // Fill in the other members the server needs from the mock, as own + // properties of the real instance (the #private members stay real). + for (const [key, value] of Object.entries(mockSM)) { + if (!(key in real)) { + (real as unknown as Record)[key] = value; + } + } + server = createServer( + buildConfig({ + sessionManager: real as unknown as ServerConfig['sessionManager'], + }), + ); + state = await server.start(); + let seen: { getter?: string; method?: string; detached?: string } = {}; + addTool('test_private_tool', async (_input, context) => { + const proxied = context.sessionManager as unknown as Pick< + PrivateSessionManager, + 'privateId' | 'readPrivateId' + >; + const { readPrivateId } = proxied; + seen = { + getter: proxied.privateId, + method: proxied.readPrivateId(), + detached: readPrivateId(), + }; + context.sessionManager.setRefMap(new Map([['e1', 'private']])); + return { ok: true, result: {} }; + }); + + const { status } = await post('/tool/test_private_tool', { + observe: 'none', + }); + + expect(status).toBe(200); + expect(seen).toStrictEqual({ + getter: 'private-session', + method: 'private-session', + detached: 'private-session', + }); + // The write went through the RefStore to the real #refMap. + expect(real.getRefMap()).toStrictEqual(new Map([['e1', 'private']])); + }); + + it('drops a late context.sessionManager.setRefMap from a timed-out request', async () => { + await startServer({ requestTimeoutMs: 50 }); + let finished!: () => void; + const done = new Promise((resolve) => { + finished = resolve; + }); + addTool('test_slow_consumer_tool', async (_input, context) => { + await sleep(120); + context.sessionManager.setRefMap(new Map([['e1', 'late']])); + finished(); + return { ok: true, result: {} }; + }); + + const { status } = await post('/tool/test_slow_consumer_tool'); + await done; + + expect(status).toBe(500); + expect(mockSM.setRefMap).not.toHaveBeenCalled(); + }); + + describe('launch and cleanup', () => { + let realLaunch: ToolFunction | undefined; + + beforeEach(() => { + realLaunch = toolRegistry.get('launch'); + }); + + afterEach(() => { + if (realLaunch) { + toolRegistry.set('launch', realLaunch); + } + }); + + it('restarts ref numbering after a launch', async () => { + await startServer(); + toolRegistry.set('launch', async () => ({ ok: true, result: {} })); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['A', 'B', 'C'])) + .mockResolvedValueOnce(buildTree(['X'])); + + await post('/tool/click', { testId: 'x' }); + await post('/launch', { observe: 'none' }); + await post('/tool/click', { testId: 'x' }); + + expect(storedRefMap).toStrictEqual(new Map([['e1', 'label:X']])); + }); + + it('does not reset the baseline when a timed-out launch completes', async () => { + await startServer({ requestTimeoutMs: 50 }); + toolRegistry.set('launch', async () => { + await sleep(120); + return { ok: true, result: {} }; + }); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Send'])) + .mockResolvedValueOnce(buildTree(['Banner', 'Send'])); + addTapTool(); + + await post('/tool/test_tap'); + const launched = await post('/launch', { observe: 'none' }); + const { body } = await post('/tool/test_tap', { observe: 'full' }); + + expect(launched.status).toBe(500); + // The baseline survived: Send keeps e1. + expect( + body.observations?.a11y.nodes.map((node) => [node.ref, node.name]), + ).toStrictEqual([ + ['e2', 'Banner'], + ['e1', 'Send'], + ]); + }); + }); + }); + + describe('P1-09: mobile settle after a mutating tool', () => { + it('re-snapshots once after 300 ms when the screen looks unchanged', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValueOnce(buildTree(['Wallet'])); + + await post('/tool/click', { testId: 'x' }); + const startedAt = Date.now(); + const { body } = await post('/tool/click', { testId: 'unlock' }); + + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(290); + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(3); + expect(body.observations?.a11y.nodes.map((node) => node.name)).toContain( + 'Wallet', + ); + }); + + it('retries an unchanged, non-busy screen only once', async () => { + await startServer(); + driver.getAccessibilityTree.mockResolvedValue(buildTree(['Same'])); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/click', { testId: 'x' }); + + // 1 (first click) + 1 + 1 retry (second click). + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(3); + }); + + it('takes at most three extra snapshots while the screen is busy', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValue(buildTree(['Password', 'Unlock, busy'])); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/click', { testId: 'x' }); + + // 1 (first click) + 1 + 3 extra (second click). + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(5); + }); + + it('keeps retrying when an unchanged screen turns busy, within the limits', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValueOnce(buildTree(['Password', 'Unlock, busy'])) + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValue(buildTree(['Wallet'])); + + await post('/tool/click', { testId: 'x' }); + const { body } = await post('/tool/click', { + testId: 'x', + observe: 'full', + }); + + // Second click: unchanged → retry (the one unchanged retry), busy → + // retry, unchanged again → no second unchanged retry: stop. + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(4); + expect( + body.observations?.a11y.nodes.map((node) => node.name), + ).toStrictEqual(['Password', 'Unlock']); + }); + + it('re-snapshots while a node is busy, even when the screen changed (Android unlock)', async () => { + await startServer(); + const busy: Tree = { + nodes: [ + { + ref: 'e1', + role: 'android.widget.EditText', + name: 'Password', + path: [], + }, + { + ref: 'e2', + role: 'android.widget.Button', + name: 'Unlock, busy', + path: [], + }, + ], + refMap: new Map([['e2', 'identifier:log-in-button']]), + }; + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Password', 'Unlock'])) + .mockResolvedValueOnce(busy) + .mockResolvedValueOnce(busy) + .mockResolvedValueOnce(buildTree(['Wallet'])); + + await post('/tool/click', { testId: 'x' }); + const { body } = await post('/tool/click', { testId: 'log-in-button' }); + + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(4); + expect( + body.observations?.a11y.nodes.map((node) => node.name), + ).toStrictEqual(['Wallet']); + }); + + it('re-snapshots a busy first observation without a baseline', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Unlock, busy'])) + .mockResolvedValueOnce(buildTree(['Wallet'])); + + await post('/tool/click', { testId: 'x' }); + + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + }); + + it('stops settling at MOBILE_SETTLE_MAX_MS', async () => { + await startServer(); + driver.getAccessibilityTree.mockImplementation(async () => { + await sleep(700); + return buildTree(['Unlock, busy']); + }); + + await post('/tool/click', { testId: 'x' }); + + // 1st snapshot, then (300 ms + 700 ms) twice = 2000 ms; a third wait + // would end after the 2000 ms cap. + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(3); + }); + + it('does not settle when the tool set refs', async () => { + await startServer(); + driver.getAccessibilityTree.mockResolvedValue(buildTree(['Same'])); + addTool('test_ref_tool', async (_input, context) => { + const tree = buildTree(['Same']); + context.setRefMap?.(tree.refMap, tree.nodes); + return { ok: true, result: {} }; + }); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/test_ref_tool'); + + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + }); + + it('still settles when the tool set refs without nodes, and keeps those refs', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Same'])) + .mockResolvedValueOnce(buildTree(['Same'])) + .mockResolvedValueOnce(buildTree(['Same', 'New'])); + addTool('test_refs_only_tool', async (_input, context) => { + // Refs-only write through the sessionManager proxy: the tool did + // not look at the screen. + context.sessionManager.setRefMap(new Map([['e1', 'tool-target']])); + return { ok: true, result: {} }; + }); + + await post('/tool/click', { testId: 'x' }); + const { body } = await post('/tool/test_refs_only_tool', { + observe: 'full', + }); + + // Settled: unchanged → one retry. + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(3); + // The refs the tool set are kept; the observation refs are added. + expect(storedRefMap.get('e1')).toBe('tool-target'); + const nodes = body.observations?.a11y.nodes ?? []; + expect(nodes.map((node) => node.name)).toStrictEqual(['Same', 'New']); + for (const node of nodes) { + expect(node.ref).not.toBe('e1'); + expect(storedRefMap.get(node.ref)).toBe(`label:${node.name}`); + } + }); + + it('does not settle after cleanup', async () => { + await startServer(); + driver.getAccessibilityTree.mockResolvedValue( + buildTree(['Unlock, busy']), + ); + mockSM.cleanup.mockResolvedValue(true); + + await post('/cleanup', { observe: 'full' }); + + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(1); + }); + + it('takes the test IDs from the same snapshot as the final tree', async () => { + await startServer(); + const screens = [ + { names: ['Password', 'Unlock'], testIds: ['login'] }, + { names: ['Password', 'Unlock, busy'], testIds: ['login-busy'] }, + { names: ['Wallet'], testIds: ['wallet'] }, + ].map(({ names, testIds }) => ({ + ...buildTree(names), + getTestIds: vi.fn(async () => + testIds.map((testId) => ({ testId, tag: 'View', visible: true })), + ), + })); + const captureScreen = vi.fn(); + for (const screen of screens) { + captureScreen.mockResolvedValueOnce(screen); + } + Object.assign(driver, { captureScreen }); + + await post('/tool/click', { testId: 'x' }); + const { body } = (await post('/tool/click', { + testId: 'x', + })) as unknown as { + body: { + observations: { + testIds: { testId: string }[]; + a11y: { nodes: Node[] }; + }; + }; + }; + + expect(captureScreen).toHaveBeenCalledTimes(3); + expect(driver.getAccessibilityTree).not.toHaveBeenCalled(); + expect(driver.getTestIds).not.toHaveBeenCalled(); + expect(screens[1]?.getTestIds).not.toHaveBeenCalled(); + expect( + body.observations.testIds.map((item) => item.testId), + ).toStrictEqual(['wallet']); + expect( + body.observations.a11y.nodes.map((node) => node.name), + ).toStrictEqual(['Wallet']); + }); + + it('does not wait when the screen changed', async () => { + await startServer(); + driver.getAccessibilityTree + .mockResolvedValueOnce(buildTree(['Login'])) + .mockResolvedValueOnce(buildTree(['Wallet'])); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/click', { testId: 'x' }); + + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + }); + + it('does not wait for a readonly tool that asked for an observation', async () => { + await startServer(); + driver.getAccessibilityTree.mockResolvedValue(buildTree(['Same'])); + + await post('/tool/click', { testId: 'x' }); + await post('/tool/get_state', { observe: 'full' }); + + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + }); + + it('stops waiting and sets no refs when the request is aborted', async () => { + await startServer({ requestTimeoutMs: 150 }); + driver.getAccessibilityTree.mockResolvedValue(buildTree(['Same'])); + + addTapTool(); + + await post('/tool/test_tap'); + mockSM.setRefMap.mockClear(); + const { status } = await post('/tool/test_tap'); + await sleep(250); + + expect(status).toBe(500); + expect(driver.getAccessibilityTree).toHaveBeenCalledTimes(2); + expect(mockSM.setRefMap).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/src/server/create-server.ts b/src/server/create-server.ts index 2e792a9..a4581a2 100644 --- a/src/server/create-server.ts +++ b/src/server/create-server.ts @@ -3,9 +3,12 @@ import { execSync } from 'node:child_process'; import { randomUUID } from 'node:crypto'; import * as fs from 'node:fs/promises'; import * as http from 'node:http'; +import { setTimeout as delay } from 'node:timers/promises'; import { writeDaemonState, removeDaemonState } from './daemon-state.js'; import { compactObservation } from './observation-compaction.js'; +import { RefStore } from './ref-store.js'; +import type { RefHandle } from './ref-store.js'; import { RequestQueue } from './request-queue.js'; import { RequestTimer, @@ -13,15 +16,19 @@ import { setRequestStart, withMeta, } from './response-timing.js'; +import { isScreenBusy, screenSignature } from './screen-settle.js'; +import type { ISessionManager } from './session-manager.js'; import pkg from '../../package.json'; import type { PortMap, WorkflowContext } from '../capabilities/context.js'; import { PlaywrightPlatformDriver } from '../platform/playwright-driver.js'; +import type { IPlatformDriver, ScreenCapture } from '../platform/types.js'; import { toolRegistry, getToolCategory, checkPlatformGate, } from '../tools/registry.js'; import type { ToolCategory } from '../tools/registry.js'; +import type { A11yNodeTrimmed, TestIdItem } from '../tools/types/discovery.js'; import { createDefaultObservation } from '../tools/types/observation.js'; import type { StepRecordObservation } from '../tools/types/observation.js'; import type { RunStepsInput } from '../tools/types/tool-inputs.js'; @@ -30,6 +37,10 @@ import { ANIMATION_SETTLE_TIMEOUT_MS, GRACEFUL_SHUTDOWN_TIMEOUT_MS, MAX_IDLE_CHECK_INTERVAL_MS, + MOBILE_SETTLE_MAX_MS, + MOBILE_SETTLE_MAX_SNAPSHOTS, + MOBILE_SETTLE_MS, + MOBILE_SETTLE_UNCHANGED_RETRIES, OBSERVATION_TESTID_LIMIT, OBSERVATION_TIMEOUT_MS, } from '../tools/utils/constants.js'; @@ -57,6 +68,40 @@ export type ServerInstance = { stop(): Promise; }; +type ObservedTree = { + testIds: TestIdItem[]; + nodes: A11yNodeTrimmed[]; + refMap: Map; +}; + +/** + * Wraps the session manager so a tool's `sessionManager.setRefMap` goes + * through the request's RefStore handle (dropped when the request is stale + * or aborted). Every other member is read from the real session manager; + * methods are bound to it, so `this` (and private fields) work. + * + * @param sessionManager - The real session manager. + * @param refs - The request's RefStore handle. + * @returns The wrapped session manager. + */ +function routeSetRefMap( + sessionManager: ISessionManager, + refs: RefHandle, +): ISessionManager { + const setRefMap = (map: Map): void => { + refs.set(map); + }; + return new Proxy(sessionManager, { + get(target, property): unknown { + if (property === 'setRefMap') { + return setRefMap; + } + const value: unknown = Reflect.get(target, property, target); + return typeof value === 'function' ? value.bind(target) : value; + }, + }); +} + /** * Merges a tool result with observation data into the HTTP response body. * @@ -162,6 +207,12 @@ export function resolveObserveLevel( export function createServer(config: ServerConfig): ServerInstance { const app = express(); const queue = new RequestQueue(config.requestTimeoutMs); + const refStore = new RefStore(config.sessionManager, () => { + const platform = config.sessionManager.getPlatformDriver?.()?.getPlatform(); + return platform === undefined || platform === 'browser' + ? 'browser' + : 'mobile'; + }); const nonce = randomUUID(); let httpServer: http.Server | null = null; @@ -218,11 +269,22 @@ export function createServer(config: ServerConfig): ServerInstance { /** * Builds a lazy ToolContext where `page` and `refMap` are only accessed * when a tool actually reads them, avoiding throws for non-session tools. + * `refMap` reads the session manager on each access; writes from older or + * aborted requests never reach it (see RefStore), so it always holds the + * refs of this request or of the last request that completed. Both + * `setRefMap` and `sessionManager.setRefMap` write through the request's + * RefStore handle. * * @param wfCtx - The current workflow context to embed in the tool context. + * @param signal - The request's abort signal. + * @param refs - The request's RefStore handle. * @returns A ToolContext with lazy page and refMap accessors. */ - function buildToolContext(wfCtx: WorkflowContext): ToolContext { + function buildToolContext( + wfCtx: WorkflowContext, + signal: AbortSignal, + refs: RefHandle, + ): ToolContext { const driver = config.sessionManager.hasActiveSession() ? (config.sessionManager.getPlatformDriver?.() ?? new PlaywrightPlatformDriver( @@ -232,7 +294,7 @@ export function createServer(config: ServerConfig): ServerInstance { : undefined; return { - sessionManager: config.sessionManager, + sessionManager: routeSetRefMap(config.sessionManager, refs), get page(): ReturnType { return config.sessionManager.getPage(); }, @@ -244,9 +306,237 @@ export function createServer(config: ServerConfig): ServerInstance { workflowContext: wfCtx, toolRegistry, driver, + signal, + setRefMap: (map, nodes) => refs.set(map, nodes), }; } + /** + * Takes one mobile snapshot: one device snapshot when the driver has + * `captureScreen`, else the tree and the test IDs separately. + * + * @param driver - The mobile platform driver. + * @returns The tree and a test ID reader. + */ + async function captureMobile( + driver: IPlatformDriver, + ): Promise { + if (driver.captureScreen) { + return driver.captureScreen(); + } + const tree = await driver.getAccessibilityTree(); + return { ...tree, getTestIds: async (limit) => driver.getTestIds(limit) }; + } + + /** + * Takes the mobile accessibility tree and test IDs. + * + * After a mutating or batch tool (`settle`), the screen may not have + * changed yet, or may show work in progress (for example a "busy" button + * while the wallet unlocks). Limits: + * + * - Busy (see `isScreenBusy`): wait `MOBILE_SETTLE_MS` and take the + * snapshot again, while the screen stays busy. + * - Unchanged (equals the last screen the agent received, ignoring refs) + * and not busy: retry at most `MOBILE_SETTLE_UNCHANGED_RETRIES` time(s). + * Many taps leave the screen as it is (for example tapping the tab that is + * already open), so more retries would only add device snapshots. On + * Android the status bar (clock, signal) usually differs between + * snapshots, so this rule rarely fires there. + * - In total: at most `MOBILE_SETTLE_MAX_SNAPSHOTS` extra snapshots, and no + * wait that would end later than `MOBILE_SETTLE_MAX_MS` after the first + * snapshot. A spinner that never stops costs at most these limits. + * + * The test IDs come from the final snapshot, so they describe the same + * screen as the tree. + * + * @param driver - The mobile platform driver. + * @param settle - Whether to settle (see above). + * @param signal - Aborted when the request or the observation times out. + * @returns The test IDs, nodes and refMap. + */ + async function snapshotMobile( + driver: IPlatformDriver, + settle: boolean, + signal: AbortSignal, + ): Promise { + let screen = await captureMobile(driver); + if (settle) { + const settleStartedAt = Date.now(); + const previous = + refStore.baselineSignature() ?? + (lastObservation ? screenSignature(lastObservation.a11y.nodes) : null); + let unchangedRetries = 0; + const shouldRetry = (nodes: A11yNodeTrimmed[]): boolean => { + if (isScreenBusy(nodes)) { + return true; + } + if ( + previous === null || + unchangedRetries >= MOBILE_SETTLE_UNCHANGED_RETRIES || + screenSignature(nodes) !== previous + ) { + return false; + } + unchangedRetries += 1; + return true; + }; + for ( + let extra = 0; + extra < MOBILE_SETTLE_MAX_SNAPSHOTS && + Date.now() - settleStartedAt + MOBILE_SETTLE_MS <= + MOBILE_SETTLE_MAX_MS && + shouldRetry(screen.nodes); + extra++ + ) { + await delay(MOBILE_SETTLE_MS, undefined, { signal }); + screen = await captureMobile(driver); + } + } + signal.throwIfAborted(); + const testIds = await screen.getTestIds(OBSERVATION_TESTID_LIMIT); + return { testIds, nodes: screen.nodes, refMap: screen.refMap }; + } + + /** + * Takes the browser test IDs and accessibility tree. Tools that may change + * the screen (mutating, batch, and a cdp call that asked for an + * observation) wait for the page to settle first. + * + * @param category - The tool category. + * @param signal - Aborted when the request or the observation times out. + * @returns The test IDs, nodes and refMap. + */ + async function snapshotBrowser( + category: ToolCategory, + signal: AbortSignal, + ): Promise { + const page = config.sessionManager.getPage(); + if (category !== 'readonly' && category !== 'discovery') { + await page.waitForLoadState('domcontentloaded').catch(() => undefined); + signal.throwIfAborted(); + await page + .waitForFunction( + async () => + new Promise((resolve) => { + requestAnimationFrame(() => { + const allSettled = document + .getAnimations() + .every((a: Animation) => a.playState !== 'running'); + resolve(allSettled); + }); + }), + { timeout: ANIMATION_SETTLE_TIMEOUT_MS }, + ) + .catch(() => undefined); + signal.throwIfAborted(); + } + const testIds = await collectTestIds(page, OBSERVATION_TESTID_LIMIT); + signal.throwIfAborted(); + const tree = await collectTrimmedA11ySnapshot(page); + return { testIds, ...tree }; + } + + /** + * Collects the post-tool observation. Its refMap is applied through the + * request's RefStore handle, which renumbers refs so unchanged elements + * keep the refs the agent already holds, drops the write when the request + * is stale or aborted, and keeps refs the tool set during this request. + * + * @param context - The tool context of the request. + * @param toolName - The tool name. + * @param category - The tool category. + * @param refs - The request's RefStore handle. + * @param signal - Aborted when the request or the observation times out. + * @returns The observation. + */ + async function collectObservation( + context: ToolContext, + toolName: string, + category: ToolCategory, + refs: RefHandle, + signal: AbortSignal, + ): Promise { + // Re-read the driver from the session manager rather than using + // context.driver, which was captured before the tool ran. The launch + // tool creates the session (and sets the platform driver) during + // execution, so the context snapshot would still be undefined then. + const driver = + config.sessionManager.getPlatformDriver?.() ?? context.driver; + const mobile = driver !== undefined && driver.getPlatform() !== 'browser'; + // When the tool set refs, the baseline is the tool's tree, so elements + // that did not change keep the tool's refs, and the observation's other + // refs are added to the tool's map. run_steps: a later step may have + // changed the screen after a step set refs, so its observation replaces + // them. + const keepToolRefs = category !== 'batch'; + // No settle when the tool set refs with nodes and keeps them: it already + // looked at the screen it returned. A refs-only write (no nodes) does + // not count: the tool did not read the screen. + const settle = + (category === 'mutating' || category === 'batch') && + toolName !== 'launch' && + toolName !== 'cleanup' && + !(refs.toolSawScreen && keepToolRefs); + const tree = mobile + ? await snapshotMobile(driver, settle, signal) + : await snapshotBrowser(category, signal); + signal.throwIfAborted(); + + const { nodes } = refs.setObservation(tree.refMap, tree.nodes, { + keepToolRefs, + }); + return createDefaultObservation(tree.testIds, nodes); + } + + /** + * Runs collectObservation bounded by `OBSERVATION_TIMEOUT_MS`. On timeout + * or request abort the collection is aborted, so it stops at the next + * checkpoint and cannot set refs later. + * + * @param context - The tool context of the request. + * @param toolName - The tool name. + * @param category - The tool category. + * @param refs - The request's RefStore handle. + * @param requestSignal - The request's abort signal. + * @returns The observation, or undefined on failure or timeout. + */ + async function observeWithTimeout( + context: ToolContext, + toolName: string, + category: ToolCategory, + refs: RefHandle, + requestSignal: AbortSignal, + ): Promise { + const controller = new AbortController(); + const forwardAbort = (): void => controller.abort(requestSignal.reason); + requestSignal.addEventListener('abort', forwardAbort, { once: true }); + let timeoutId: ReturnType | undefined; + try { + return await Promise.race([ + collectObservation( + context, + toolName, + category, + refs, + controller.signal, + ), + new Promise((resolve) => { + timeoutId = setTimeout(() => { + controller.abort(new Error('Observation timed out')); + resolve(undefined); + }, OBSERVATION_TIMEOUT_MS); + }), + ]); + } catch { + // non-fatal: observation failure must not block the tool response + return undefined; + } finally { + clearTimeout(timeoutId); + requestSignal.removeEventListener('abort', forwardAbort); + } + } + /** * Shared tool executor — validates input, runs through the queue, * and collects observations. @@ -350,83 +640,40 @@ export function createServer(config: ServerConfig): ServerInstance { try { timer.markEnqueued(); const { toolResult, observeLevel, observations } = await queue.enqueue( - async () => { + async (signal) => { timer.markTaskStart(); - const context = buildToolContext(currentWorkflowContext); + const refs = refStore.begin(signal); + const context = buildToolContext( + currentWorkflowContext, + signal, + refs, + ); const result = await tool(omitObserve(validatedInput), context); timer.markToolEnd(); + if ( + (toolName === 'launch' || toolName === 'cleanup') && + refs.isCurrent() + ) { + // A new or closed session: earlier refs are no baseline, and ref + // numbering starts again. + refStore.resetBaseline(); + } const level = resolveObserveLevel(category, inputRecord, result); let obs: StepRecordObservation | undefined; - if (level !== 'none' && config.sessionManager.hasActiveSession()) { + if ( + level !== 'none' && + !signal.aborted && + config.sessionManager.hasActiveSession() + ) { timer.markObservationStart(); - try { - obs = await Promise.race([ - (async (): Promise => { - // Re-read the driver from the session manager rather than - // using context.driver, which was captured before the tool - // ran. The launch tool creates the session (and sets the - // platform driver) during execution, so the context snapshot - // would still be undefined at that point. - const currentDriver = - config.sessionManager.getPlatformDriver?.() ?? - context.driver; - - if ( - currentDriver && - currentDriver.getPlatform() !== 'browser' - ) { - const testIds = await currentDriver.getTestIds( - OBSERVATION_TESTID_LIMIT, - ); - const { nodes, refMap: newRefMap } = - await currentDriver.getAccessibilityTree(); - config.sessionManager.setRefMap(newRefMap); - return createDefaultObservation(testIds, nodes); - } - - const page = config.sessionManager.getPage(); - - // Tools that may change the screen (mutating, batch, and a - // cdp call that asked for an observation) wait for the page - // to settle first. - if (category !== 'readonly' && category !== 'discovery') { - await page - .waitForLoadState('domcontentloaded') - .catch(() => undefined); - await page - .waitForFunction( - async () => - new Promise((resolve) => { - requestAnimationFrame(() => { - const allSettled = document - .getAnimations() - .every( - (a: Animation) => a.playState !== 'running', - ); - resolve(allSettled); - }); - }), - { timeout: ANIMATION_SETTLE_TIMEOUT_MS }, - ) - .catch(() => undefined); - } - const testIds = await collectTestIds( - page, - OBSERVATION_TESTID_LIMIT, - ); - const { nodes, refMap: newRefMap } = - await collectTrimmedA11ySnapshot(page); - config.sessionManager.setRefMap(newRefMap); - return createDefaultObservation(testIds, nodes); - })(), - new Promise((resolve) => - setTimeout(() => resolve(undefined), OBSERVATION_TIMEOUT_MS), - ), - ]); - } catch { - // non-fatal: observation failure must not block the tool response - } + obs = await observeWithTimeout( + context, + toolName, + category, + refs, + signal, + ); timer.markObservationEnd(); } diff --git a/src/server/ref-stability.test.ts b/src/server/ref-stability.test.ts new file mode 100644 index 0000000..433d066 --- /dev/null +++ b/src/server/ref-stability.test.ts @@ -0,0 +1,451 @@ +import { describe, it, expect } from 'vitest'; + +import { maxRefIndex, refKey, stabilizeRefs } from './ref-stability.js'; +import type { RefTree } from './ref-stability.js'; +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; + +type Spec = [ref: string, role: string, name: string, testId?: string]; + +function tree(specs: Spec[]): RefTree { + const nodes = specs.map( + ([ref, role, name, testId]): A11yNodeTrimmed => ({ + ref, + role, + name, + path: [], + ...(testId ? { testId } : {}), + }), + ); + return { + nodes, + refMap: new Map( + specs.map(([ref, role, name]) => [ref, `${role}:${name}`] as const), + ), + }; +} + +const refsOf = (result: RefTree): string[] => + result.nodes.map((node) => node.ref); + +const targetOf = (result: RefTree, name: string): string | undefined => { + const found = result.nodes.find((node) => node.name === name); + return found ? result.refMap.get(found.ref) : undefined; +}; + +describe('refKey', () => { + const button: A11yNodeTrimmed = { + ref: 'e1', + role: 'button', + name: 'Go', + path: [], + }; + + it('uses role, name and testId in browser mode', () => { + expect(refKey(button, 'browser')).not.toBe( + refKey({ ...button, testId: 'go' }, 'browser'), + ); + expect(refKey(button, 'browser')).not.toBe( + refKey({ ...button, name: 'Stop' }, 'browser'), + ); + }); + + it('uses the testId alone in mobile mode when present', () => { + expect(refKey({ ...button, testId: 'go' }, 'mobile')).toBe( + refKey({ ...button, name: 'Renamed', testId: 'go' }, 'mobile'), + ); + }); + + it('uses role and name in mobile mode without a testId', () => { + expect(refKey(button, 'mobile')).toBe( + refKey({ ...button, ref: 'e9' }, 'mobile'), + ); + expect(refKey(button, 'mobile')).not.toBe( + refKey({ ...button, role: 'text' }, 'mobile'), + ); + }); +}); + +describe('stabilizeRefs', () => { + it('keeps refs when an element is inserted before existing ones', () => { + const previous = tree([ + ['e1', 'button', 'Send'], + ['e2', 'button', 'Swap'], + ]); + const next = tree([ + ['e1', 'dialog', 'Notice'], + ['e2', 'button', 'Send'], + ['e3', 'button', 'Swap'], + ]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e3', 'e1', 'e2']); + expect(result.refMap).toStrictEqual( + new Map([ + ['e3', 'dialog:Notice'], + ['e1', 'button:Send'], + ['e2', 'button:Swap'], + ]), + ); + }); + + it('keeps refs of remaining elements and never reuses a removed ref', () => { + const previous = tree([ + ['e1', 'button', 'Send'], + ['e2', 'button', 'Swap'], + ['e3', 'button', 'Bridge'], + ]); + const next = tree([ + ['e1', 'button', 'Send'], + ['e2', 'button', 'Bridge'], + ['e3', 'button', 'Buy'], + ]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e1', 'e3', 'e4']); + expect(targetOf(result, 'Bridge')).toBe('button:Bridge'); + expect(targetOf(result, 'Buy')).toBe('button:Buy'); + expect(result.refMap.has('e2')).toBe(false); + }); + + it('gives a renamed element a fresh ref in browser mode', () => { + const previous = tree([ + ['e1', 'button', 'Unlock'], + ['e2', 'button', 'Help'], + ]); + const next = tree([ + ['e1', 'button', 'Unlocking…'], + ['e2', 'button', 'Help'], + ]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e3', 'e2']); + }); + + it('keeps the ref of a renamed element with the same testId in mobile mode', () => { + const previous = tree([['e1', 'Button', 'Unlock', 'unlock-button']]); + const next = tree([ + ['e1', 'Text', 'Welcome'], + ['e2', 'Button', 'Unlocking…', 'unlock-button'], + ]); + + const result = stabilizeRefs(previous, next, 'mobile'); + + expect(refsOf(result)).toStrictEqual(['e2', 'e1']); + expect(targetOf(result, 'Unlocking…')).toBe('Button:Unlocking…'); + }); + + it('matches a key that appears more often in the new tree in screen order; extra ones get fresh refs', () => { + const previous = tree([ + ['e1', 'button', 'Copy'], + ['e2', 'button', 'Done'], + ]); + const next = tree([ + ['e1', 'button', 'Copy'], + ['e2', 'button', 'Copy'], + ['e3', 'button', 'Done'], + ]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e1', 'e3', 'e2']); + }); + + it('matches a key that appeared more often in the previous tree in screen order', () => { + const previous = tree([ + ['e1', 'button', 'Copy'], + ['e2', 'button', 'Copy'], + ]); + const next = tree([['e1', 'button', 'Copy']]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e1']); + }); + + describe('same key: matched in screen order', () => { + const containers = (count: number, start = 1): Spec[] => + Array.from( + { length: count }, + (_, index): Spec => [`e${start + index}`, 'android.view.ViewGroup', ''], + ); + + it('keeps the refs of 20 identical unnamed containers when a unique node is inserted', () => { + const previous = tree(containers(20)); + const next = tree([ + ['e1', 'android.widget.TextView', 'Banner'], + ...containers(20, 2), + ]); + + const result = stabilizeRefs(previous, next, 'mobile'); + + expect(refsOf(result)).toStrictEqual([ + 'e21', + ...refsOf(previous).slice(0, 20), + ]); + expect(targetOf(result, 'Banner')).toBe('android.widget.TextView:Banner'); + }); + + it('shifts only the tail of a duplicate run when a container is inserted in its middle', () => { + // previous: Header, 5 containers (e2–e6), Footer. + const previous = tree([ + ['e1', 'text', 'Header'], + ...containers(5, 2), + ['e7', 'text', 'Footer'], + ]); + // A container is inserted after the 2nd one. + const next = tree([ + ['e1', 'text', 'Header'], + ...containers(6, 2), + ['e8', 'text', 'Footer'], + ]); + + const result = stabilizeRefs(previous, next, 'mobile'); + + // The first two containers and the unique nodes keep their refs. From + // the inserted one on, each container takes its predecessor's ref and + // the last one gets a fresh ref: acceptable for unnamed containers. + expect(refsOf(result)).toStrictEqual([ + 'e1', + 'e2', + 'e3', + 'e4', + 'e5', + 'e6', + 'e8', + 'e7', + ]); + }); + + it('gives a fresh ref only to a node whose name changed (status bar clock)', () => { + const previous = tree([ + ['e1', 'android.widget.TextView', '1:01 AM', 'clock'], + ['e2', 'android.widget.FrameLayout', 'T-Mobile, no signal.'], + ...containers(4, 3), + ]); + const next = tree([ + ['e1', 'android.widget.TextView', '1:02 AM', 'clock'], + ['e2', 'android.widget.FrameLayout', 'T-Mobile, signal full.'], + ...containers(4, 3), + ]); + + const result = stabilizeRefs(previous, next, 'mobile'); + + // The clock keeps its ref (same testId); the unnamed signal label is + // keyed by role + name, so it gets a fresh ref. The containers keep + // theirs. + expect(refsOf(result)).toStrictEqual([ + 'e1', + 'e7', + 'e3', + 'e4', + 'e5', + 'e6', + ]); + }); + + it('never produces duplicate refs with mixed duplicate and unique keys', () => { + const previous = tree([ + ['e1', 'button', 'Copy'], + ['e2', 'button', 'Copy'], + ['e3', 'button', 'Copy'], + ['e4', 'button', 'Done'], + ]); + const next = tree([ + ['e1', 'button', 'Done'], + ['e2', 'button', 'Copy'], + ['e3', 'button', 'New'], + ['e4', 'button', 'Copy'], + ['e5', 'button', 'Copy'], + ['e6', 'button', 'Copy'], + ]); + + const result = stabilizeRefs(previous, next, 'browser', 10); + const refs = refsOf(result); + + expect(refs).toStrictEqual(['e4', 'e1', 'e11', 'e2', 'e3', 'e12']); + expect(new Set(refs).size).toBe(refs.length); + expect([...result.refMap.keys()]).toStrictEqual(refs); + expect(result.refMap.get('e12')).toBe('button:Copy'); + }); + }); + + it('never produces duplicate refs', () => { + const previous = tree([ + ['e1', 'button', 'A'], + ['e2', 'button', 'B'], + ['e3', 'button', 'C'], + ]); + const next = tree([ + ['e1', 'button', 'C'], + ['e2', 'button', 'X'], + ['e3', 'button', 'A'], + ['e4', 'button', 'Y'], + ['e5', 'button', 'B'], + ]); + + const result = stabilizeRefs(previous, next, 'browser'); + const refs = refsOf(result); + + expect(new Set(refs).size).toBe(refs.length); + expect(new Set(result.refMap.keys()).size).toBe(result.refMap.size); + expect(refs).toStrictEqual(['e3', 'e4', 'e1', 'e5', 'e2']); + }); + + it('rewrites collapsed runs consistently and gives their members fresh refs', () => { + const previous = tree([['e1', 'button', 'Send']]); + const next: RefTree = { + nodes: [ + { ref: 'e1', role: 'img', name: 'icon', path: [] }, + { + ref: 'e1\u2013e3', + role: 'img', + name: '\u2026 2 more "icon" (refs e1\u2013e3)', + path: [], + }, + { ref: 'e4', role: 'button', name: 'Send', path: [] }, + ], + refMap: new Map([ + ['e1', 'img1'], + ['e2', 'img2'], + ['e3', 'img3'], + ['e4', 'button:Send'], + ]), + }; + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e2', 'e2\u2013e4', 'e1']); + expect(result.nodes[1].name).toBe('\u2026 2 more "icon" (refs e2\u2013e4)'); + expect(result.refMap).toStrictEqual( + new Map([ + ['e2', 'img1'], + ['e3', 'img2'], + ['e4', 'img3'], + ['e1', 'button:Send'], + ]), + ); + }); + + it('does not match previous nodes inside a collapsed run', () => { + const previous: RefTree = { + nodes: [ + { ref: 'e1', role: 'img', name: 'icon', path: [] }, + { ref: 'e1\u2013e3', role: 'img', name: 'more', path: [] }, + ], + refMap: new Map([ + ['e1', 'a'], + ['e2', 'b'], + ['e3', 'c'], + ]), + }; + const next = tree([['e1', 'img', 'icon']]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e4']); + }); + + it('gives fresh refs to refMap entries without a node', () => { + const previous = tree([['e1', 'button', 'Send']]); + const next: RefTree = { + nodes: [{ ref: 'e1', role: 'button', name: 'Send', path: [] }], + refMap: new Map([ + ['e1', 'button:Send'], + ['e2', 'hidden'], + ]), + }; + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(result.refMap).toStrictEqual( + new Map([ + ['e1', 'button:Send'], + ['e2', 'hidden'], + ]), + ); + }); + + it('leaves refs that are not eN untouched', () => { + const previous = tree([['e1', 'button', 'Send']]); + const next: RefTree = { + nodes: [{ ref: 'custom', role: 'button', name: 'Send', path: [] }], + refMap: new Map([['custom', 'x']]), + }; + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['custom']); + expect(result.refMap).toStrictEqual(new Map([['custom', 'x']])); + }); + + it('does not change the inputs', () => { + const previous = tree([['e1', 'button', 'Send']]); + const next = tree([ + ['e1', 'button', 'New'], + ['e2', 'button', 'Send'], + ]); + + stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(next)).toStrictEqual(['e1', 'e2']); + expect([...next.refMap.keys()]).toStrictEqual(['e1', 'e2']); + }); + + it('gives options fresh contiguous refs so option ranges stay valid', () => { + const previous = tree([ + ['e1', 'combobox', 'Language'], + ['e2', 'option', 'English'], + ['e3', 'option', 'French'], + ]); + const next = tree([ + ['e1', 'combobox', 'Language'], + ['e2', 'option', 'German'], + ['e3', 'option', 'English'], + ['e4', 'option', 'French'], + ]); + + const result = stabilizeRefs(previous, next, 'browser'); + + expect(refsOf(result)).toStrictEqual(['e1', 'e4', 'e5', 'e6']); + }); + + it('starts fresh refs above the floor', () => { + const previous = tree([['e1', 'button', 'Send']]); + const next = tree([ + ['e1', 'button', 'Send'], + ['e2', 'button', 'New'], + ]); + + const result = stabilizeRefs(previous, next, 'browser', 10); + + expect(refsOf(result)).toStrictEqual(['e1', 'e11']); + }); + + it('ignores a floor below the previous maximum', () => { + const previous = tree([['e5', 'button', 'Send']]); + const next = tree([['e1', 'button', 'New']]); + + expect(refsOf(stabilizeRefs(previous, next, 'browser', 2))).toStrictEqual([ + 'e6', + ]); + }); +}); + +describe('maxRefIndex', () => { + it('reads nodes, ranges and map keys', () => { + expect( + maxRefIndex({ + nodes: [ + { ref: 'e2', role: 'img', name: '', path: [] }, + { ref: 'e3\u2013e7', role: 'img', name: '', path: [] }, + ], + refMap: new Map([['e4', 'x']]), + }), + ).toBe(7); + expect(maxRefIndex({ nodes: [], refMap: new Map([['e9', 'x']]) })).toBe(9); + expect(maxRefIndex({ nodes: [], refMap: new Map() })).toBe(0); + }); +}); diff --git a/src/server/ref-stability.ts b/src/server/ref-stability.ts new file mode 100644 index 0000000..ac8369c --- /dev/null +++ b/src/server/ref-stability.ts @@ -0,0 +1,188 @@ +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; + +/** + * How elements are matched between two trees. + * + * - `browser`: role + name + testId. + * - `mobile`: the stable identifier (testId) when present, else role + name. + */ +export type RefKeyMode = 'browser' | 'mobile'; + +export type RefTree = { + nodes: A11yNodeTrimmed[]; + refMap: Map; +}; + +const SIMPLE_REF = /^e(?\d+)$/u; +const RANGE_REF = /^e(?\d+)\u2013e(?\d+)$/u; + +/** + * Builds the match key of a node. + * + * @param node - The node. + * @param mode - The key mode. + * @returns The key. + */ +export function refKey(node: A11yNodeTrimmed, mode: RefKeyMode): string { + if (mode === 'mobile') { + return node.testId + ? `id\u0000${node.testId}` + : `rn\u0000${node.role}\u0000${node.name}`; + } + return `${node.role}\u0000${node.name}\u0000${node.testId ?? ''}`; +} + +/** + * Lists the refs covered by collapsed-run summary nodes (`e3–e6`). + * + * @param nodes - The nodes. + * @returns The covered refs, including the first node of each run. + */ +function rangeMembers(nodes: A11yNodeTrimmed[]): Set { + const members = new Set(); + for (const node of nodes) { + const groups = RANGE_REF.exec(node.ref)?.groups; + if (groups) { + for (let i = Number(groups.start); i <= Number(groups.end); i++) { + members.add(`e${i}`); + } + } + } + return members; +} + +/** + * Groups the refs of the matchable nodes by key, in tree (screen) order. Not + * matchable: refs inside collapsed runs, and `option` nodes (observations + * summarize options as a `first–last` ref range, which must stay + * contiguous). + * + * @param nodes - The nodes. + * @param mode - The key mode. + * @returns The refs per key, each list in tree order. + */ +function refsByKey( + nodes: A11yNodeTrimmed[], + mode: RefKeyMode, +): Map { + const members = rangeMembers(nodes); + const byKey = new Map(); + for (const node of nodes) { + if ( + node.role !== 'option' && + SIMPLE_REF.test(node.ref) && + !members.has(node.ref) + ) { + const key = refKey(node, mode); + const refs = byKey.get(key); + if (refs) { + refs.push(node.ref); + } else { + byKey.set(key, [node.ref]); + } + } + } + return byKey; +} + +/** + * Returns the highest `eN` index used by a tree (nodes, ranges and map keys). + * + * @param tree - The tree. + * @returns The highest index, or 0. + */ +export function maxRefIndex(tree: RefTree): number { + let max = 0; + const refs = [...tree.nodes.map((node) => node.ref), ...tree.refMap.keys()]; + for (const ref of refs) { + const groups = SIMPLE_REF.exec(ref)?.groups ?? RANGE_REF.exec(ref)?.groups; + for (const value of Object.values(groups ?? {})) { + max = Math.max(max, Number(value)); + } + } + return max; +} + +/** + * Renumbers the refs of a new tree so an element that was in the previous + * tree keeps its previous ref. + * + * Elements are matched by key (see `RefKeyMode`). A key found once in each + * tree matches that element. When a key is found more than once (for + * example many unnamed containers on mobile), elements that share it are + * matched in screen order: the k-th element with key K in the new tree gets + * the ref of the k-th element with key K in the previous tree, for k up to + * the smaller of the two counts. So an element inserted in the middle of + * such a group shifts the refs of the later elements of that group by one + * (each gets the ref of its predecessor, and the last one a fresh ref); + * elements before it, and elements with other keys, keep their refs. + * + * Every other element (new, an extra occurrence of a key, an option, or + * inside a collapsed run) gets a fresh ref above both the highest ref of the + * previous tree and `floor`. Each previous ref is used at most once, so the + * result has no duplicates. Pass the highest ref ever handed out as `floor` + * so a ref that was removed earlier is never reused for a different element. + * Node refs, collapsed-run ranges and the refMap are rewritten together. + * + * @param previous - The tree the agent last received. + * @param next - The new tree. + * @param mode - How to match elements. + * @param floor - Fresh refs start above this index (default 0). + * @returns The new tree with stabilized refs (new objects; inputs unchanged). + */ +export function stabilizeRefs( + previous: RefTree, + next: RefTree, + mode: RefKeyMode, + floor = 0, +): RefTree { + const previousKeys = refsByKey(previous.nodes, mode); + const remap = new Map(); + for (const [key, refs] of refsByKey(next.nodes, mode)) { + const previousRefs = previousKeys.get(key) ?? []; + // Same key: matched in screen order. A unique key is the case k = 0. + const matched = Math.min(refs.length, previousRefs.length); + for (let index = 0; index < matched; index++) { + remap.set(refs[index], previousRefs[index]); + } + } + + let counter = Math.max(floor, maxRefIndex(previous)); + const assignFresh = (ref: string): void => { + if (!remap.has(ref)) { + counter += 1; + remap.set(ref, `e${counter}`); + } + }; + + // Tree order, so collapsed runs (first node, then its summary) stay + // contiguous. + for (const node of next.nodes) { + const groups = RANGE_REF.exec(node.ref)?.groups; + if (groups) { + for (let i = Number(groups.start); i <= Number(groups.end); i++) { + assignFresh(`e${i}`); + } + } else if (SIMPLE_REF.test(node.ref)) { + assignFresh(node.ref); + } + } + for (const ref of next.refMap.keys()) { + if (SIMPLE_REF.test(ref)) { + assignFresh(ref); + } + } + + const nodes = next.nodes.map((node): A11yNodeTrimmed => { + const groups = RANGE_REF.exec(node.ref)?.groups; + if (!groups) { + return { ...node, ref: remap.get(node.ref) ?? node.ref }; + } + const ref = `${remap.get(`e${groups.start}`)}\u2013${remap.get(`e${groups.end}`)}`; + return { ...node, ref, name: node.name.split(node.ref).join(ref) }; + }); + const refMap = new Map( + [...next.refMap].map(([ref, target]) => [remap.get(ref) ?? ref, target]), + ); + return { nodes, refMap }; +} diff --git a/src/server/ref-store.test.ts b/src/server/ref-store.test.ts new file mode 100644 index 0000000..9c1472c --- /dev/null +++ b/src/server/ref-store.test.ts @@ -0,0 +1,429 @@ +import { describe, it, expect, vi } from 'vitest'; + +import { RefStore } from './ref-store.js'; +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; + +function createSessionManager() { + let stored = new Map(); + return { + setRefMap: vi.fn((map: Map) => { + stored = map; + }), + getRefMap: vi.fn(() => stored), + }; +} + +const node = (ref: string, name: string): A11yNodeTrimmed => ({ + ref, + role: 'button', + name, + path: [], +}); + +describe('RefStore', () => { + it('increments the generation on each begin()', () => { + const store = new RefStore(createSessionManager()); + + const first = store.begin(); + const second = store.begin(); + + expect(first.generation).toBe(1); + expect(second.generation).toBe(2); + expect(store.generation).toBe(2); + }); + + it('forwards set() from the current handle to the session manager', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + const map = new Map([['e1', 'role=button[name="Go"]']]); + + const handle = store.begin(); + + expect(handle.isCurrent()).toBe(true); + expect(handle.set(map)).toStrictEqual({ applied: true, refMap: map }); + expect(handle.toolSetRefs).toBe(true); + expect(sessionManager.setRefMap).toHaveBeenCalledWith(map); + }); + + it('drops set() from a handle of an older generation', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + + const stale = store.begin(); + store.begin(); + + expect(stale.isCurrent()).toBe(false); + expect(stale.set(new Map([['e1', 'old']])).applied).toBe(false); + expect(stale.setObservation(new Map([['e1', 'old']]), []).applied).toBe( + false, + ); + expect(stale.toolSetRefs).toBe(false); + expect(sessionManager.setRefMap).not.toHaveBeenCalled(); + }); + + it('drops set() after the request signal aborts', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + const controller = new AbortController(); + + const handle = store.begin(controller.signal); + controller.abort(); + + expect(handle.isCurrent()).toBe(false); + expect(handle.set(new Map([['e1', 'late']])).applied).toBe(false); + expect(sessionManager.setRefMap).not.toHaveBeenCalled(); + }); + + it('applies the observation map when the tool did not set refs', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + const map = new Map([['e1', 'obs']]); + + const handle = store.begin(); + + expect(handle.setObservation(map, [node('e1', 'Go')]).applied).toBe(true); + expect(handle.toolSetRefs).toBe(false); + expect(sessionManager.setRefMap).toHaveBeenCalledWith(map); + }); + + it('keeps the refs the tool set and adds the other observation refs', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + const toolMap = new Map([ + ['e1', 'tool-send'], + ['e2', 'tool-swap'], + ]); + + const handle = store.begin(); + handle.set(toolMap, [node('e1', 'Send'), node('e2', 'Swap')]); + // Banner is new; Send and Swap match the tool's tree. + const written = handle.setObservation( + new Map([ + ['e1', 'obs-banner'], + ['e2', 'obs-send'], + ['e3', 'obs-swap'], + ]), + [node('e1', 'Banner'), node('e2', 'Send'), node('e3', 'Swap')], + ); + + expect(written.applied).toBe(true); + expect(written.nodes.map((each) => each.ref)).toStrictEqual([ + 'e3', + 'e1', + 'e2', + ]); + // Tool entries win on conflict; the new ref is added. + expect(sessionManager.getRefMap()).toStrictEqual( + new Map([ + ['e1', 'tool-send'], + ['e2', 'tool-swap'], + ['e3', 'obs-banner'], + ]), + ); + // Every ref of the observation resolves. + for (const each of written.nodes) { + expect(sessionManager.getRefMap().has(each.ref)).toBe(true); + } + // The baseline is the observation nodes with the merged map. + expect(store.baseline()).toStrictEqual({ + nodes: written.nodes, + refMap: written.refMap, + }); + }); + + describe('toolSetRefs and toolSawScreen', () => { + it('sets only toolSetRefs for a refs-only write', () => { + const handle = new RefStore(createSessionManager()).begin(); + + handle.set(new Map([['e1', 'tool']])); + + expect(handle.toolSetRefs).toBe(true); + expect(handle.toolSawScreen).toBe(false); + }); + + it('sets both for a write with nodes, and keeps toolSawScreen after a later refs-only write', () => { + const handle = new RefStore(createSessionManager()).begin(); + + handle.set(new Map([['e1', 'go']]), [node('e1', 'Go')]); + expect(handle.toolSawScreen).toBe(true); + + handle.set(new Map([['e1', 'tool']])); + expect(handle.toolSetRefs).toBe(true); + expect(handle.toolSawScreen).toBe(true); + }); + + it('sets neither when the write is dropped', () => { + const store = new RefStore(createSessionManager()); + const stale = store.begin(); + store.begin(); + + stale.set(new Map([['e1', 'go']]), [node('e1', 'Go')]); + + expect(stale.toolSetRefs).toBe(false); + expect(stale.toolSawScreen).toBe(false); + }); + }); + + it('lets the tool map win when an observation ref collides with it', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + const handle = store.begin(); + // A tool map without nodes: no baseline, so the observation is not + // stabilized and its e1 collides with the tool's e1. + handle.set(new Map([['e1', 'tool']])); + + handle.setObservation( + new Map([ + ['e1', 'obs'], + ['e2', 'obs2'], + ]), + [node('e1', 'A'), node('e2', 'B')], + ); + + expect(sessionManager.getRefMap().get('e1')).toBe('tool'); + }); + + it('applies the observation map over tool refs when keepToolRefs is false', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + const observationMap = new Map([['e1', 'obs']]); + + const handle = store.begin(); + handle.set(new Map([['e1', 'tool']])); + + const written = handle.setObservation(observationMap, [], { + keepToolRefs: false, + }); + + expect(written.applied).toBe(true); + // Not merged with the tool map. (e1 was used by the tool write, so the + // observation's entry gets a fresh ref.) + expect(sessionManager.getRefMap()).toStrictEqual(new Map([['e2', 'obs']])); + }); + + describe('baseline', () => { + it('is null before any write', () => { + expect(new RefStore(createSessionManager()).baseline()).toBeNull(); + }); + + it('returns the nodes and map of the last write', () => { + const store = new RefStore(createSessionManager()); + const nodes = [node('e1', 'Go')]; + const map = new Map([['e1', 'go']]); + + store.begin().set(map, nodes); + + expect(store.baseline()).toStrictEqual({ nodes, refMap: map }); + }); + + it('is null when the last write had no nodes', () => { + const store = new RefStore(createSessionManager()); + + store.begin().set(new Map([['e1', 'go']])); + + expect(store.baseline()).toBeNull(); + }); + + it('is null when the session manager map was replaced elsewhere', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + + store.begin().set(new Map([['e1', 'go']]), [node('e1', 'Go')]); + sessionManager.setRefMap(new Map([['e1', 'other']])); + + expect(store.baseline()).toBeNull(); + }); + + it('is null when the session manager map has a different size', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + + store.begin().set(new Map([['e1', 'go']]), [node('e1', 'Go')]); + sessionManager.setRefMap(new Map()); + + expect(store.baseline()).toBeNull(); + }); + + it('is null when getRefMap throws (no session)', () => { + const sessionManager = createSessionManager(); + const store = new RefStore(sessionManager); + + store.begin().set(new Map([['e1', 'go']]), [node('e1', 'Go')]); + sessionManager.getRefMap.mockImplementation(() => { + throw new Error('No active session'); + }); + + expect(store.baseline()).toBeNull(); + }); + + it('is cleared by resetBaseline()', () => { + const store = new RefStore(createSessionManager()); + + store.begin().set(new Map([['e1', 'go']]), [node('e1', 'Go')]); + store.resetBaseline(); + + expect(store.baseline()).toBeNull(); + }); + + it('is not moved by a dropped write', () => { + const store = new RefStore(createSessionManager()); + const nodes = [node('e1', 'Go')]; + const map = new Map([['e1', 'go']]); + + const stale = store.begin(); + store.begin().set(map, nodes); + stale.set(new Map([['e1', 'late']]), [node('e1', 'Late')]); + + expect(store.baseline()).toStrictEqual({ nodes, refMap: map }); + }); + }); + + describe('stable refs and the high-water mark', () => { + const tree = ( + names: string[], + ): { nodes: A11yNodeTrimmed[]; map: Map } => ({ + nodes: names.map((name, index) => node(`e${index + 1}`, name)), + map: new Map(names.map((name, index) => [`e${index + 1}`, name])), + }); + const names = (count: number, prefix = 'B'): string[] => + Array.from({ length: count }, (_, index) => `${prefix}${index + 1}`); + + it('never reuses a removed ref across three observations', () => { + const store = new RefStore(createSessionManager()); + + // 1: e1–e10. + const first = tree(names(10)); + store.begin().setObservation(first.map, first.nodes); + // 2: B8–B10 (e8–e10) are gone. + const second = tree(names(7)); + const afterDrop = store.begin().setObservation(second.map, second.nodes); + // 3: a new element appears after B7. + const third = tree([...names(7), 'New']); + const written = store.begin().setObservation(third.map, third.nodes); + + expect(afterDrop.nodes.map((each) => each.ref)).toStrictEqual( + names(7).map((_, index) => `e${index + 1}`), + ); + const fresh = written.nodes.find((each) => each.name === 'New'); + expect(fresh?.ref).toBe('e11'); + expect(written.refMap.get('e11')).toBe('New'); + expect(written.refMap.has('e8')).toBe(false); + expect(store.highWater).toBe(11); + }); + + it('bumps the high-water mark on tool writes too', () => { + const store = new RefStore(createSessionManager()); + + store.begin().set(new Map([['e40', 'x']])); + + expect(store.highWater).toBe(40); + // No baseline (no nodes), but fresh refs still start above 40. + const written = store + .begin() + .setObservation(new Map([['e1', 'a']]), [node('e1', 'A')]); + expect(written.nodes[0]?.ref).toBe('e41'); + }); + + it('does not move the high-water mark on a dropped write', () => { + const store = new RefStore(createSessionManager()); + const stale = store.begin(); + store.begin(); + + stale.set(new Map([['e9', 'late']])); + + expect(store.highWater).toBe(0); + }); + + it('resetBaseline() restarts ref numbering', () => { + const store = new RefStore(createSessionManager()); + const first = tree(names(5)); + store.begin().set(first.map, first.nodes); + + store.resetBaseline(); + const written = store + .begin() + .setObservation(new Map([['e1', 'a']]), [node('e1', 'A')]); + + expect(store.highWater).toBe(1); + expect(written.nodes[0]?.ref).toBe('e1'); + }); + + it('stabilizes tool writes with nodes against the baseline', () => { + const store = new RefStore(createSessionManager()); + const first = tree(['Send', 'Swap']); + store.begin().setObservation(first.map, first.nodes); + + const second = tree(['Banner', 'Send', 'Swap']); + const written = store.begin().set(second.map, second.nodes); + + expect(written.nodes?.map((each) => [each.ref, each.name])).toStrictEqual( + [ + ['e3', 'Banner'], + ['e1', 'Send'], + ['e2', 'Swap'], + ], + ); + expect(written.refMap).toStrictEqual( + new Map([ + ['e3', 'Banner'], + ['e1', 'Send'], + ['e2', 'Swap'], + ]), + ); + }); + + it('returns a tool write without nodes unchanged', () => { + const store = new RefStore(createSessionManager()); + const first = tree(['Send']); + store.begin().setObservation(first.map, first.nodes); + const map = new Map([['e1', 'other']]); + + const written = store.begin().set(map); + + expect(written).toStrictEqual({ applied: true, refMap: map }); + expect(store.baseline()).toBeNull(); + }); + + it('uses the key mode the store was created with', () => { + const store = new RefStore(createSessionManager(), () => 'mobile'); + store + .begin() + .setObservation(new Map([['e1', 'id']]), [ + { ...node('e1', 'Unlock'), testId: 'unlock' }, + ]); + + // Mobile: same testId, new name → same ref. + const written = store + .begin() + .setObservation(new Map([['e1', 'id']]), [ + { ...node('e1', 'Unlock, busy'), testId: 'unlock' }, + ]); + + expect(written.nodes[0]?.ref).toBe('e1'); + }); + }); + + describe('baselineSignature', () => { + it('is null without a valid baseline', () => { + expect( + new RefStore(createSessionManager()).baselineSignature(), + ).toBeNull(); + }); + + it('is computed once per baseline and ignores refs', () => { + const store = new RefStore(createSessionManager()); + const nodes = [node('e1', 'Go')]; + store.begin().set(new Map([['e1', 'go']]), nodes); + + const first = store.baselineSignature(); + // A cached value: changing the stored nodes in place does not change + // it. + nodes.push(node('e2', 'Later')); + + expect(first).toBe( + JSON.stringify([{ role: 'button', name: 'Go', path: [] }]), + ); + expect(store.baselineSignature()).toBe(first); + }); + }); +}); diff --git a/src/server/ref-store.ts b/src/server/ref-store.ts new file mode 100644 index 0000000..e8cb771 --- /dev/null +++ b/src/server/ref-store.ts @@ -0,0 +1,317 @@ +import { maxRefIndex, stabilizeRefs } from './ref-stability.js'; +import type { RefKeyMode, RefTree } from './ref-stability.js'; +import { screenSignature } from './screen-settle.js'; +import type { ISessionManager } from './session-manager.js'; +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; +import { debugWarn } from '../utils/logger.js'; + +/** + * The result of a write through a `RefHandle`. + */ +export type RefWrite = { + /** True when the map reached the session manager. */ + applied: boolean; + /** + * The map that was written (or would have been, when dropped). Refs may + * differ from the input: they are stabilized against the baseline. + */ + refMap: Map; + /** The nodes with the same refs as `refMap`, when nodes were passed. */ + nodes?: A11yNodeTrimmed[]; +}; + +/** + * A request's write access to the refMap. Get one from `RefStore.begin()`. + */ +export type RefHandle = { + /** The generation this handle was created with. */ + readonly generation: number; + /** + * True once a tool set refs through this handle during the request (with + * or without nodes). The observation then keeps those refs (see + * `setObservation`). + */ + readonly toolSetRefs: boolean; + /** + * True once a tool set refs with nodes during the request: the tool read + * the screen it returned, so the observation does not need to wait for + * the screen to settle. A refs-only write (for example through + * `sessionManager.setRefMap`) does not count. + */ + readonly toolSawScreen: boolean; + /** + * @returns True while no newer request has started and the request is not + * aborted. + */ + isCurrent(): boolean; + /** + * Sets refs from a tool. Dropped when the handle is not current. When + * nodes are passed, refs are first stabilized against the baseline, so + * elements the agent already holds keep their refs. + * + * @param map - The new ref map. + * @param nodes - The nodes the map was built from, when known. + * @returns The written map and nodes (with stabilized refs). + */ + set(map: Map, nodes?: A11yNodeTrimmed[]): RefWrite; + /** + * Sets refs from the post-tool observation. Refs are stabilized against + * the baseline. Dropped when the handle is not current. When the tool + * already set refs during this request (and `keepToolRefs` is not false), + * the result is the tool's map plus the observation entries whose ref the + * tool map does not have, so every ref of the observation resolves and the + * tool's refs stay valid. + * + * @param map - The new ref map. + * @param nodes - The observation nodes. + * @param options - Options. + * @param options.keepToolRefs - Keep refs the tool set (default true). + * `run_steps` passes false: a later step may have changed the screen after + * a step set refs. + * @returns The written map and the observation nodes (stabilized refs). + */ + setObservation( + map: Map, + nodes: A11yNodeTrimmed[], + options?: { keepToolRefs?: boolean }, + ): RefWrite & { nodes: A11yNodeTrimmed[] }; +}; + +/** + * The refs the agent last received, used to keep refs stable across + * observations. + */ +export type RefBaseline = { + nodes: A11yNodeTrimmed[]; + refMap: Map; +}; + +type StoredBaseline = RefBaseline & { signature?: string }; + +const EMPTY_TREE: RefTree = { nodes: [], refMap: new Map() }; + +/** + * Checks whether two maps have the same entries. + * + * @param left - First map. + * @param right - Second map. + * @returns True when both maps hold the same key/value pairs. + */ +function sameEntries( + left: Map, + right: Map, +): boolean { + if (left.size !== right.size) { + return false; + } + for (const [key, value] of left) { + if (right.get(key) !== value) { + return false; + } + } + return true; +} + +/** + * Guards `sessionManager.setRefMap` with a request generation and keeps refs + * stable. + * + * Each request calls `begin()` when it starts. A write from an older request + * (for example one that timed out and kept running) or from an aborted + * request is dropped before it reaches the session manager, so it cannot + * replace the refs of the request that runs now. + * + * Writes with nodes are stabilized against the last applied tree (the + * baseline): an element that is still on screen keeps its ref. New refs + * start above the highest ref ever applied in the session, so a removed ref + * is never reused for another element until `resetBaseline()`. + */ +export class RefStore { + readonly #sessionManager: Pick; + + readonly #keyMode: () => RefKeyMode; + + #generation = 0; + + #baseline: StoredBaseline | null = null; + + #highWater = 0; + + /** + * @param sessionManager - The session manager that owns the refMap. + * @param keyMode - Returns how elements are matched (browser or mobile) + * for the current session. Defaults to browser. + */ + constructor( + sessionManager: Pick, + keyMode: () => RefKeyMode = () => 'browser', + ) { + this.#sessionManager = sessionManager; + this.#keyMode = keyMode; + } + + /** + * @returns The generation of the most recent request. + */ + get generation(): number { + return this.#generation; + } + + /** + * @returns The highest ref index applied since the last reset. + */ + get highWater(): number { + return this.#highWater; + } + + /** + * Starts a new generation. Handles from earlier generations stop writing. + * + * @param signal - The request's abort signal. + * @returns The handle for this request. + */ + begin(signal?: AbortSignal): RefHandle { + this.#generation += 1; + const generation = this.#generation; + let toolMap: Map | null = null; + let toolSawScreen = false; + + const isCurrent = (): boolean => + generation === this.#generation && signal?.aborted !== true; + + const apply = ( + source: string, + map: Map, + nodes: A11yNodeTrimmed[] | undefined, + ): boolean => { + if (!isCurrent()) { + debugWarn( + 'ref-store', + `dropped ${source} refMap from generation ${generation} (current ${this.#generation}, aborted ${String(signal?.aborted === true)})`, + ); + return false; + } + this.#sessionManager.setRefMap(map); + this.#baseline = nodes ? { nodes, refMap: new Map(map) } : null; + this.#highWater = Math.max( + this.#highWater, + maxRefIndex({ nodes: nodes ?? [], refMap: map }), + ); + return true; + }; + + return { + generation, + get toolSetRefs(): boolean { + return toolMap !== null; + }, + get toolSawScreen(): boolean { + return toolSawScreen; + }, + isCurrent, + set: (map, nodes): RefWrite => { + const tree = nodes ? this.#stabilize(map, nodes) : undefined; + const refMap = tree?.refMap ?? map; + const applied = apply('tool', refMap, tree?.nodes); + if (applied) { + toolMap = refMap; + toolSawScreen ||= nodes !== undefined; + } + return { applied, refMap, ...(tree ? { nodes: tree.nodes } : {}) }; + }, + setObservation: ( + map, + nodes, + options, + ): RefWrite & { nodes: A11yNodeTrimmed[] } => { + const tree = this.#stabilize(map, nodes); + let { refMap } = tree; + if (toolMap && options?.keepToolRefs !== false) { + const merged = new Map(toolMap); + for (const [ref, target] of tree.refMap) { + if (!merged.has(ref)) { + merged.set(ref, target); + } + } + refMap = merged; + } + const applied = apply('observation', refMap, tree.nodes); + return { applied, refMap, nodes: tree.nodes }; + }, + }; + } + + /** + * Forgets the baseline and the highest ref (for example after a launch or + * cleanup, which start or end a session). + */ + resetBaseline(): void { + this.#baseline = null; + this.#highWater = 0; + } + + /** + * Returns the nodes and map of the last write, if the session manager still + * holds that map. Returns null when something else (for example code + * calling the real `sessionManager.setRefMap` directly) replaced it. + * + * @returns The baseline, or null. + */ + baseline(): RefBaseline | null { + return this.#validBaseline(); + } + + /** + * Returns the screen signature of the baseline nodes (see + * `screenSignature`). Computed once per baseline and cached. + * + * @returns The signature, or null when there is no valid baseline. + */ + baselineSignature(): string | null { + const baseline = this.#validBaseline(); + if (!baseline) { + return null; + } + baseline.signature ??= screenSignature(baseline.nodes); + return baseline.signature; + } + + /** + * @returns The stored baseline when the session manager still holds its + * map, else null. + */ + #validBaseline(): StoredBaseline | null { + if (!this.#baseline) { + return null; + } + let current: Map; + try { + current = this.#sessionManager.getRefMap(); + } catch { + return null; + } + return sameEntries(this.#baseline.refMap, current) ? this.#baseline : null; + } + + /** + * Stabilizes a tree against the baseline, with fresh refs above the + * highest ref ever applied. Returns the tree unchanged at the start of a + * session (no baseline and no refs applied yet). + * + * @param map - The new ref map. + * @param nodes - The new nodes. + * @returns The stabilized tree. + */ + #stabilize(map: Map, nodes: A11yNodeTrimmed[]): RefTree { + const baseline = this.#validBaseline(); + if (!baseline && this.#highWater === 0) { + return { nodes, refMap: map }; + } + return stabilizeRefs( + baseline ?? EMPTY_TREE, + { nodes, refMap: map }, + this.#keyMode(), + this.#highWater, + ); + } +} diff --git a/src/server/request-queue.test.ts b/src/server/request-queue.test.ts index 2ff4dee..c5bb14e 100644 --- a/src/server/request-queue.test.ts +++ b/src/server/request-queue.test.ts @@ -139,4 +139,97 @@ describe('RequestQueue', () => { expect(order).toStrictEqual([1, 2]); }); }); + + describe('abort signal', () => { + it('passes a signal that is not aborted on the normal path', async () => { + const queue = new RequestQueue(1_000); + let seen: AbortSignal | undefined; + + await queue.enqueue(async (signal) => { + seen = signal; + return 'ok'; + }); + + expect(seen).toBeInstanceOf(AbortSignal); + expect(seen?.aborted).toBe(false); + }); + + it('aborts the signal with the timeout error before rejecting', async () => { + const queue = new RequestQueue(20); + let seen: AbortSignal | undefined; + + const rejection = await queue + .enqueue(async (signal) => { + seen = signal; + return sleep(200); + }) + .catch((error: unknown) => error); + + expect(seen?.aborted).toBe(true); + expect(seen?.reason).toBe(rejection); + expect((rejection as Error).message).toBe( + 'Tool execution timed out after 20ms', + ); + }); + + it('lets a timed-out task stop at its next checkpoint', async () => { + const queue = new RequestQueue(20); + const steps: string[] = []; + + await queue + .enqueue(async (signal) => { + steps.push('step 1'); + await sleep(60); + if (signal.aborted) { + steps.push('stopped'); + return; + } + steps.push('step 2'); + }) + .catch(() => undefined); + + expect(steps).toStrictEqual(['step 1', 'stopped']); + }); + + it('gives each task its own signal', async () => { + const queue = new RequestQueue(20); + const signals: AbortSignal[] = []; + + await queue + .enqueue(async (signal) => { + signals.push(signal); + return sleep(60); + }) + .catch(() => undefined); + await queue.enqueue(async (signal) => { + signals.push(signal); + }); + + expect(signals[0].aborted).toBe(true); + expect(signals[1].aborted).toBe(false); + expect(signals[0]).not.toBe(signals[1]); + }); + + it('waits for the timed-out task to settle before the next task starts', async () => { + const queue = new RequestQueue(20); + const order: string[] = []; + + const first = queue + .enqueue(async () => { + await sleep(60); + order.push('first settled'); + }) + .catch(() => order.push('first rejected')); + const second = queue.enqueue(async () => { + order.push('second started'); + }); + + await Promise.all([first, second]); + + expect(order).toContain('first rejected'); + expect(order.indexOf('first settled')).toBeLessThan( + order.indexOf('second started'), + ); + }); + }); }); diff --git a/src/server/request-queue.ts b/src/server/request-queue.ts index 5914174..f28d390 100644 --- a/src/server/request-queue.ts +++ b/src/server/request-queue.ts @@ -20,13 +20,19 @@ export class RequestQueue { /** * Enqueues an async task for serial execution with a timeout. * - * @param fn - The async function to execute. + * Each task gets its own `AbortSignal`. When the task times out, the signal + * is aborted (with the timeout error as its reason) before the returned + * promise rejects, so the task can stop its remaining work. The queue still + * waits up to `QUEUE_SETTLE_TIMEOUT_MS` for the task to settle before the + * next task starts. + * + * @param fn - The async function to execute. It receives the task's signal. * @param options - Optional configuration for the enqueued task. * @param options.timeoutMs - Override timeout in milliseconds for this task. * @returns The resolved value of the provided function. */ async enqueue( - fn: () => Promise, + fn: (signal: AbortSignal) => Promise, options?: { timeoutMs?: number }, ): Promise { let release!: () => void; @@ -37,35 +43,36 @@ export class RequestQueue { this.#queue = next; await prev; const effectiveTimeout = options?.timeoutMs ?? this.#timeoutMs; + const controller = new AbortController(); let timer: ReturnType | undefined; - const fnPromise = fn(); + const fnPromise = fn(controller.signal); try { return await Promise.race([ fnPromise, new Promise((_resolve, reject) => { - timer = setTimeout( - () => - reject( - new Error( - `Tool execution timed out after ${effectiveTimeout}ms`, - ), - ), - effectiveTimeout, - ); + timer = setTimeout(() => { + const error = new Error( + `Tool execution timed out after ${effectiveTimeout}ms`, + ); + controller.abort(error); + reject(error); + }, effectiveTimeout); }), ]); } finally { if (timer !== undefined) { clearTimeout(timer); } + let settleTimer: ReturnType | undefined; await Promise.race([ fnPromise.catch((error) => { debugWarn('request-queue.enqueue', error); }), - new Promise((resolve) => - setTimeout(resolve, QUEUE_SETTLE_TIMEOUT_MS), - ), + new Promise((resolve) => { + settleTimer = setTimeout(resolve, QUEUE_SETTLE_TIMEOUT_MS); + }), ]); + clearTimeout(settleTimer); release(); } } diff --git a/src/server/screen-settle.test.ts b/src/server/screen-settle.test.ts new file mode 100644 index 0000000..07f92d1 --- /dev/null +++ b/src/server/screen-settle.test.ts @@ -0,0 +1,83 @@ +import { describe, it, expect } from 'vitest'; + +import { isBusyNode, isScreenBusy, screenSignature } from './screen-settle.js'; +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; + +const node = ( + role: string, + name: string, + extra: Partial = {}, +): A11yNodeTrimmed => ({ ref: 'e1', role, name, path: [], ...extra }); + +describe('isBusyNode', () => { + it('matches the React Native busy label (Android unlock proof)', () => { + // From proof/p1-08/android/click-unlock.txt: the login button while the + // wallet unlocks. + expect( + isBusyNode( + node('android.widget.Button', 'Unlock, busy', { + disabled: true, + testId: 'log-in-button', + }), + ), + ).toBe(true); + }); + + it('matches the busy label on iOS element types', () => { + expect(isBusyNode(node('Button', 'Unlock, busy'))).toBe(true); + expect(isBusyNode(node('Button', 'Confirm,busy'))).toBe(true); + expect(isBusyNode(node('Button', 'Send, Busy'))).toBe(true); + }); + + it('matches Android progress bars', () => { + expect(isBusyNode(node('android.widget.ProgressBar', ''))).toBe(true); + expect( + isBusyNode(node('androidx.core.widget.ContentLoadingProgressBar', '')), + ).toBe(true); + }); + + it('matches iOS activity and progress indicators (Appium and idb types)', () => { + expect(isBusyNode(node('ActivityIndicator', ''))).toBe(true); + expect(isBusyNode(node('ProgressIndicator', ''))).toBe(true); + expect(isBusyNode(node('UIActivityIndicatorView', ''))).toBe(true); + }); + + it('does not match other nodes', () => { + expect(isBusyNode(node('ActivityIndicatorContainer', ''))).toBe(false); + expect(isBusyNode(node('StaticText', 'ProgressIndicator'))).toBe(false); + expect(isBusyNode(node('android.widget.Button', 'Unlock'))).toBe(false); + expect(isBusyNode(node('android.widget.TextView', 'busy'))).toBe(false); + expect(isBusyNode(node('android.widget.TextView', 'Busy day'))).toBe(false); + expect(isBusyNode(node('android.widget.TextView', 'Not busy now'))).toBe( + false, + ); + // The spinner in the proof has no busy role or label; the button does. + expect( + isBusyNode( + node('android.view.ViewGroup', '', { testId: 'spinner-animated-view' }), + ), + ).toBe(false); + expect(isBusyNode(node('_UIScrollViewScrollIndicator', ''))).toBe(false); + }); +}); + +describe('isScreenBusy', () => { + it('is true when any node is busy', () => { + expect( + isScreenBusy([node('Text', 'Welcome'), node('Button', 'Unlock, busy')]), + ).toBe(true); + expect(isScreenBusy([node('Text', 'Welcome')])).toBe(false); + expect(isScreenBusy([])).toBe(false); + }); +}); + +describe('screenSignature', () => { + it('ignores refs', () => { + expect(screenSignature([{ ...node('Button', 'Go'), ref: 'e1' }])).toBe( + screenSignature([{ ...node('Button', 'Go'), ref: 'e9' }]), + ); + expect(screenSignature([node('Button', 'Go')])).not.toBe( + screenSignature([node('Button', 'Stop')]), + ); + }); +}); diff --git a/src/server/screen-settle.ts b/src/server/screen-settle.ts new file mode 100644 index 0000000..d7c8182 --- /dev/null +++ b/src/server/screen-settle.ts @@ -0,0 +1,56 @@ +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; + +/** + * React Native adds ", busy" to the accessibility label of an element with + * `accessibilityState={{ busy: true }}`. Seen on Android as the login + * button `"Unlock, busy"` while the wallet unlocks. + */ +const BUSY_LABEL = /,\s*busy$/iu; + +/** + * Roles of progress indicators, matched as a suffix of the node role (the + * mobile driver uses the backend element type as the role, see + * `normalizeSnapshot`): + * + * - Android (adb/Appium `class`): `android.widget.ProgressBar` and + * subclasses whose class name ends in `ProgressBar`. + * - iOS via Appium XCUITest: `XCUIElementTypeActivityIndicator` and + * `XCUIElementTypeProgressIndicator`, reported without the + * `XCUIElementType` prefix (`ActivityIndicator`, `ProgressIndicator`). + * - iOS via idb `describe-all`: the element `type`, either the AX role + * without `AX` (`ProgressIndicator`) or the UIKit class name + * (`UIActivityIndicatorView`, used by React Native `ActivityIndicator`). + */ +const BUSY_ROLE = + /(?:ProgressBar|ActivityIndicator(?:View)?|ProgressIndicator)$/u; + +/** + * Builds a ref-independent signature of a screen, used to tell whether the + * screen changed. + * + * @param nodes - The accessibility nodes. + * @returns The signature. + */ +export function screenSignature(nodes: A11yNodeTrimmed[]): string { + return JSON.stringify(nodes.map(({ ref: _ref, ...rest }) => rest)); +} + +/** + * Checks whether a node shows that work is in progress: its label ends in + * ", busy" (React Native busy state), or it is a progress indicator (see + * `BUSY_ROLE`). + * + * @param node - The accessibility node. + * @returns True when the node is busy. + */ +export function isBusyNode(node: A11yNodeTrimmed): boolean { + return BUSY_LABEL.test(node.name) || BUSY_ROLE.test(node.role); +} + +/** + * @param nodes - The accessibility nodes of a screen. + * @returns True when any node is busy (see `isBusyNode`). + */ +export function isScreenBusy(nodes: A11yNodeTrimmed[]): boolean { + return nodes.some(isBusyNode); +} diff --git a/src/tools/batch.test.ts b/src/tools/batch.test.ts index 742d4cd..9706b66 100644 --- a/src/tools/batch.test.ts +++ b/src/tools/batch.test.ts @@ -854,4 +854,34 @@ describe('runStepsTool', () => { expect(hermesTargetsHandler).toHaveBeenCalled(); } }); + + it('stops before the next step when the request is aborted', async () => { + const controller = new AbortController(); + const clickHandler = vi.fn().mockImplementation(async () => { + controller.abort(); + return { ok: true, result: 'clicked' }; + }); + const context = { + ...createMockContext({ + toolRegistry: new Map([['click', clickHandler]]), + }), + signal: controller.signal, + } as ToolContext; + + const result = await runStepsTool( + { + steps: [ + { tool: 'click', args: { testId: 'a' } }, + { tool: 'click', args: { testId: 'b' } }, + ], + }, + context, + ); + + expect(clickHandler).toHaveBeenCalledTimes(1); + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error.code).toBe(ErrorCodes.MM_REQUEST_ABORTED); + } + }); }); diff --git a/src/tools/batch.ts b/src/tools/batch.ts index c4a44bc..6e575b7 100644 --- a/src/tools/batch.ts +++ b/src/tools/batch.ts @@ -1,7 +1,11 @@ import { checkPlatformGate } from './registry.js'; import type { RunStepsInput, RunStepsResult, StepResult } from './types'; import { ErrorCodes } from './types'; -import { createToolError, createToolSuccess } from './utils.js'; +import { + abortedResponse, + createToolError, + createToolSuccess, +} from './utils.js'; import type { ToolContext, ToolFunction, ToolResponse } from '../types/http.js'; import { extractErrorMessage } from '../utils'; import type { ToolName } from '../validation/schemas.js'; @@ -125,6 +129,10 @@ export async function runStepsTool( : undefined; for (const stepInput of stepInputs) { + const aborted = abortedResponse(context); + if (aborted) { + return aborted; + } if (batchDeadline && Date.now() > batchDeadline) { const remainingIndex = stepInputs.indexOf(stepInput); for (const remaining of stepInputs.slice(remainingIndex)) { diff --git a/src/tools/discovery-tools.test.ts b/src/tools/discovery-tools.test.ts index 62f3ae5..a922403 100644 --- a/src/tools/discovery-tools.test.ts +++ b/src/tools/discovery-tools.test.ts @@ -115,21 +115,22 @@ describe('discovery-tools', () => { ); }); - it('updates refMap in session manager', async () => { + it('does not take an a11y snapshot or change the refMap (no refs in its result)', async () => { const context = createMockContext(); - const mockRefMap = new Map([['e1', 'role=button[name="Submit"]']]); + const setRefMap = vi.fn(); vi.spyOn(discoveryModule, 'collectTestIds').mockResolvedValue([]); - vi.spyOn(discoveryModule, 'collectTrimmedA11ySnapshot').mockResolvedValue( - { - nodes: [], - refMap: mockRefMap, - }, - ); + const snapshot = vi.spyOn(discoveryModule, 'collectTrimmedA11ySnapshot'); - await listTestIdsTool({}, context); + const result = await listTestIdsTool({}, { + ...context, + setRefMap, + } as ToolContext); - expect(context.sessionManager.setRefMap).toHaveBeenCalledWith(mockRefMap); + expect(result).toStrictEqual({ ok: true, result: { items: [] } }); + expect(snapshot).not.toHaveBeenCalled(); + expect(setRefMap).not.toHaveBeenCalled(); + expect(context.sessionManager.setRefMap).not.toHaveBeenCalled(); }); it('returns error when no active session', async () => { @@ -433,6 +434,75 @@ describe('discovery-tools', () => { expect(context.sessionManager.setRefMap).toHaveBeenCalledWith(mockRefMap); }); + it('uses one captureScreen snapshot for test IDs and tree when the driver has it', async () => { + const context = createMockContext(); + const screenNodes: A11yNodeTrimmed[] = [ + { ref: 'e1', role: 'Button', name: 'Send', path: [] }, + ]; + const screenRefMap = new Map([['e1', 'label:Send|type:Button']]); + const screenTestIds: TestIdItem[] = [ + { testId: 'wallet-screen', tag: 'View', visible: true }, + ]; + const getTestIds = vi.fn().mockResolvedValue(screenTestIds); + const captureScreen = vi.fn().mockResolvedValue({ + nodes: screenNodes, + refMap: screenRefMap, + getTestIds, + }); + const driver = { + getAppState: vi.fn().mockResolvedValue({ isLoaded: true }), + getTestIds: vi.fn(), + getAccessibilityTree: vi.fn(), + captureScreen, + }; + + const result = await describeScreenTool({}, { + ...context, + driver, + } as unknown as ToolContext); + + expect(result.ok).toBe(true); + if (result.ok) { + expect(result.result.testIds.items).toStrictEqual(screenTestIds); + expect(result.result.a11y.nodes).toStrictEqual(screenNodes); + } + expect(captureScreen).toHaveBeenCalledTimes(1); + expect(getTestIds).toHaveBeenCalledWith(150); + expect(driver.getTestIds).not.toHaveBeenCalled(); + expect(driver.getAccessibilityTree).not.toHaveBeenCalled(); + expect(context.sessionManager.setRefMap).toHaveBeenCalledWith( + screenRefMap, + ); + }); + + it('stops after captureScreen when the request is aborted', async () => { + const controller = new AbortController(); + const context = createMockContext(); + const driver = { + getAppState: vi.fn().mockResolvedValue({ isLoaded: true }), + captureScreen: vi.fn().mockResolvedValue({ + nodes: [], + refMap: new Map(), + getTestIds: vi.fn(async () => { + controller.abort(); + return []; + }), + }), + }; + + const result = await describeScreenTool({}, { + ...context, + driver, + signal: controller.signal, + } as unknown as ToolContext); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error.code).toBe(ErrorCodes.MM_REQUEST_ABORTED); + } + expect(context.sessionManager.setRefMap).not.toHaveBeenCalled(); + }); + it('returns error when no active session', async () => { const context = createMockContext({ hasActive: false }); @@ -459,4 +529,103 @@ describe('discovery-tools', () => { } }); }); + + describe('refs and abort', () => { + const nodes: A11yNodeTrimmed[] = [ + { ref: 'e1', role: 'button', name: 'Go', path: [] }, + ]; + const refMap = new Map([['e1', 'role=button[name="Go"]']]); + + beforeEach(() => { + vi.spyOn(discoveryModule, 'collectTestIds').mockResolvedValue([]); + vi.spyOn(discoveryModule, 'collectTrimmedA11ySnapshot').mockResolvedValue( + { nodes, refMap }, + ); + }); + + const nodesOf = (result: unknown): A11yNodeTrimmed[] | undefined => { + const value = (result as { result?: Record }).result; + return ( + (value?.nodes as A11yNodeTrimmed[] | undefined) ?? + (value?.a11y as { nodes: A11yNodeTrimmed[] } | undefined)?.nodes + ); + }; + + it.each([ + ['accessibility_snapshot', accessibilitySnapshotTool], + ['describe_screen', describeScreenTool], + ] as const)( + '%s sets refs through context.setRefMap and returns the written refs', + async (_name, tool) => { + const context = createMockContext(); + const remapped: A11yNodeTrimmed[] = [{ ...nodes[0], ref: 'e7' }]; + const setRefMap = vi.fn().mockReturnValue({ + refMap: new Map([['e7', 'role=button[name="Go"]']]), + nodes: remapped, + }); + const result = await tool({}, { ...context, setRefMap } as ToolContext); + + expect(result.ok).toBe(true); + expect(setRefMap).toHaveBeenCalledWith(refMap, nodes); + expect(nodesOf(result)).toStrictEqual(remapped); + expect(context.sessionManager.setRefMap).not.toHaveBeenCalled(); + }, + ); + + it.each([ + ['accessibility_snapshot', accessibilitySnapshotTool], + ['describe_screen', describeScreenTool], + ] as const)( + '%s returns the input nodes when context.setRefMap returns no nodes', + async (_name, tool) => { + const context = createMockContext(); + const setRefMap = vi.fn().mockReturnValue({ refMap }); + const result = await tool({}, { ...context, setRefMap } as ToolContext); + + expect(nodesOf(result)).toStrictEqual(nodes); + }, + ); + + it.each([ + ['accessibility_snapshot', accessibilitySnapshotTool], + ['describe_screen', describeScreenTool], + ] as const)( + '%s falls back to sessionManager.setRefMap and returns the input unchanged', + async (_name, tool) => { + const context = createMockContext(); + const result = await tool({}, context); + + expect(context.sessionManager.setRefMap).toHaveBeenCalledWith(refMap); + expect(nodesOf(result)).toStrictEqual(nodes); + }, + ); + + it.each([['describe_screen', describeScreenTool]] as const)( + '%s stops before the a11y snapshot when the request is aborted', + async (_name, tool) => { + const controller = new AbortController(); + const context = createMockContext(); + vi.spyOn(discoveryModule, 'collectTestIds').mockImplementation( + async () => { + controller.abort(); + return []; + }, + ); + + const result = await tool({}, { + ...context, + signal: controller.signal, + } as ToolContext); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error.code).toBe(ErrorCodes.MM_REQUEST_ABORTED); + } + expect( + discoveryModule.collectTrimmedA11ySnapshot, + ).not.toHaveBeenCalled(); + expect(context.sessionManager.setRefMap).not.toHaveBeenCalled(); + }, + ); + }); }); diff --git a/src/tools/discovery-tools.ts b/src/tools/discovery-tools.ts index 1d690d7..1952127 100644 --- a/src/tools/discovery-tools.ts +++ b/src/tools/discovery-tools.ts @@ -10,9 +10,11 @@ import type { import { ErrorCodes } from './types'; import { DEFAULT_TESTID_LIMIT } from './utils/constants.js'; import { + abortedResponse, createToolError, createToolSuccess, requireActiveSession, + setToolRefMap, } from './utils.js'; import type { ToolContext, ToolResponse } from '../types/http.js'; @@ -61,10 +63,9 @@ export async function listTestIdsTool( const limit = input.limit ?? DEFAULT_TESTID_LIMIT; + // The result has no refs, so the refMap is left as it is. try { const items = await driver.getTestIds(limit); - const { refMap } = await driver.getAccessibilityTree(); - context.sessionManager.setRefMap(refMap); return createToolSuccess({ items }); } catch (error) { const errorInfo = classifyDiscoveryError(error); @@ -95,10 +96,8 @@ export async function accessibilitySnapshotTool( const { driver } = driverResult; try { - const { nodes, refMap } = await driver.getAccessibilityTree( - input.rootSelector, - ); - context.sessionManager.setRefMap(refMap); + const tree = await driver.getAccessibilityTree(input.rootSelector); + const { nodes } = setToolRefMap(context, tree.refMap, tree.nodes); return createToolSuccess({ nodes }); } catch (error) { const errorInfo = classifyDiscoveryError(error); @@ -129,10 +128,20 @@ export async function describeScreenTool( try { const state = await driver.getAppState(); - const testIds = await driver.getTestIds(DEFAULT_TESTID_LIMIT); - const { nodes, refMap } = await driver.getAccessibilityTree(); - - context.sessionManager.setRefMap(refMap); + // Mobile: one device snapshot for both the test IDs and the tree, when + // the driver supports it. + const screen = driver.captureScreen + ? await driver.captureScreen() + : undefined; + const testIds = screen + ? await screen.getTestIds(DEFAULT_TESTID_LIMIT) + : await driver.getTestIds(DEFAULT_TESTID_LIMIT); + const aborted = abortedResponse(context); + if (aborted) { + return aborted; + } + const tree = screen ?? (await driver.getAccessibilityTree()); + const { nodes } = setToolRefMap(context, tree.refMap, tree.nodes); const trackedPages = context.sessionManager.getTrackedPages(); let activeTab: DescribeScreenResult['activeTab']; diff --git a/src/tools/driver-delegation.test.ts b/src/tools/driver-delegation.test.ts index 31b15d8..df2c349 100644 --- a/src/tools/driver-delegation.test.ts +++ b/src/tools/driver-delegation.test.ts @@ -392,7 +392,8 @@ describe('tool delegation to context.driver', () => { expect(result.ok).toBe(true); expect(driver.getTestIds).toHaveBeenCalledOnce(); - expect(driver.getAccessibilityTree).toHaveBeenCalledOnce(); + // Its result has no refs, so it does not read the a11y tree. + expect(driver.getAccessibilityTree).not.toHaveBeenCalled(); }); it('accessibilitySnapshotTool delegates to driver.getAccessibilityTree', async () => { diff --git a/src/tools/interaction.test.ts b/src/tools/interaction.test.ts index 2cba370..9d9c1c1 100644 --- a/src/tools/interaction.test.ts +++ b/src/tools/interaction.test.ts @@ -1240,4 +1240,45 @@ describe('interaction', () => { } }); }); + + describe('abort', () => { + it.each([ + [ + 'click', + async (context: ToolContext) => clickTool({ testId: 'a' }, context), + ], + [ + 'type', + async (context: ToolContext) => + typeTool({ testId: 'a', text: 'x' }, context), + ], + [ + 'wait_for', + async (context: ToolContext) => waitForTool({ testId: 'a' }, context), + ], + [ + 'get_text', + async (context: ToolContext) => getTextTool({ testId: 'a' }, context), + ], + ])( + '%s returns MM_REQUEST_ABORTED without touching the page when aborted', + async (_name, run) => { + const controller = new AbortController(); + controller.abort(); + const context = { + ...createMockContext(), + signal: controller.signal, + } as ToolContext; + const waitSpy = vi.spyOn(discoveryModule, 'waitForTarget'); + + const result = await run(context); + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.error.code).toBe(ErrorCodes.MM_REQUEST_ABORTED); + } + expect(waitSpy).not.toHaveBeenCalled(); + }, + ); + }); }); diff --git a/src/tools/interaction.ts b/src/tools/interaction.ts index cb2fbc9..7dcc3b4 100644 --- a/src/tools/interaction.ts +++ b/src/tools/interaction.ts @@ -24,6 +24,7 @@ import { isValidTargetSelection, } from './utils/type-guards.js'; import { + abortedResponse, createToolError, createToolSuccess, requireActiveSession, @@ -58,6 +59,11 @@ function validateInteraction( return { error: missingSession }; } + const aborted = abortedResponse(context); + if (aborted) { + return { error: aborted }; + } + if (!context.driver) { return { error: createToolError( diff --git a/src/tools/types/errors.ts b/src/tools/types/errors.ts index 5aac6a3..ed42a79 100644 --- a/src/tools/types/errors.ts +++ b/src/tools/types/errors.ts @@ -49,6 +49,8 @@ export const ErrorCodes = { MM_TOOL_NOT_SUPPORTED_ON_PLATFORM: 'MM_TOOL_NOT_SUPPORTED_ON_PLATFORM', // tool gated off the active platform (browser-only / mobile-only) + MM_REQUEST_ABORTED: 'MM_REQUEST_ABORTED', // the request timed out; the tool stopped at its next checkpoint + MM_UNKNOWN_TOOL: 'MM_UNKNOWN_TOOL', MM_INTERNAL_ERROR: 'MM_INTERNAL_ERROR', } as const; diff --git a/src/tools/utils.ts b/src/tools/utils.ts index 6eae6a2..2825585 100644 --- a/src/tools/utils.ts +++ b/src/tools/utils.ts @@ -1,4 +1,5 @@ import { ErrorCodes } from './types'; +import type { A11yNodeTrimmed } from './types/discovery.js'; import type { ToolContext, ToolResponse } from '../types/http.js'; /** @@ -51,3 +52,46 @@ export function requireActiveSession( return undefined; } + +/** + * Returns an error response when the request was aborted (it timed out). + * Tools call it between steps so a timed-out request stops early. + * + * @param context - The tool execution context. + * @returns An error response when aborted, or undefined. + */ +export function abortedResponse( + context: ToolContext, +): ToolResponse | undefined { + if (context.signal?.aborted) { + return createToolError( + ErrorCodes.MM_REQUEST_ABORTED, + 'The request timed out and was cancelled.', + ); + } + return undefined; +} + +/** + * Sets the refMap through the request's guarded setter when the daemon + * provides one, else directly on the session manager. + * + * @param context - The tool execution context. + * @param map - The new ref map. + * @param nodes - The nodes the map was built from. + * @returns The map and nodes to return to the agent. The daemon may renumber + * refs so elements the agent already holds keep their refs; without the + * daemon the input is returned unchanged. + */ +export function setToolRefMap( + context: ToolContext, + map: Map, + nodes: A11yNodeTrimmed[], +): { refMap: Map; nodes: A11yNodeTrimmed[] } { + if (context.setRefMap) { + const written = context.setRefMap(map, nodes); + return { refMap: written.refMap, nodes: written.nodes ?? nodes }; + } + context.sessionManager.setRefMap(map); + return { refMap: map, nodes }; +} diff --git a/src/tools/utils/constants.ts b/src/tools/utils/constants.ts index a9a5204..817c0af 100644 --- a/src/tools/utils/constants.ts +++ b/src/tools/utils/constants.ts @@ -22,6 +22,29 @@ export const MAX_IDLE_CHECK_INTERVAL_MS = 60_000; /** Hard timeout for post-tool observation collection (testIds, a11y) */ export const OBSERVATION_TIMEOUT_MS = 8_000; +/** + * Mobile observation settle: after a mutating or batch tool, when a snapshot + * still shows the last screen the agent received or a busy element, wait + * this long and take the snapshot again. + */ +export const MOBILE_SETTLE_MS = 300; + +/** + * Mobile observation settle: no wait starts that would end later than this + * after the first snapshot (waits and snapshots both count). + */ +export const MOBILE_SETTLE_MAX_MS = 2_000; + +/** Mobile observation settle: maximum number of extra snapshots. */ +export const MOBILE_SETTLE_MAX_SNAPSHOTS = 3; + +/** + * Mobile observation settle: extra snapshots allowed because the screen + * still equals the last screen the agent received (no busy element). Busy + * screens are retried up to `MOBILE_SETTLE_MAX_SNAPSHOTS`. + */ +export const MOBILE_SETTLE_UNCHANGED_RETRIES = 1; + /** Timeout for graceful HTTP server close during daemon shutdown */ export const GRACEFUL_SHUTDOWN_TIMEOUT_MS = 10_000; diff --git a/src/types/http.ts b/src/types/http.ts index ff73d91..7ccede0 100644 --- a/src/types/http.ts +++ b/src/types/http.ts @@ -9,6 +9,7 @@ import type { Page } from '@playwright/test'; import type { PortMap, WorkflowContext } from '../capabilities/context.js'; import type { IPlatformDriver } from '../platform/types.js'; import type { ISessionManager } from '../server/session-manager.js'; +import type { A11yNodeTrimmed } from '../tools/types/discovery.js'; /** * Context passed to standalone tool functions. @@ -29,6 +30,25 @@ export type ToolContext = { toolRegistry: Map>; /** Platform driver for cross-platform element interaction (browser, iOS, Android) */ driver?: IPlatformDriver; + /** + * Aborted when the request times out. Tools check it between steps and + * stop. Absent when the context was not built by the daemon. + */ + signal?: AbortSignal; + /** + * Sets the refMap for this request. A write from a request that timed out + * or was replaced by a newer request is dropped. Pass the nodes the map + * was built from: refs are then renumbered so elements the agent already + * holds keep their refs, and the returned `refMap` and `nodes` carry the + * refs that were written. Return those to the agent, not the input. + * Absent when the context was not built by the daemon; fall back to + * `sessionManager.setRefMap`. (In the daemon, `sessionManager.setRefMap` + * is routed here too, without nodes.) + */ + setRefMap?: ( + map: Map, + nodes?: A11yNodeTrimmed[], + ) => { refMap: Map; nodes?: A11yNodeTrimmed[] }; }; /** From 7d0c5654e7b23be7dc31bcaec663fbc250d31d81 Mon Sep 17 00:00:00 2001 From: cryptotavares Date: Sat, 24 Oct 2026 02:19:06 +0100 Subject: [PATCH 2/3] fix: clamp queued task timeouts CodeQL (js/resource-exhaustion) flagged the queue timer because its duration comes from the request body. Clamp it to MAX_QUEUE_TIMEOUT_MS. --- src/server/request-queue.test.ts | 26 +++++++++++++++++++++++++- src/server/request-queue.ts | 14 ++++++++++++-- src/tools/utils/constants.ts | 3 +++ 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/src/server/request-queue.test.ts b/src/server/request-queue.test.ts index c5bb14e..4b9416b 100644 --- a/src/server/request-queue.test.ts +++ b/src/server/request-queue.test.ts @@ -1,6 +1,10 @@ -import { describe, it, expect } from 'vitest'; +import { describe, it, expect, vi } from 'vitest'; import { RequestQueue } from './request-queue.js'; +import { + MAX_QUEUE_TIMEOUT_MS, + QUEUE_SETTLE_TIMEOUT_MS, +} from '../tools/utils/constants.js'; async function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); @@ -232,4 +236,24 @@ describe('RequestQueue', () => { ); }); }); + + it('clamps an oversized timeout to MAX_QUEUE_TIMEOUT_MS', async () => { + vi.useFakeTimers(); + try { + const queue = new RequestQueue(); + const pending = queue + .enqueue(async () => new Promise(() => undefined), { + timeoutMs: Number.MAX_SAFE_INTEGER, + }) + .catch((error: unknown) => error); + await vi.advanceTimersByTimeAsync( + MAX_QUEUE_TIMEOUT_MS + QUEUE_SETTLE_TIMEOUT_MS, + ); + expect(String(await pending)).toContain( + `timed out after ${MAX_QUEUE_TIMEOUT_MS}ms`, + ); + } finally { + vi.useRealTimers(); + } + }); }); diff --git a/src/server/request-queue.ts b/src/server/request-queue.ts index f28d390..6a939e9 100644 --- a/src/server/request-queue.ts +++ b/src/server/request-queue.ts @@ -1,4 +1,7 @@ -import { QUEUE_SETTLE_TIMEOUT_MS } from '../tools/utils/constants.js'; +import { + MAX_QUEUE_TIMEOUT_MS, + QUEUE_SETTLE_TIMEOUT_MS, +} from '../tools/utils/constants.js'; import { debugWarn } from '../utils'; /** @@ -42,7 +45,14 @@ export class RequestQueue { const prev = this.#queue; this.#queue = next; await prev; - const effectiveTimeout = options?.timeoutMs ?? this.#timeoutMs; + // Clamp caller-supplied timeouts (they come from request bodies). + let effectiveTimeout = options?.timeoutMs ?? this.#timeoutMs; + if (!(effectiveTimeout > 0)) { + effectiveTimeout = this.#timeoutMs; + } + if (effectiveTimeout > MAX_QUEUE_TIMEOUT_MS) { + effectiveTimeout = MAX_QUEUE_TIMEOUT_MS; + } const controller = new AbortController(); let timer: ReturnType | undefined; const fnPromise = fn(controller.signal); diff --git a/src/tools/utils/constants.ts b/src/tools/utils/constants.ts index 817c0af..686e11e 100644 --- a/src/tools/utils/constants.ts +++ b/src/tools/utils/constants.ts @@ -13,6 +13,9 @@ export const DEFAULT_INTERACTION_TIMEOUT_MS = 15000; /** Maximum time the RequestQueue finally block waits for a hung task to settle */ export const QUEUE_SETTLE_TIMEOUT_MS = 5_000; +/** Upper bound for any queued task timeout (largest schema timeout + budget). */ +export const MAX_QUEUE_TIMEOUT_MS = 600_000; + /** Timeout for waiting for CSS/Web animations to settle after a mutation */ export const ANIMATION_SETTLE_TIMEOUT_MS = 3_000; From dcc7a0aae33e8f34034918d524d155e623764895 Mon Sep 17 00:00:00 2001 From: cryptotavares Date: Mon, 26 Oct 2026 18:11:54 +0000 Subject: [PATCH 3/3] refactor: simplify ref map merging and test types --- src/server/create-server.test.ts | 1 + src/server/ref-store.ts | 7 +++---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/server/create-server.test.ts b/src/server/create-server.test.ts index b950418..9da9306 100644 --- a/src/server/create-server.test.ts +++ b/src/server/create-server.test.ts @@ -18,6 +18,7 @@ import { toolRegistry } from '../tools/registry.js'; import type { DaemonState, ServerConfig, + ToolContext, ToolFunction, ToolResponse, } from '../types/http.js'; diff --git a/src/server/ref-store.ts b/src/server/ref-store.ts index e8cb771..5a0955d 100644 --- a/src/server/ref-store.ts +++ b/src/server/ref-store.ts @@ -227,13 +227,12 @@ export class RefStore { const tree = this.#stabilize(map, nodes); let { refMap } = tree; if (toolMap && options?.keepToolRefs !== false) { - const merged = new Map(toolMap); + refMap = new Map(toolMap); for (const [ref, target] of tree.refMap) { - if (!merged.has(ref)) { - merged.set(ref, target); + if (!refMap.has(ref)) { + refMap.set(ref, target); } } - refMap = merged; } const applied = apply('observation', refMap, tree.nodes); return { applied, refMap, nodes: tree.nodes };