EVM plugin wallet provider and transfer action improvements

packages/plugin-evm/src/actions/transfer.ts

@@ -1,5 +1,6 @@
 import { ByteArray, formatEther, parseEther, type Hex } from "viem";
 import {
+    Action,
     composeContext,
     generateObjectDeprecated,
     HandlerCallback,
@@ -13,8 +14,6 @@ import { initWalletProvider, WalletProvider } from "../providers/wallet";
 import type { Transaction, TransferParams } from "../types";
 import { transferTemplate } from "../templates";
 
-export { transferTemplate };
-
 // Exported for tests
 export class TransferAction {
     constructor(private walletProvider: WalletProvider) {}
@@ -72,21 +71,17 @@ const buildTransferDetails = async (
     runtime: IAgentRuntime,
     wp: WalletProvider
 ): Promise<TransferParams> => {
+    const chains = Object.keys(wp.chains);
+    state.supportedChains = chains.map((item) => `"${item}"`).join("|");
+
     const context = composeContext({
         state,
         template: transferTemplate,
     });
 
-    const chains = Object.keys(wp.chains);
-
-    const contextWithChains = context.replace(
-        "SUPPORTED_CHAINS",
-        chains.map((item) => `"${item}"`).join("|")
-    );
-
     const transferDetails = (await generateObjectDeprecated({
         runtime,
-        context: contextWithChains,
+        context,
         modelClass: ModelClass.SMALL,
     })) as TransferParams;
 
@@ -104,16 +99,22 @@ const buildTransferDetails = async (
     return transferDetails;
 };
 
-export const transferAction = {
+export const transferAction: Action = {
     name: "transfer",
     description: "Transfer tokens between addresses on the same chain",
     handler: async (
         runtime: IAgentRuntime,
-        _message: Memory,
+        message: Memory,
         state: State,
         _options: any,
         callback?: HandlerCallback
     ) => {
+        if (!state) {
+            state = (await runtime.composeState(message)) as State;
+        } else {
+            state = await runtime.updateRecentMessageState(state);
+        }
+
         console.log("Transfer action handler called");
         const walletProvider = await initWalletProvider(runtime);
         const action = new TransferAction(walletProvider);
@@ -151,7 +152,6 @@ export const transferAction = {
             return false;
         }
     },
-    template: transferTemplate,
     validate: async (runtime: IAgentRuntime) => {
         const privateKey = runtime.getSetting("EVM_PRIVATE_KEY");
         return typeof privateKey === "string" && privateKey.startsWith("0x");

packages/plugin-evm/src/providers/wallet.ts

@@ -5,7 +5,7 @@ import {
     http,
 } from "viem";
 import { privateKeyToAccount } from "viem/accounts";
-import type { IAgentRuntime, Provider, Memory, State } from "@elizaos/core";
+import { type IAgentRuntime, type Provider, type Memory, type State, type ICacheManager, elizaLogger } from "@elizaos/core";
 import type {
     Address,
     WalletClient,
@@ -17,16 +17,22 @@ import type {
 } from "viem";
 import * as viemChains from "viem/chains";
 import { DeriveKeyProvider, TEEMode } from "@elizaos/plugin-tee";
+import NodeCache from "node-cache";
+import * as path from "path";
 
 import type { SupportedChain } from "../types";
 
 export class WalletProvider {
+    private cache: NodeCache;
+    private cacheKey: string = "evm/wallet";
     private currentChain: SupportedChain = "mainnet";
+    private CACHE_EXPIRY_SEC = 5;
     chains: Record<string, Chain> = { mainnet: viemChains.mainnet };
     account: PrivateKeyAccount;
 
     constructor(
         accountOrPrivateKey: PrivateKeyAccount | `0x${string}`,
+        private cacheManager: ICacheManager,
         chains?: Record<string, Chain>
     ) {
         this.setAccount(accountOrPrivateKey);
@@ -35,6 +41,8 @@ export class WalletProvider {
         if (chains && Object.keys(chains).length > 0) {
             this.setCurrentChain(Object.keys(chains)[0] as SupportedChain);
         }
+
+        this.cache = new NodeCache({ stdTTL: this.CACHE_EXPIRY_SEC });
     }
 
     getAddress(): Address {
@@ -80,12 +88,22 @@ export class WalletProvider {
     }
 
     async getWalletBalance(): Promise<string | null> {
+        const cacheKey = "walletBalance_" + this.currentChain;
+        const cachedData = await this.getCachedData<string>(cacheKey);
+        if (cachedData) {
+            elizaLogger.log("Returning cached wallet balance for chain: " + this.currentChain);
+            return cachedData;
+        }
+
         try {
             const client = this.getPublicClient(this.currentChain);
             const balance = await client.getBalance({
                 address: this.account.address,
             });
-            return formatUnits(balance, 18);
+            const balanceFormatted = formatUnits(balance, 18);
+            this.setCachedData<string>(cacheKey, balanceFormatted);
+            elizaLogger.log("Wallet balance cached for chain: ", this.currentChain);
+            return balanceFormatted;
         } catch (error) {
             console.error("Error getting wallet balance:", error);
             return null;
@@ -122,6 +140,45 @@ export class WalletProvider {
         this.setCurrentChain(chainName);
     }
 
+    private async readFromCache<T>(key: string): Promise<T | null> {
+        const cached = await this.cacheManager.get<T>(
+            path.join(this.cacheKey, key)
+        );
+        return cached;
+    }
+
+    private async writeToCache<T>(key: string, data: T): Promise<void> {
+        await this.cacheManager.set(path.join(this.cacheKey, key), data, {
+            expires: Date.now() + this.CACHE_EXPIRY_SEC * 1000,
+        });
+    }
+
+    private async getCachedData<T>(key: string): Promise<T | null> {
+        // Check in-memory cache first
+        const cachedData = this.cache.get<T>(key);
+        if (cachedData) {
+            return cachedData;
+        }
+
+        // Check file-based cache
+        const fileCachedData = await this.readFromCache<T>(key);
+        if (fileCachedData) {
+            // Populate in-memory cache
+            this.cache.set(key, fileCachedData);
+            return fileCachedData;
+        }
+
+        return null;
+    }
+
+    private async setCachedData<T>(cacheKey: string, data: T): Promise<void> {
+        // Set in-memory cache
+        this.cache.set(cacheKey, data);
+
+        // Write to file-based cache
+        await this.writeToCache(cacheKey, data);
+    }
+
     private setAccount = (
         accountOrPrivateKey: PrivateKeyAccount | `0x${string}`
     ) => {
@@ -226,30 +283,31 @@ export const initWalletProvider = async (runtime: IAgentRuntime) => {
             walletSecretSalt,
             runtime.agentId
         );
-        return new WalletProvider(deriveKeyResult.keypair, chains);
+        return new WalletProvider(deriveKeyResult.keypair, runtime.cacheManager, chains);
     } else {
         const privateKey = runtime.getSetting(
             "EVM_PRIVATE_KEY"
         ) as `0x${string}`;
         if (!privateKey) {
             throw new Error("EVM_PRIVATE_KEY is missing");
         }
-        return new WalletProvider(privateKey, chains);
+        return new WalletProvider(privateKey, runtime.cacheManager, chains);
     }
 };
 
 export const evmWalletProvider: Provider = {
     async get(
         runtime: IAgentRuntime,
         _message: Memory,
-        _state?: State
+        state?: State
     ): Promise<string | null> {
         try {
             const walletProvider = await initWalletProvider(runtime);
             const address = walletProvider.getAddress();
             const balance = await walletProvider.getWalletBalance();
             const chain = walletProvider.getCurrentChain();
-            return `EVM Wallet Address: ${address}\nBalance: ${balance} ${chain.nativeCurrency.symbol}\nChain ID: ${chain.id}, Name: ${chain.name}`;
+            const agentName = state?.agentName || "The agent";
+            return `${agentName}'s EVM Wallet Address: ${address}\nBalance: ${balance} ${chain.nativeCurrency.symbol}\nChain ID: ${chain.id}, Name: ${chain.name}`;
         } catch (error) {
             console.error("Error in EVM wallet provider:", error);
             return null;

packages/plugin-evm/src/templates/index.ts

@@ -1,25 +1,60 @@
-export const transferTemplate = `Given the recent messages and wallet information below:
+export const transferTemplate = `You are an AI assistant specialized in processing cryptocurrency transfer requests. Your task is to extract specific information from user messages and format it into a structured JSON response.
 
+First, review the recent messages from the conversation:
+
+<recent_messages>
 {{recentMessages}}
+</recent_messages>
 
-{{walletInfo}}
+Here's a list of supported chains:
+<supported_chains>
+{{supportedChains}}
+</supported_chains>
+
+Your goal is to extract the following information about the requested transfer:
+1. Chain to execute on (must be one of the supported chains)
+2. Amount to transfer (in ETH, without the coin symbol)
+3. Recipient address (must be a valid Ethereum address)
+4. Token symbol or address (if not a native token transfer)
+
+Before providing the final JSON output, show your reasoning process inside <analysis> tags. Follow these steps:
+
+1. Identify the relevant information from the user's message:
+   - Quote the part of the message mentioning the chain.
+   - Quote the part mentioning the amount.
+   - Quote the part mentioning the recipient address.
+   - Quote the part mentioning the token (if any).
 
-Extract the following information about the requested transfer:
-- Chain to execute on (like in viem/chains)
-- Amount to transfer: Must be a string representing the amount in ETH (only number without coin symbol, e.g., "0.1")
-- Recipient address: Must be a valid Ethereum address starting with "0x"
-- Token symbol or address (if not native token): Optional, leave as null for ETH transfers
+2. Validate each piece of information:
+   - Chain: List all supported chains and check if the mentioned chain is in the list.
+   - Amount: Attempt to convert the amount to a number to verify it's valid.
+   - Address: Check that it starts with "0x" and count the number of characters (should be 42).
+   - Token: Note whether it's a native transfer or if a specific token is mentioned.
 
-Respond with a JSON markdown block containing only the extracted values. All fields except 'token' are required:
+3. If any information is missing or invalid, prepare an appropriate error message.
+
+4. If all information is valid, summarize your findings.
+
+5. Prepare the JSON structure based on your analysis.
+
+After your analysis, provide the final output in a JSON markdown block. All fields except 'token' are required. The JSON should have this structure:
 
 \`\`\`json
 {
-    "fromChain": SUPPORTED_CHAINS,
+    "fromChain": string,
     "amount": string,
     "toAddress": string,
     "token": string | null
 }
 \`\`\`
+
+Remember:
+- The chain name must be a string and must exactly match one of the supported chains.
+- The amount should be a string representing the number without any currency symbol.
+- The recipient address must be a valid Ethereum address starting with "0x".
+- If no specific token is mentioned (i.e., it's a native token transfer), set the "token" field to null.
+
+Now, process the user's request and provide your response.
 `;
 
 export const bridgeTemplate = `Given the recent messages and wallet information below:

packages/plugin-evm/src/tests/transfer.test.ts

@@ -1,18 +1,32 @@
-import { describe, it, expect, beforeEach } from "vitest";
+import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
 import { generatePrivateKey, privateKeyToAccount } from "viem/accounts";
 import { Account, Chain } from "viem";
 
 import { TransferAction } from "../actions/transfer";
 import { WalletProvider } from "../providers/wallet";
 
+// Mock the ICacheManager
+const mockCacheManager = {
+    get: vi.fn().mockResolvedValue(null),
+    set: vi.fn(),
+};
+
 describe("Transfer Action", () => {
     let wp: WalletProvider;
 
     beforeEach(async () => {
+        vi.clearAllMocks();
+        mockCacheManager.get.mockResolvedValue(null);
+
         const pk = generatePrivateKey();
         const customChains = prepareChains();
-        wp = new WalletProvider(pk, customChains);
+        wp = new WalletProvider(pk, mockCacheManager as any, customChains);
     });
+
+    afterEach(() => {
+        vi.clearAllTimers();
+    });
+
     describe("Constructor", () => {
         it("should initialize with wallet provider", () => {
             const ta = new TransferAction(wp);