feat(extension): honor AbortSignal in tab tools and tab-load wait (#590)

packages/core/src/PageAgentCore.ts

@@ -23,7 +23,7 @@ import type {
 } from './types'
 import { assert, fetchLlmsTxt, normalizeResponse, suppress, uid, waitFor } from './utils'
 
-export { tool, type PageAgentTool } from './tools'
+export { tool, type PageAgentTool, type ToolContext } from './tools'
 export type * from './types'
 
 export type PageAgentCoreConfig = AgentConfig & { pageController: PageController }

packages/extension/package.json

@@ -7,6 +7,7 @@
         "dev": "wxt",
         "build:ext": "wxt build",
         "zip": "wxt zip",
+        "test": "vitest run",
         "postinstall": "wxt prepare"
     },
     "devDependencies": {

packages/extension/src/agent/TabsController.test.ts

@@ -0,0 +1,78 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { TabsController } from './TabsController'
+
+describe('TabsController.waitUntilTabLoaded', () => {
+	interface TabRow {
+		id: number
+		isInitial: boolean
+		status?: 'loading' | 'unloaded' | 'complete'
+	}
+
+	// White-box helper: build a controller with a known tab list and a stubbed `syncTabs`
+	// (the only chrome-backed dependency the wait touches), so tab-status transitions can be
+	// driven deterministically without the background service worker.
+	function makeController(
+		tabs: TabRow[],
+		onSync: () => void = () => {}
+	): { controller: TabsController; syncCount: () => number } {
+		const controller = new TabsController()
+		;(controller as unknown as { tabs: TabRow[] }).tabs = tabs
+		let syncs = 0
+		;(controller as unknown as { syncTabs: () => Promise<void> }).syncTabs = async () => {
+			syncs += 1
+			onSync()
+		}
+		return { controller, syncCount: () => syncs }
+	}
+
+	it('throws for an unknown tab id', async () => {
+		const { controller } = makeController([{ id: 1, isInitial: false, status: 'complete' }])
+		await expect(controller.waitUntilTabLoaded(999)).rejects.toThrow('not found')
+	})
+
+	it('resolves once a loading tab transitions to complete during the wait', async () => {
+		const tabs: TabRow[] = [{ id: 1, isInitial: false, status: 'loading' }]
+		const { controller, syncCount } = makeController(tabs, () => {
+			// The background reports the tab finished loading after a couple of polls.
+			if (syncCount() >= 2) tabs[0].status = 'complete'
+		})
+
+		await expect(controller.waitUntilTabLoaded(1)).resolves.toBeUndefined()
+		expect(syncCount()).toBeGreaterThanOrEqual(2)
+	})
+
+	it('throws when the tab ends up unloaded', async () => {
+		const tabs: TabRow[] = [{ id: 1, isInitial: false, status: 'loading' }]
+		const { controller } = makeController(tabs, () => {
+			tabs[0].status = 'unloaded'
+		})
+
+		await expect(controller.waitUntilTabLoaded(1)).rejects.toThrow('unloaded')
+	})
+
+	it('rejects with an AbortError when aborted while the tab is still loading', async () => {
+		// syncTabs never leaves the tab in `loading`, so only the signal can end the wait.
+		const tabs: TabRow[] = [{ id: 1, isInitial: false, status: 'loading' }]
+		const { controller } = makeController(tabs)
+		const ac = new AbortController()
+
+		const promise = controller.waitUntilTabLoaded(1, { signal: ac.signal })
+		setTimeout(() => ac.abort(), 20)
+
+		await expect(promise).rejects.toMatchObject({ name: 'AbortError' })
+	})
+
+	it('rejects immediately without polling when the signal is already aborted', async () => {
+		const tabs: TabRow[] = [{ id: 1, isInitial: false, status: 'loading' }]
+		const { controller, syncCount } = makeController(tabs)
+		const ac = new AbortController()
+		ac.abort()
+
+		await expect(controller.waitUntilTabLoaded(1, { signal: ac.signal })).rejects.toMatchObject({
+			name: 'AbortError',
+		})
+		// The already-aborted signal must short-circuit before the wait polls (no syncTabs).
+		expect(syncCount()).toBe(0)
+	})
+})

packages/extension/src/agent/TabsController.ts

@@ -129,7 +129,7 @@ export class TabsController {
 		await this.updateCurrentTabId(this.currentTabId)
 	}
 
-	async openNewTab(url: string): Promise<string> {
+	async openNewTab(url: string, options: { signal?: AbortSignal } = {}): Promise<string> {
 		debug('openNewTab', url)
 
 		const result = await sendMessage({
@@ -161,7 +161,7 @@ export class TabsController {
 			})
 		}
 
-		await this.waitUntilTabLoaded(tabId)
+		await this.waitUntilTabLoaded(tabId, options)
 
 		return `✅ Opened new tab ID ${tabId} with URL ${url}`
 	}
@@ -293,7 +293,7 @@ export class TabsController {
 		return summaries.join('\n')
 	}
 
-	async waitUntilTabLoaded(tabId: number): Promise<void> {
+	async waitUntilTabLoaded(tabId: number, options: { signal?: AbortSignal } = {}): Promise<void> {
 		const tab = this.tabs.find((t) => t.id === tabId)
 		if (!tab) throw new Error(`Tab ID ${tabId} not found in tab list.`)
 		if (tab.status === 'complete') return
@@ -303,11 +303,16 @@ export class TabsController {
 		// Finding the latest tab object is the only way to know if it's closed.
 
 		debug('waitUntilTabLoaded', tabId)
-		await waitUntil(async () => {
-			await this.syncTabs()
-			const latest = this.tabs.find((t) => t.id === tabId)
-			return !latest || latest.status !== 'loading'
-		}, 4_000)
+		await waitUntil(
+			async () => {
+				await this.syncTabs()
+				const latest = this.tabs.find((t) => t.id === tabId)
+				return !latest || latest.status !== 'loading'
+			},
+			4_000,
+			false,
+			options.signal
+		)
 
 		const latest = this.tabs.find((t) => t.id === tabId)
 		if (latest?.status === 'unloaded') throw new Error(`Tab ID ${tabId} is unloaded.`)
@@ -426,31 +431,24 @@ function randomColor(): TabGroupColor {
  * @returns Returns when condition becomes true, false if timeout
  * @param timeoutMS Timeout in milliseconds, default 1 minutes
  * @param throwIfTimeout Reject on timeout instead of resolving with `false`
+ * @param signal Abort the wait early; rejects with the signal's reason (an `AbortError`).
+ *   Observed once per poll iteration, not during an in-flight `check()`.
  */
 async function waitUntil(
 	check: () => boolean | Promise<boolean>,
 	timeoutMS = 60_000,
-	throwIfTimeout = false
+	throwIfTimeout = false,
+	signal?: AbortSignal
 ): Promise<boolean> {
-	if (await check()) return true
-
-	return new Promise((resolve, reject) => {
-		const start = Date.now()
-		const poll = async () => {
-			try {
-				if (await check()) return resolve(true)
-				if (Date.now() - start > timeoutMS) {
-					if (throwIfTimeout) {
-						return reject(new Error(`waitUntil timed out after ${timeoutMS}ms`))
-					} else {
-						return resolve(false)
-					}
-				}
-				setTimeout(poll, 100)
-			} catch (err) {
-				reject(err instanceof Error ? err : new Error(String(err)))
-			}
+	const start = Date.now()
+	while (true) {
+		signal?.throwIfAborted()
+		if (await check()) return true
+		signal?.throwIfAborted()
+		if (Date.now() - start > timeoutMS) {
+			if (throwIfTimeout) throw new Error(`waitUntil timed out after ${timeoutMS}ms`)
+			return false
 		}
-		setTimeout(poll, 100)
-	})
+		await new Promise((resolve) => setTimeout(resolve, 100))
+	}
 }

packages/extension/src/agent/tabTools.ts

@@ -6,6 +6,7 @@
  * - switch_to_tab: Switch to an existing tab
  * - close_tab: Close a tab (optionally switch to another)
  */
+import type { ToolContext } from '@page-agent/core'
 import * as z from 'zod/v4'
 
 import type { TabsController } from './TabsController'
@@ -14,7 +15,7 @@ import type { TabsController } from './TabsController'
 interface TabTool {
 	description: string
 	inputSchema: z.ZodType
-	execute: (input: unknown) => Promise<string>
+	execute: (input: unknown, ctx: ToolContext) => Promise<string>
 }
 
 /**
@@ -29,11 +30,13 @@ export function createTabTools(tabsController: TabsController): Record<string, T
 			inputSchema: z.object({
 				url: z.string().describe('The URL to open in the new tab'),
 			}),
-			execute: async (input: unknown) => {
+			execute: async (input: unknown, { signal }: ToolContext) => {
 				const { url } = input as { url: string }
 				try {
-					return await tabsController.openNewTab(url)
+					return await tabsController.openNewTab(url, { signal })
 				} catch (error) {
+					// Let cancellation propagate instead of masking it as a tool failure.
+					if (signal.aborted) throw error
 					return `❌ Failed: ${error instanceof Error ? error.message : String(error)}`
 				}
 			},

packages/extension/vitest.config.js

@@ -0,0 +1,9 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+	test: {
+		name: 'ext',
+		include: ['src/**/*.test.ts'],
+		silent: 'passed-only',
+	},
+})