diff --git a/src/lib/agent/BrowserAgent.ts b/src/lib/agent/BrowserAgent.ts index c29757721..a57e0f83a 100644 --- a/src/lib/agent/BrowserAgent.ts +++ b/src/lib/agent/BrowserAgent.ts @@ -59,6 +59,7 @@ import { createGetSelectedTabsTool } from '@/lib/tools/tab/GetSelectedTabsTool'; import { createClassificationTool } from '@/lib/tools/classification/ClassificationTool'; import { createValidatorTool } from '@/lib/tools/validation/ValidatorTool'; import { createScreenshotTool } from '@/lib/tools/utils/ScreenshotTool'; +import { createStorageTool } from '@/lib/tools/utils/StorageTool'; import { createExtractTool } from '@/lib/tools/extraction/ExtractTool'; import { createResultTool } from '@/lib/tools/result/ResultTool'; import { createHumanInputTool } from '@/lib/tools/utils/HumanInputTool'; @@ -281,6 +282,7 @@ export class BrowserAgent { // util tools this.toolManager.register(createScreenshotTool(this.executionContext)); + this.toolManager.register(createStorageTool(this.executionContext)); this.toolManager.register(createExtractTool(this.executionContext)); this.toolManager.register(createHumanInputTool(this.executionContext)); diff --git a/src/lib/runtime/StorageManager.ts b/src/lib/runtime/StorageManager.ts new file mode 100644 index 000000000..6a7f4b7b8 --- /dev/null +++ b/src/lib/runtime/StorageManager.ts @@ -0,0 +1,58 @@ +/** + * StorageManager handles persistent key-value storage for the agent + * using chrome.storage.local with automatic JSON serialization + */ +export class StorageManager { + private static readonly STORAGE_PREFIX = 'nxtscape_agent_' + + /** + * Store a value with automatic JSON serialization + */ + static async set(key: string, value: any): Promise { + const storageKey = this.STORAGE_PREFIX + key + + try { + const serialized = JSON.stringify(value) + await chrome.storage.local.set({ [storageKey]: serialized }) + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error) + throw new Error(`Failed to store value: ${errorMessage}`) + } + } + + /** + * Get a value with automatic JSON deserialization + */ + static async get(key: string): Promise { + const storageKey = this.STORAGE_PREFIX + key + + try { + const result = await chrome.storage.local.get(storageKey) + const serialized = result[storageKey] + + if (serialized === undefined) { + return null + } + + return JSON.parse(serialized) + } catch (error) { + // If JSON parse fails, return null + console.warn(`Failed to parse stored value for key ${key}:`, error) + return null + } + } + + /** + * Clear all agent storage + */ + static async clearAll(): Promise { + const allKeys = await chrome.storage.local.get(null) + const agentKeys = Object.keys(allKeys).filter(k => + k.startsWith(this.STORAGE_PREFIX) + ) + + if (agentKeys.length > 0) { + await chrome.storage.local.remove(agentKeys) + } + } +} \ No newline at end of file diff --git a/src/lib/tools/utils/StorageTool.test.ts b/src/lib/tools/utils/StorageTool.test.ts new file mode 100644 index 000000000..2b98911cc --- /dev/null +++ b/src/lib/tools/utils/StorageTool.test.ts @@ -0,0 +1,101 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { StorageTool } from './StorageTool' +import { StorageManager } from '@/lib/runtime/StorageManager' +import { ExecutionContext } from '@/lib/runtime/ExecutionContext' + +// Mock the StorageManager +vi.mock('@/lib/runtime/StorageManager', () => ({ + StorageManager: { + set: vi.fn(), + get: vi.fn(), + clearAll: vi.fn() + } +})) + +describe('StorageTool', () => { + let storageTool: StorageTool + let mockExecutionContext: ExecutionContext + + beforeEach(() => { + vi.clearAllMocks() + mockExecutionContext = {} as ExecutionContext + storageTool = new StorageTool(mockExecutionContext) + }) + + it('tests that the tool can store a value', async () => { + const input = { + action: 'set' as const, + key: 'test_key', + value: { data: 'test_value' } + } + + vi.mocked(StorageManager.set).mockResolvedValue(undefined) + + const result = await storageTool.execute(input) + + expect(result.ok).toBe(true) + expect(result.output).toBe('Stored value for key: test_key') + expect(StorageManager.set).toHaveBeenCalledWith('test_key', { data: 'test_value' }) + }) + + it('tests that the tool can retrieve a stored value', async () => { + const input = { + action: 'get' as const, + key: 'test_key', + value: undefined + } + + const storedValue = { data: 'test_value' } + vi.mocked(StorageManager.get).mockResolvedValue(storedValue) + + const result = await storageTool.execute(input) + + expect(result.ok).toBe(true) + expect(result.output).toBe(JSON.stringify(storedValue)) + expect(StorageManager.get).toHaveBeenCalledWith('test_key') + }) + + it('tests that the tool handles missing values gracefully', async () => { + const input = { + action: 'get' as const, + key: 'missing_key', + value: undefined + } + + vi.mocked(StorageManager.get).mockResolvedValue(null) + + const result = await storageTool.execute(input) + + expect(result.ok).toBe(true) + expect(result.output).toBe('No value found for key: missing_key') + }) + + it('tests that the tool requires a value for set operations', async () => { + const input = { + action: 'set' as const, + key: 'test_key', + value: undefined + } + + const result = await storageTool.execute(input) + + expect(result.ok).toBe(false) + expect(result.output).toBe('Value is required for set operation') + expect(StorageManager.set).not.toHaveBeenCalled() + }) + + it('tests that the tool handles storage errors gracefully', async () => { + const input = { + action: 'set' as const, + key: 'test_key', + value: 'test_value' + } + + vi.mocked(StorageManager.set).mockRejectedValue(new Error('Storage error')) + + const result = await storageTool.execute(input) + + expect(result.ok).toBe(false) + expect(result.output).toBe('Storage operation failed: Storage error') + }) +}) \ No newline at end of file diff --git a/src/lib/tools/utils/StorageTool.ts b/src/lib/tools/utils/StorageTool.ts new file mode 100644 index 000000000..7e666b018 --- /dev/null +++ b/src/lib/tools/utils/StorageTool.ts @@ -0,0 +1,68 @@ +import { z } from 'zod' +import { DynamicStructuredTool } from '@langchain/core/tools' +import { ExecutionContext } from '@/lib/runtime/ExecutionContext' +import { StorageManager } from '@/lib/runtime/StorageManager' +import { toolSuccess, toolError, type ToolOutput } from '@/lib/tools/Tool.interface' + +// Input schema for storage operations +const StorageToolInputSchema = z.object({ + action: z.enum(['get', 'set']), // Storage operation + key: z.string().min(1).max(100), // Storage key + value: z.any().optional() // Value to store (for set operation) +}) + +type StorageToolInput = z.infer + +export class StorageTool { + constructor(private executionContext: ExecutionContext) {} + + async execute(input: StorageToolInput): Promise { + const { action, key, value } = input + + try { + if (action === 'set') { + return await this._handleSet(key, value) + } else { + return await this._handleGet(key) + } + } catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error) + return toolError(`Storage operation failed: ${errorMessage}`) + } + } + + private async _handleSet(key: string, value: any): Promise { + if (value === undefined) { + return toolError('Value is required for set operation') + } + + await StorageManager.set(key, value) + return toolSuccess(`Stored value for key: ${key}`) + } + + private async _handleGet(key: string): Promise { + const value = await StorageManager.get(key) + + if (value === null) { + return toolSuccess(`No value found for key: ${key}`) + } + + // Return the value as JSON string + return toolSuccess(JSON.stringify(value)) + } +} + +// Factory function to create StorageTool +export function createStorageTool(executionContext: ExecutionContext): DynamicStructuredTool { + const storageTool = new StorageTool(executionContext) + + return new DynamicStructuredTool({ + name: 'storage_tool', + description: 'Store and retrieve JSON values persistently using get/set operations', + schema: StorageToolInputSchema, + func: async (args: StorageToolInput): Promise => { + const result = await storageTool.execute(args) + return JSON.stringify(result) + } + }) +} \ No newline at end of file diff --git a/src/lib/tools/utils/index.ts b/src/lib/tools/utils/index.ts index 492dba584..9e01091dc 100644 --- a/src/lib/tools/utils/index.ts +++ b/src/lib/tools/utils/index.ts @@ -1,2 +1,3 @@ export { createDoneTool } from './DoneTool'; -export { createScreenshotTool } from './ScreenshotTool'; \ No newline at end of file +export { createScreenshotTool } from './ScreenshotTool'; +export { createStorageTool } from './StorageTool'; \ No newline at end of file