From 71d1cd4df37062e1da83a44d0b52098e7671099b Mon Sep 17 00:00:00 2001
From: Simon <10131203+gaomeng1900@users.noreply.github.com>
Date: Tue, 21 Oct 2025 20:05:39 +0800
Subject: [PATCH] feat: add lifecycle hooks for task and step events

---
 ROADMAP.md          |  5 +++--
 src/PageAgent.ts    | 25 +++++++++++++++++++------
 src/config/index.ts | 17 +++++++++++++++--
 3 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/ROADMAP.md b/ROADMAP.md
index affde59..18e75e8 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -17,8 +17,9 @@ The development progress and future plans for PageAgent.
 - [x] **Free evaluation plan?**
 - [x] **Custom actions and HITL**
 - [ ] **Hooks and Events**
-	- [x] **Step lifecycle**
-- [ ] **Hijacking `page_open` event**
+	- [x] **lifecycle**
+- [ ] **Pause and intervene**
+- [ ] **Hijacking `page_open/page_change/page_unload` event**
 - [ ] **Custom knowledge base and instructions**
 - [ ] **Black/white-list safeguard**
 - [ ] **Data-masking**
diff --git a/src/PageAgent.ts b/src/PageAgent.ts
index 668d649..20c412d 100644
--- a/src/PageAgent.ts
+++ b/src/PageAgent.ts
@@ -134,6 +134,13 @@ export class PageAgent extends EventTarget {
 		if (!task) throw new Error('Task is required')
 		this.task = task
 
+		const onBeforeStep = this.config.onBeforeStep || (() => void 0)
+		const onAfterStep = this.config.onAfterStep || (() => void 0)
+		const onBeforeTask = this.config.onBeforeTask || (() => void 0)
+		const onAfterTask = this.config.onAfterTask || (() => void 0)
+
+		await onBeforeTask.call(this, task)
+
 		// Show mask and panel
 		this.mask.show()
 
@@ -142,7 +149,7 @@ export class PageAgent extends EventTarget {
 
 		this.bus.emit('panel:update', {
 			type: 'input',
-			displayText: task,
+			displayText: this.task,
 		})
 
 		if (this.#abortController) {
@@ -156,7 +163,7 @@ export class PageAgent extends EventTarget {
 			let step = 0
 
 			while (true) {
-				if (this.config.onBeforeStep) await this.config.onBeforeStep.call(this, step)
+				await onBeforeStep.call(this, step)
 
 				console.group(`step: ${step + 1}`)
 
@@ -211,37 +218,43 @@ export class PageAgent extends EventTarget {
 				console.log(chalk.green('Step finished:'), actionName)
 				console.groupEnd()
 
-				if (this.config.onAfterStep) await this.config.onAfterStep.call(this, step, this.history)
+				await onAfterStep.call(this, step, this.history)
 
 				step++
 				if (step > MAX_STEPS) {
 					this.#onDone('Step count exceeded maximum limit', false)
-					return {
+					const result: ExecutionResult = {
 						success: false,
 						data: 'Step count exceeded maximum limit',
 						history: this.history,
 					}
+					await onAfterTask.call(this, task, result)
+					return result
 				}
 				if (actionName === 'done') {
 					const success = action.input?.success ?? false
 					const text = action.input?.text || 'no text provided'
 					console.log(chalk.green.bold('Task completed'), success, text)
 					this.#onDone(text, success)
-					return {
+					const result: ExecutionResult = {
 						success,
 						data: text,
 						history: this.history,
 					}
+					await onAfterTask.call(this, task, result)
+					return result
 				}
 			}
 		} catch (error: unknown) {
 			console.error('Task failed', error)
 			this.#onDone(String(error), false)
-			return {
+			const result: ExecutionResult = {
 				success: false,
 				data: String(error),
 				history: this.history,
 			}
+			await onAfterTask.call(this, task, result)
+			return result
 		}
 	}
 
diff --git a/src/config/index.ts b/src/config/index.ts
index 3859fde..5d806ae 100644
--- a/src/config/index.ts
+++ b/src/config/index.ts
@@ -1,4 +1,4 @@
-import type { AgentHistory, PageAgent } from '@/PageAgent'
+import type { AgentHistory, ExecutionResult, PageAgent } from '@/PageAgent'
 import type { DomConfig } from '@/dom'
 import type { SupportedLanguage } from '@/i18n'
 import type { PageAgentTool } from '@/tools'
@@ -32,6 +32,7 @@ export interface UIConfig {
 	 * @see [tools](../tools/index.ts)
 	 *
 	 * @example
+	 * // override internal tool
 	 * import { tool } from 'page-agent'
 	 * const customTools = {
 	 * ask_user: tool({
@@ -46,13 +47,25 @@ export interface UIConfig {
 	 * 	},
 	 * })
 	 * }
+	 *
+	 * @example
+	 * // remove internal tool
+	 * const customTools = {
+	 * 	ask_user: null // never ask user questions
+	 * }
 	 */
 	customTools?: Record<string, PageAgentTool | null>
 
-	// hooks
+	// lifecycle hooks
 
 	onBeforeStep?: (this: PageAgent, stepCnt: number) => Promise<void> | void
 	onAfterStep?: (this: PageAgent, stepCnt: number, history: AgentHistory[]) => Promise<void> | void
+	onBeforeTask?: (this: PageAgent, task: string) => Promise<void> | void
+	onAfterTask?: (this: PageAgent, task: string, result: ExecutionResult) => Promise<void> | void
+
+	// page behavior hooks
+
+	onPageUnload?: (this: PageAgent) => Promise<void> | void
 }
 
 export type PageAgentConfig = LLMConfig & DomConfig & UIConfig