From 5cfaa292d315733506a8232a4604fbc3a827a96a Mon Sep 17 00:00:00 2001
From: Simon <10131203+gaomeng1900@users.noreply.github.com>
Date: Thu, 29 Jan 2026 19:21:57 +0800
Subject: [PATCH] feat(ext): add agent heart beat check

---
 .../extension/src/agent/MultiPageAgent.ts     | 30 +++++++++++++++++++
 .../src/agent/RemotePageController.content.ts | 17 +++++------
 .../extension/src/entrypoints/background.ts   |  2 +-
 3 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/packages/extension/src/agent/MultiPageAgent.ts b/packages/extension/src/agent/MultiPageAgent.ts
index 342b71d..f53461e 100644
--- a/packages/extension/src/agent/MultiPageAgent.ts
+++ b/packages/extension/src/agent/MultiPageAgent.ts
@@ -5,12 +5,26 @@ import { TabsController } from './TabsController'
 import SYSTEM_PROMPT from './system_prompt.md?raw'
 import { createTabTools } from './tabTools'
 
+/**
+ * MultiPageAgent
+ * - use with extension
+ * - can be used from a side panel or a content script
+ */
 export class MultiPageAgent extends PageAgentCore {
 	constructor(config: Omit<PageAgentConfig, 'pageController'>) {
 		const tabsController = new TabsController()
 		const pageController = new RemotePageController(tabsController)
 		const customTools = createTabTools(tabsController)
 
+		/**
+		 * When the agent is in side-panel and user closed the side-panel.
+		 * There is no chance for isAgentRunning to be set false.
+		 * (unload event doesn't work well in side panel.)
+		 * (I'm trying not to use long-lived connection because the lifecycle of a sw is hard to predict.)
+		 * This heartbeat mechanism acts as a backup.
+		 */
+		let heartBeatInterval: null | number = null
+
 		super({
 			...config,
 			pageController: pageController as any,
@@ -20,18 +34,34 @@ export class MultiPageAgent extends PageAgentCore {
 			onBeforeTask: async (agent) => {
 				await tabsController.init(agent.task)
 
+				heartBeatInterval = window.setInterval(() => {
+					chrome.storage.local.set({
+						agentHeartbeat: Date.now(),
+					})
+				}, 1_000)
+
 				await chrome.storage.local.set({
 					isAgentRunning: true,
 				})
 			},
 
 			onAfterTask: async () => {
+				if (heartBeatInterval) {
+					window.clearInterval(heartBeatInterval)
+					heartBeatInterval = null
+				}
+
 				await chrome.storage.local.set({
 					isAgentRunning: false,
 				})
 			},
 
 			onDispose: () => {
+				if (heartBeatInterval) {
+					window.clearInterval(heartBeatInterval)
+					heartBeatInterval = null
+				}
+
 				chrome.storage.local.set({
 					isAgentRunning: false,
 				})
diff --git a/packages/extension/src/agent/RemotePageController.content.ts b/packages/extension/src/agent/RemotePageController.content.ts
index b200fd3..a49784b 100644
--- a/packages/extension/src/agent/RemotePageController.content.ts
+++ b/packages/extension/src/agent/RemotePageController.content.ts
@@ -21,17 +21,14 @@ export function initPageController() {
 	}
 
 	intervalID = window.setInterval(async () => {
+		const agentHeartbeat = (await chrome.storage.local.get('agentHeartbeat')).agentHeartbeat
+		const now = Date.now()
+		const agentInTouch = typeof agentHeartbeat === 'number' && now - agentHeartbeat < 2_000
+
 		const isAgentRunning = (await chrome.storage.local.get('isAgentRunning')).isAgentRunning
 		const currentTabId = (await chrome.storage.local.get('currentTabId')).currentTabId
 
-		const shouldShowMask = isAgentRunning && currentTabId === (await myTabIdPromise)
-
-		// console.log('[RemotePageController] polling:', {
-		// 	isAgentRunning,
-		// 	currentTabId,
-		// 	myTabId: await myTabIdPromise,
-		// 	shouldShowMask,
-		// })
+		const shouldShowMask = isAgentRunning && agentInTouch && currentTabId === (await myTabIdPromise)
 
 		if (shouldShowMask) {
 			const pc = getPC()
@@ -45,13 +42,13 @@ export function initPageController() {
 			}
 		}
 
-		if (!isAgentRunning) {
+		if (!isAgentRunning && agentInTouch) {
 			if (pageController) {
 				pageController.dispose()
 				pageController = null
 			}
 		}
-	}, 1_000)
+	}, 500)
 
 	chrome.runtime.onMessage.addListener((message, sender, sendResponse): true | undefined => {
 		if (message.type !== 'PAGE_CONTROL') {
diff --git a/packages/extension/src/entrypoints/background.ts b/packages/extension/src/entrypoints/background.ts
index 434d0f5..bec06d2 100644
--- a/packages/extension/src/entrypoints/background.ts
+++ b/packages/extension/src/entrypoints/background.ts
@@ -1,7 +1,7 @@
 import { handlePageControlMessage } from '@/agent/RemotePageController.background'
 import { handleTabControlMessage } from '@/agent/TabsController.background'
 
-chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
+chrome.runtime.onMessage.addListener((message, sender, sendResponse): true | undefined => {
 	if (message.type === 'TAB_CONTROL') {
 		return handleTabControlMessage(message, sender, sendResponse)
 	} else if (message.type === 'PAGE_CONTROL') {