feat(ext): improve token validation and error handling

packages/extension/src/agent/MultiPageAgent.ts

@@ -35,6 +35,9 @@ export class MultiPageAgent extends PageAgentCore {
 				chrome.storage.local.set({
 					isAgentRunning: false,
 				})
+
+				// no need to dispose tabsController and pageController
+				// as they do not keep references
 			},
 		})
 	}

packages/extension/src/entrypoints/content.ts

@@ -13,13 +13,22 @@ export default defineContentScript({
 
 		// if auth token matches, expose agent to page
 		chrome.storage.local.get('PageAgentExtUserAuthToken').then((result) => {
-			if (!result.PageAgentExtUserAuthToken) return
+			// extension side token.
-			if (!localStorage.getItem('PageAgentExtUserAuthToken')) return
+			// @note this is isolated world. it is safe to assume user script cannot access it
-			if (localStorage.getItem('PageAgentExtUserAuthToken') !== result.PageAgentExtUserAuthToken)
+			const extToken = result.PageAgentExtUserAuthToken
-				return
+			if (!extToken) return
 
-			exposeAgentToPage()
+			// page side token
-			injectScript('/main-world.js')
+			const pageToken = localStorage.getItem('PageAgentExtUserAuthToken')
+			if (!pageToken) return
+
+			if (pageToken !== extToken) return
+
+			// add isolated world script
+			exposeAgentToPage().then(
+				// add main-world script
+				() => injectScript('/main-world.js')
+			)
 		})
 	},
 })
@@ -42,9 +51,24 @@ async function exposeAgentToPage() {
 
 		switch (action) {
 			case 'execute': {
-				if (!multiPageAgent) multiPageAgent = new MultiPageAgent(DEMO_CONFIG)
+				if (multiPageAgent && multiPageAgent.status === 'running') {
+					window.postMessage(
+						{
+							channel: 'PAGE_AGENT_EXT_RESPONSE',
+							id,
+							action: 'execute_result',
+							error: 'Agent is already running a task. Please wait until it finishes.',
+						},
+						'*'
+					)
+					return
+				}
+
+				try {
+					multiPageAgent = new MultiPageAgent(DEMO_CONFIG)
 
 					const result = await multiPageAgent.execute(payload)
+
 					window.postMessage(
 						{
 							channel: 'PAGE_AGENT_EXT_RESPONSE',
@@ -54,6 +78,18 @@ async function exposeAgentToPage() {
 						},
 						'*'
 					)
+				} catch (error) {
+					window.postMessage(
+						{
+							channel: 'PAGE_AGENT_EXT_RESPONSE',
+							id,
+							action: 'execute_result',
+							error: (error as Error).message,
+						},
+						'*'
+					)
+				}
+
 				break
 			}
 

packages/extension/src/entrypoints/main-world.ts

@@ -10,7 +10,7 @@ export default defineUnlistedScript(() => {
 	w.execute = async (task: string) => {
 		const id = getId()
 
-		const promise = new Promise((resolve) => {
+		const promise = new Promise((resolve, reject) => {
 			function handleMessage(e: MessageEvent) {
 				const data = e.data
 				if (typeof data !== 'object' || data === null) return
@@ -19,8 +19,13 @@ export default defineUnlistedScript(() => {
 				if (data.id !== id) return
 
 				window.removeEventListener('message', handleMessage)
+
+				if (data.error) {
+					reject(new Error(data.error))
+				} else {
 					resolve(data.payload)
 				}
+			}
 
 			window.addEventListener('message', handleMessage)
 		})