From d6323f148801b4fbf811617856cc1c5e07b93c61 Mon Sep 17 00:00:00 2001
From: Simon <10131203+gaomeng1900@users.noreply.github.com>
Date: Wed, 28 Jan 2026 19:59:55 +0800
Subject: [PATCH] feat(ext): improve token validation and error handling

---
 .../extension/src/agent/MultiPageAgent.ts     |  3 +
 packages/extension/src/entrypoints/content.ts | 70 ++++++++++++++-----
 .../extension/src/entrypoints/main-world.ts   |  9 ++-
 3 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/packages/extension/src/agent/MultiPageAgent.ts b/packages/extension/src/agent/MultiPageAgent.ts
index 8c86c3d..342b71d 100644
--- a/packages/extension/src/agent/MultiPageAgent.ts
+++ b/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
 			},
 		})
 	}
diff --git a/packages/extension/src/entrypoints/content.ts b/packages/extension/src/entrypoints/content.ts
index 28bcb52..c664ef9 100644
--- a/packages/extension/src/entrypoints/content.ts
+++ b/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
-			if (!localStorage.getItem('PageAgentExtUserAuthToken')) return
-			if (localStorage.getItem('PageAgentExtUserAuthToken') !== result.PageAgentExtUserAuthToken)
-				return
+			// extension side token.
+			// @note this is isolated world. it is safe to assume user script cannot access it
+			const extToken = result.PageAgentExtUserAuthToken
+			if (!extToken) return
 
-			exposeAgentToPage()
-			injectScript('/main-world.js')
+			// page side token
+			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,18 +51,45 @@ 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',
+							id,
+							action: 'execute_result',
+							payload: result,
+						},
+						'*'
+					)
+				} catch (error) {
+					window.postMessage(
+						{
+							channel: 'PAGE_AGENT_EXT_RESPONSE',
+							id,
+							action: 'execute_result',
+							error: (error as Error).message,
+						},
+						'*'
+					)
+				}
 
-				const result = await multiPageAgent.execute(payload)
-				window.postMessage(
-					{
-						channel: 'PAGE_AGENT_EXT_RESPONSE',
-						id,
-						action: 'execute_result',
-						payload: result,
-					},
-					'*'
-				)
 				break
 			}
 
diff --git a/packages/extension/src/entrypoints/main-world.ts b/packages/extension/src/entrypoints/main-world.ts
index f5a5a6b..2be06aa 100644
--- a/packages/extension/src/entrypoints/main-world.ts
+++ b/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,7 +19,12 @@ export default defineUnlistedScript(() => {
 				if (data.id !== id) return
 
 				window.removeEventListener('message', handleMessage)
-				resolve(data.payload)
+
+				if (data.error) {
+					reject(new Error(data.error))
+				} else {
+					resolve(data.payload)
+				}
 			}
 
 			window.addEventListener('message', handleMessage)