diff --git a/packages/core/src/PageAgentCore.ts b/packages/core/src/PageAgentCore.ts index 86af83f..c3dba4c 100644 --- a/packages/core/src/PageAgentCore.ts +++ b/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 } diff --git a/packages/extension/package.json b/packages/extension/package.json index c80b8c2..b693cff 100644 --- a/packages/extension/package.json +++ b/packages/extension/package.json @@ -7,6 +7,7 @@ "dev": "wxt", "build:ext": "wxt build", "zip": "wxt zip", + "test": "vitest run", "postinstall": "wxt prepare" }, "devDependencies": { diff --git a/packages/extension/src/agent/TabsController.test.ts b/packages/extension/src/agent/TabsController.test.ts new file mode 100644 index 0000000..d58046a --- /dev/null +++ b/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 }).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) + }) +}) diff --git a/packages/extension/src/agent/TabsController.ts b/packages/extension/src/agent/TabsController.ts index 32286e4..e58e8dc 100644 --- a/packages/extension/src/agent/TabsController.ts +++ b/packages/extension/src/agent/TabsController.ts @@ -129,7 +129,7 @@ export class TabsController { await this.updateCurrentTabId(this.currentTabId) } - async openNewTab(url: string): Promise { + async openNewTab(url: string, options: { signal?: AbortSignal } = {}): Promise { 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 { + async waitUntilTabLoaded(tabId: number, options: { signal?: AbortSignal } = {}): Promise { 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, timeoutMS = 60_000, - throwIfTimeout = false + throwIfTimeout = false, + signal?: AbortSignal ): Promise { - 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)) + } } diff --git a/packages/extension/src/agent/tabTools.ts b/packages/extension/src/agent/tabTools.ts index 25fe1db..43fb1b8 100644 --- a/packages/extension/src/agent/tabTools.ts +++ b/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 + execute: (input: unknown, ctx: ToolContext) => Promise } /** @@ -29,11 +30,13 @@ export function createTabTools(tabsController: TabsController): Record { + 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)}` } }, diff --git a/packages/extension/vitest.config.js b/packages/extension/vitest.config.js new file mode 100644 index 0000000..d56cdc7 --- /dev/null +++ b/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', + }, +})