feat(browser): add scoped observations and shared wait budgets

Author: waleedlatif1

apps/desktop/src/main/browser-agent/driver.test.ts

@@ -3257,6 +3257,53 @@ describe('credential protection', () => {
     })
   })
 
+  it.each(['browser_snapshot', 'browser_find'] as const)(
+    'omits an absent scope from the serialized %s page call',
+    async (tool) => {
+      const contents = await openPage()
+      vi.mocked(contents.executeJavaScript).mockClear()
+      await driver.executeTool('chat-test', tool, { query: 'Continue' })
+      const expressions = vi
+        .mocked(contents.executeJavaScript)
+        .mock.calls.map(([expression]) => expression)
+        .filter((expression) => isPageCall(expression, 'collectSnapshot'))
+      expect(expressions).toHaveLength(1)
+      expect(expressions[0]).toContain('.apply(null, [1])')
+    }
+  )
+
+  it.each(['browser_snapshot', 'browser_find'] as const)(
+    'passes the current root ref to %s and invalidates previous refs',
+    async (tool) => {
+      const contents = await openPage()
+      respondWith(contents, {
+        collectSnapshot: {
+          url: 'https://example.com/login',
+          title: 'Scoped',
+          scoped: true,
+          outline: '- button "Save" [ref=1]',
+          truncated: false,
+          refIds: [1],
+          refLineIndexes: { 1: 0 },
+          nextElementId: 2,
+        },
+      })
+      const result = await driver.executeTool('chat-test', tool, { elementId: 0, query: 'Save' })
+      expect(result).toMatchObject({ ok: true, result: { scoped: true } })
+      expect(
+        vi
+          .mocked(contents.executeJavaScript)
+          .mock.calls.some(
+            ([expression]) =>
+              isPageCall(expression, 'collectSnapshot') &&
+              expression.includes('.apply(null, [1,0])')
+          )
+      ).toBe(true)
+      const stale = await driver.executeTool('chat-test', 'browser_click', { elementId: 0 })
+      expect(stale).toMatchObject({ ok: false })
+    }
+  )
+
   it.each([
     [1, false, 1],
     [100, false, 50],
@@ -3405,6 +3452,47 @@ describe('credential protection', () => {
     })
   })
 
