refactor: zod tree-shaking; better error handling in agent steps

packages/core/src/PageAgentCore.ts

@@ -5,13 +5,13 @@
 import { InvokeError, LLM, type Tool } from '@page-agent/llms'
 import type { PageController } from '@page-agent/page-controller'
 import chalk from 'chalk'
-import zod from 'zod'
+import * as zod from 'zod'
 
 import { type PageAgentConfig } from './config'
 import { DEFAULT_MAX_STEPS } from './config/constants'
 import SYSTEM_PROMPT from './prompts/system_prompt.md?raw'
 import { tools } from './tools'
-import {
+import type {
 	AgentActivity,
 	AgentReflection,
 	AgentStatus,
@@ -21,8 +21,7 @@ import {
 	MacroToolInput,
 	MacroToolResult,
 } from './types'
-import { normalizeResponse, trimLines, uid, waitFor } from './utils'
-import { assert } from './utils/assert'
+import { assert, normalizeResponse, trimLines, uid, waitFor } from './utils'
 
 export { type PageAgentConfig }
 export { tool, type PageAgentTool } from './tools'
@@ -209,20 +208,20 @@ export class PageAgentCore extends EventTarget {
 			lastURL: '',
 		}
 
-		try {
-			let step = 0
+		let step = 0
+
+		while (true) {
+			try {
+				console.group(`step: ${step}`)
 
-			while (true) {
 				await this.#generateObservations(step)
 
 				await onBeforeStep?.(this, step)
 
-				console.group(`step: ${step}`)
-
 				// abort
 				if (this.#abortController.signal.aborted) throw new Error('AbortError')
 
-				// Thinking
+				// status
 				console.log(chalk.blue('Thinking...'))
 				this.#emitActivity({ type: 'thinking' })
 
@@ -270,22 +269,13 @@ export class PageAgentCore extends EventTarget {
 
 				//
 
+				await onAfterStep?.(this, this.history)
+
 				console.log(chalk.green('Step finished:'), actionName)
 				console.groupEnd()
 
-				await onAfterStep?.(this, this.history)
+				// finish task if done
 
-				step++
-				if (step > this.config.maxSteps) {
-					this.#onDone(false)
-					const result: ExecutionResult = {
-						success: false,
-						data: 'Step count exceeded maximum limit',
-						history: this.history,
-					}
-					await onAfterTask?.(this, result)
-					return result
-				}
 				if (actionName === 'done') {
 					const success = action.input?.success ?? false
 					const text = action.input?.text || 'no text provided'
@@ -299,19 +289,32 @@ export class PageAgentCore extends EventTarget {
 					await onAfterTask?.(this, result)
 					return result
 				}
+			} catch (error: unknown) {
+				console.groupEnd() // to prevent nested groups
+				console.error('Task failed', error)
+				const errorMessage = String(error)
+				this.#emitActivity({ type: 'error', message: errorMessage })
+				this.#onDone(false)
+				const result: ExecutionResult = {
+					success: false,
+					data: errorMessage,
+					history: this.history,
+				}
+				await onAfterTask?.(this, result)
+				return result
 			}
-		} catch (error: unknown) {
-			console.error('Task failed', error)
-			const errorMessage = String(error)
-			this.#emitActivity({ type: 'error', message: errorMessage })
-			this.#onDone(false)
-			const result: ExecutionResult = {
-				success: false,
-				data: errorMessage,
-				history: this.history,
+
+			step++
+			if (step > this.config.maxSteps) {
+				this.#onDone(false)
+				const result: ExecutionResult = {
+					success: false,
+					data: 'Step count exceeded maximum limit',
+					history: this.history,
+				}
+				await onAfterTask?.(this, result)
+				return result
 			}
-			await onAfterTask?.(this, result)
-			return result
 		}
 	}
 

packages/core/src/tools/index.ts

@@ -2,7 +2,7 @@
  * Internal tools for PageAgent.
  * @note Adapted from browser-use
  */
-import zod, { type z } from 'zod'
+import * as zod from 'zod'
 
 import type { PageAgentCore } from '../PageAgentCore'
 import { waitFor } from '../utils'
@@ -13,7 +13,7 @@ import { waitFor } from '../utils'
 export interface PageAgentTool<TParams = any> {
 	// name: string
 	description: string
-	inputSchema: z.ZodType<TParams>
+	inputSchema: zod.ZodType<TParams>
 	execute: (this: PageAgentCore, args: TParams) => Promise<string>
 }
 

packages/core/src/utils/assert.ts

@@ -1,17 +0,0 @@
-import chalk from 'chalk'
-
-/**
- * Simple assertion function that throws an error if the condition is falsy
- * @param condition - The condition to assert
- * @param message - Optional error message
- * @throws Error if condition is falsy
- */
-export function assert(condition: unknown, message?: string, silent?: boolean): asserts condition {
-	if (!condition) {
-		const errorMessage = message ?? 'Assertion failed'
-
-		if (!silent) console.error(chalk.red(`❌ assert: ${errorMessage}`))
-
-		throw new Error(errorMessage)
-	}
-}

packages/core/src/utils/index.ts

@@ -1,3 +1,5 @@
+import chalk from 'chalk'
+
 export { normalizeResponse } from './autoFixer'
 
 /**
@@ -85,3 +87,19 @@ export function uid() {
 	ids.push(id)
 	return id
 }
+
+/**
+ * Simple assertion function that throws an error if the condition is falsy
+ * @param condition - The condition to assert
+ * @param message - Optional error message
+ * @throws Error if condition is falsy
+ */
+export function assert(condition: unknown, message?: string, silent?: boolean): asserts condition {
+	if (!condition) {
+		const errorMessage = message ?? 'Assertion failed'
+
+		if (!silent) console.error(chalk.red(`❌ assert: ${errorMessage}`))
+
+		throw new Error(errorMessage)
+	}
+}