+  it.each([true, false])(
+    'polls delayed checkable state without redispatching input (updates=%s)',
+    async (updates) => {
+      const contents = await openPage()
+      let reads = 0
+      vi.mocked(contents.executeJavaScript).mockImplementation(async (expression: string) => {
+        if (isPageCall(expression, 'readCheckableElementState')) {
+          reads++
+          return { checked: updates && reads >= 4, kind: 'input:checkbox' }
+        }
+        if (isPageCall(expression, 'clickElement'))
+          return { dispatched: false, x: 24, y: 48, element: 'Checkbox' }
+        if (isPageCall(expression, 'readPageActionState'))
+          return {
+            url: 'https://example.com/login',
+            title: 'Example',
+            focus: 'body',
+            mutationRevision: 0,
+            dialogs: [],
+            scroll: [0],
+          }
+        if (isPageCall(expression, 'readActiveElementState')) return {}
+      })
+      vi.useFakeTimers()
+      try {
+        const pending = driver.executeTool('chat-test', 'browser_set_checked', {
+          elementId: 0,
+          checked: true,
+        })
+        await vi.advanceTimersByTimeAsync(2000)
+        const result = await pending
+        expect(result.ok).toBe(updates)
+        expect(reads).toBeGreaterThanOrEqual(4)
+        expect(cdpCalls(contents, 'Input.dispatchMouseEvent')).toHaveLength(3)
+        if (!updates) expect(result.error).toContain('did not reach the requested checked state')
+      } finally {
+        vi.useRealTimers()
+      }
+    }
+  )
+
   it.each(['hidden', 'detached'])(
     'does not treat a failed probe as element state %s',
     async (state) => {

apps/desktop/src/main/browser-agent/driver.ts

@@ -1203,6 +1203,11 @@ function unwrapPageResult(result: unknown): unknown {
         'Element-state waits are limited to the top page. Use a text or URL condition for framed content.'
       )
     }
+    if (code === 'framed-snapshot') {
+      throw new ToolError(
+        'Scoped snapshots require a top-page element. Omit elementId to capture framed content.'
+      )
+    }
     if (code === 'no-option') {
       const options = (result as { options?: string[] }).options ?? []
       throw new ToolError(
@@ -1951,12 +1956,21 @@ function validateSnapshotRefs(
  * policy intentionally lacks; password redaction still runs inside every frame
  * before any result crosses back to the driver.
  */
-async function captureSnapshot(contents: WebContents, notAfter?: number): Promise<unknown> {
+async function captureSnapshot(
+  contents: WebContents,
+  notAfter?: number,
+  elementId?: number
+): Promise<unknown> {
   const state = driverScopeState()
   const tab = session.requireAutomationTab()
   if (tab.view.webContents !== contents) {
     throw new ToolError('The active tab changed before the snapshot started. Try again.')
   }
+  if (elementId !== undefined && pageTargetForElement(contents, elementId) !== contents) {
+    throw new ToolError(
+      'Scoped snapshots require a top-page element. Omit elementId to capture framed content.'
+    )
+  }
   invalidateSnapshot(state)
   const captureEpoch = state.snapshotCaptureEpoch
   const capturedTabId = tab.id
@@ -1977,12 +1991,14 @@ async function captureSnapshot(contents: WebContents, notAfter?: number): Promis
   }
 
   const mainStartingElementId = state.nextElementRefId
-  const mainSnapshot = await execInPage(
-    contents,
-    collectSnapshot,
-    [mainStartingElementId],
-    false,
-    notAfter
+  const mainSnapshot = unwrapPageResult(
+    await execInPage(
+      contents,
+      collectSnapshot,
+      elementId === undefined ? [mainStartingElementId] : [mainStartingElementId, elementId],
+      false,
+      notAfter
+    )
   )
   if (!stillCurrent()) {
     throw new ToolError('The tab changed while its snapshot was being captured. Try again.')
@@ -2012,7 +2028,7 @@ async function captureSnapshot(contents: WebContents, notAfter?: number): Promis
   let capturedCrossOriginFrames = 0
   let unreadableCrossOriginFrames = 0
   let hiddenCrossOriginFrames = 0
-  const boundaryFrames = crossOriginBoundaryFrames(contents)
+  const boundaryFrames = elementId === undefined ? crossOriginBoundaryFrames(contents) : []
   const frames = boundaryFrames.slice(0, MAX_CROSS_ORIGIN_SCAN_FRAMES)
   if (boundaryFrames.length > frames.length) truncated = true
 
@@ -2432,7 +2448,7 @@ async function executeToolInner(
     case 'browser_snapshot': {
       const contents = session.requireAutomationTab().view.webContents
       assertCurrentExecution()
-      return await captureSnapshot(contents, executionDeadline)
+      return await captureSnapshot(contents, executionDeadline, num(params, 'elementId'))
     }
 
     case 'browser_find': {
@@ -2441,7 +2457,9 @@ async function executeToolInner(
       const requestedMax = num(params, 'maxResults')
       const maxResults = Math.min(50, Math.max(1, Math.floor(requestedMax ?? 20)))
       const contents = session.requireAutomationTab().view.webContents
-      const snapshot = toRecord(await captureSnapshot(contents, executionDeadline))
+      const snapshot = toRecord(
+        await captureSnapshot(contents, executionDeadline, num(params, 'elementId'))
+      )
       const outline = typeof snapshot.outline === 'string' ? snapshot.outline : ''
       const needle = query.toLowerCase()
       const matches = outline.split('\n').flatMap((line) => {
@@ -2457,6 +2475,7 @@ async function executeToolInner(
         truncated: snapshot.truncated === true || matches.length > maxResults,
         url: snapshot.url,
         title: snapshot.title,
+        ...(snapshot.scoped === true ? { scoped: true } : {}),
       }
     }
 
@@ -3656,16 +3675,31 @@ async function executeToolInner(
           invocationEpoch
         )
       )
-      await sleep(100)
-      const after = toRecord(
-        unwrapPageResult(
-          await execInPage(target, readCheckableElementState, [elementId], false, executionDeadline)
-        )
+      const readbackDeadline = Math.min(
+        Date.now() + SETTLE_GRACE_MS,
+        executionDeadline ?? Number.POSITIVE_INFINITY
       )
-      if (after.checked !== checked) {
-        throw new ToolError(
-          'The control did not retain the requested checked state. Take a fresh browser_snapshot and inspect the page.'
+      let after: Record<string, unknown>
+      for (;;) {
+        assertCurrentExecution()
+        after = toRecord(
+          unwrapPageResult(
+            await execInPage(
+              target,
+              readCheckableElementState,
+              [elementId],
+              false,
+              executionDeadline
+            )
+          )
         )
+        if (after.checked === checked) break
+        if (Date.now() + SETTLE_PROBE_INTERVAL_MS > readbackDeadline) {
+          throw new ToolError(
+            'The control did not reach the requested checked state. Take a fresh browser_snapshot and inspect the page.'
+          )
+        }
+        await sleep(SETTLE_PROBE_INTERVAL_MS)
       }
       return {
         checked,

apps/desktop/src/main/browser-agent/page-functions.test.ts

@@ -954,6 +954,87 @@ describe('collectSnapshot', () => {
 })
 
 describe('semantic control state', () => {
+  it('scopes a fresh snapshot to one card without reading sibling geometry', () => {
+    document.body.innerHTML =
+      '<div role="group" tabindex="0" aria-label="Selected card"><button>Save card</button><input type="password" value="private" /></div><button>Outside card</button>'
+    for (const element of document.querySelectorAll('*')) visible(element)
+    const full = outlineOf(collectSnapshot())
+    const oldRootRef = refFor(full, 'Selected card')
+    const outside = document.querySelector('body > button')!
+    const outsideGeometry = vi.spyOn(outside, 'getBoundingClientRect')
+
+    const scoped = runSerialized(collectSnapshot, [100, oldRootRef]) as {
+      scoped: boolean
+      outline: string
+      refIds: number[]
+    }
+    expect(scoped.scoped).toBe(true)
+    expect(scoped.outline).toContain('Selected card')
+    expect(scoped.outline).toContain('Save card')
+    expect(scoped.outline).not.toContain('Outside card')
+    expect(scoped.outline).not.toContain('private')
+    expect(scoped.refIds.every((id) => id >= 100)).toBe(true)
+    expect(window.__simAgentResolveElement?.(oldRootRef)).toBeNull()
+    expect(outsideGeometry).not.toHaveBeenCalled()
+  })
+
+  it('rejects stale or framed snapshot roots without replacing the page registry', () => {
+    const button = document.createElement('button')
+    document.body.append(visible(button))
+    register(button)
+    button.remove()
+    expect(collectSnapshot(10, 0)).toMatchObject({ error: 'stale' })
+    expect(window.__simAgentElements?.[0]).toBe(button)
+
+    const frame = document.createElement('iframe')
+    document.body.append(frame)
+    const child = frame.contentDocument!.createElement('button')
+    frame.contentDocument!.body.append(visible(child))
+    register(child)
+    expect(collectSnapshot(10, 0)).toEqual({ error: 'framed-snapshot' })
+  })
+
+  it('marks unreadable scoped frame content truncated', () => {
+    const root = visible(document.createElement('div'))
+    const frame = visible(document.createElement('iframe'))
+    root.append(frame)
+    document.body.append(root)
+    register(root)
+    Object.defineProperty(frame, 'contentDocument', { value: null })
+    expect(collectSnapshot(10, 0)).toMatchObject({ scoped: true, truncated: true })
+  })
+
+  it('does not recover scoped refs into a different card after the original closes', () => {
+    document.body.innerHTML =
+      '<div tabindex="0" aria-label="Selected card"><button id="save">Save card</button></div>'
+    for (const element of document.querySelectorAll('*')) visible(element)
+    const root = document.querySelector('body > div')!
+    const rootRef = refFor(outlineOf(collectSnapshot()), 'Selected card')
+    const saveRef = refFor(outlineOf(collectSnapshot(10, rootRef)), 'Save card')
+    const replacement = root.cloneNode(true) as HTMLElement
+    for (const element of [replacement, ...replacement.querySelectorAll('*')]) visible(element)
+    root.replaceWith(replacement)
+
+    expect(window.__simAgentResolveElement?.(saveRef)).toBeNull()
+    expect(window.__simAgentStaleReason).toContain('scoped snapshot root')
+  })
+
+  it('keeps scoped snapshots bounded when a selected container is very large', () => {
+    const root = document.createElement('div')
+    root.tabIndex = 0
+    root.setAttribute('aria-label', 'Large card')
+    document.body.append(visible(root))
+    register(root)
+    for (let index = 0; index < 400; index++) {
+      const button = visible(document.createElement('button'))
+      button.textContent = `Action ${index}`
+      root.append(button)
+    }
+    const scoped = collectSnapshot(10, 0) as { refIds: number[]; truncated: boolean }
+    expect(scoped.refIds).toHaveLength(300)
+    expect(scoped.truncated).toBe(true)
+  })
+
   it('distinguishes hidden registered nodes from detached nodes without action recovery', () => {
     const button = visible(document.createElement('button'))
     document.body.append(button)