Add support for paymasters in zk stack networks

Author: 0xaguspunk

typescript/examples/vercel-ai/abstract/.env.template

@@ -0,0 +1,4 @@
+OPENAI_API_KEY=
+WALLET_PRIVATE_KEY=
+PAYMASTER_ADDRESS=
+RPC_PROVIDER_URL=

typescript/examples/vercel-ai/abstract/README.md

@@ -0,0 +1,15 @@
+# Vercel AI Abstract Example
+
+## Setup
+
+Copy the `.env.template` and populate with your values.
+
+```
+cp .env.template .env
+```
+
+## Usage
+
+```
+npx ts-node index.ts
+```

typescript/examples/vercel-ai/abstract/index.ts

@@ -0,0 +1,43 @@
+import { openai } from "@ai-sdk/openai";
+import { generateText } from "ai";
+
+import { http } from "viem";
+import { createWalletClient } from "viem";
+import { privateKeyToAccount } from "viem/accounts";
+import { abstractTestnet } from "viem/chains";
+
+import { getOnChainTools } from "@goat-sdk/adapter-vercel-ai";
+import { PEPE, USDC, erc20 } from "@goat-sdk/plugin-erc20";
+
+import { sendETH } from "@goat-sdk/core";
+import { viem } from "@goat-sdk/wallet-viem";
+
+require("dotenv").config();
+
+const account = privateKeyToAccount(
+    process.env.WALLET_PRIVATE_KEY as `0x${string}`
+);
+
+const walletClient = createWalletClient({
+    account: account,
+    transport: http(process.env.RPC_PROVIDER_URL),
+    chain: abstractTestnet,
+});
+
+(async () => {
+    const tools = await getOnChainTools({
+        wallet: viem(walletClient, {
+            defaultPaymaster: process.env.PAYMASTER_ADDRESS as `0x${string}`,
+        }),
+        plugins: [sendETH(), erc20({ tokens: [USDC, PEPE] })],
+    });
+
+    const result = await generateText({
+        model: openai("gpt-4o-mini"),
+        tools: tools,
+        maxSteps: 5,
+        prompt: "Send 1 USDC to 0x016c0803FFC6880a9a871ba104709cDBf341A90a",
+    });
+
+    console.log(result.text);
+})();

typescript/examples/vercel-ai/abstract/package.json

@@ -0,0 +1,21 @@
+{
+    "name": "goat-examples-vercel-ai-abstract",
+    "version": "0.1.0",
+    "description": "",
+    "private": true,
+    "scripts": {
+        "test": "vitest run --passWithNoTests"
+    },
+    "author": "",
+    "license": "MIT",
+    "dependencies": {
+        "@ai-sdk/openai": "^1.0.4",
+        "@goat-sdk/adapter-vercel-ai": "workspace:*",
+        "@goat-sdk/core": "workspace:*",
+        "@goat-sdk/plugin-erc20": "workspace:*",
+        "@goat-sdk/wallet-viem": "workspace:*",
+        "ai": "^4.0.3",
+        "dotenv": "^16.4.5",
+        "viem": "2.21.49"
+    }
+}

typescript/examples/vercel-ai/abstract/tsconfig.json

@@ -0,0 +1,8 @@
+{
+    "extends": "../../../tsconfig.base.json",
+    "compilerOptions": {
+        "outDir": "dist"
+    },
+    "include": ["index.ts"],
+    "exclude": ["node_modules", "dist"]
+}

typescript/packages/core/src/wallets/evm.ts

@@ -11,6 +11,12 @@ export type EVMTransaction = {
     args?: unknown[];
     value?: bigint;
     abi?: Abi;
+    options?: EVMTransactionOptions;
+};
+
+export type EVMTransactionOptions = {
+    paymaster?: string;
+    paymasterInput?: string;
 };
 
 export type EVMReadRequest = {

typescript/packages/wallets/viem/src/index.ts

@@ -1,12 +1,38 @@
-import type { EVMReadRequest, EVMTransaction, EVMTypedData, EVMWalletClient } from "@goat-sdk/core";
-
-import { publicActions } from "viem";
-import type { WalletClient as ViemWalletClient } from "viem";
+import type {
+    EVMReadRequest,
+    EVMTransaction,
+    EVMTypedData,
+    EVMWalletClient,
+} from "@goat-sdk/core";
+import {
+    publicActions,
+    encodeFunctionData,
+    type WalletClient as ViemWalletClient,
+} from "viem";
+import { mainnet } from "viem/chains";
 import { normalize } from "viem/ens";
+import { eip712WalletActions } from "viem/zksync";
+
+export type ViemOptions = {
+    // Only used for zkSync Stack networks
+    defaultPaymaster?: string;
+    defaultPaymasterInput?: string;
+};
+
+export function viem(
+    client: ViemWalletClient,
+    options?: ViemOptions
+): EVMWalletClient {
+    const defaultPaymaster = options?.defaultPaymaster;
+    const defaultPaymasterInput = options?.defaultPaymasterInput;
 
-export function viem(client: ViemWalletClient): EVMWalletClient {
     const publicClient = client.extend(publicActions);
 
+    const waitForReceipt = async (hash: `0x${string}`) => {
+        const receipt = await publicClient.waitForTransactionReceipt({ hash });
+        return { hash: receipt.transactionHash, status: receipt.status };
+    };
+
     return {
         getAddress: () => client.account?.address ?? "",
         getChain() {
@@ -16,9 +42,8 @@ export function viem(client: ViemWalletClient): EVMWalletClient {
             };
         },
         async resolveAddress(address: string) {
-            if (/^0x[a-fA-F0-9]{40}$/.test(address)) {
+            if (/^0x[a-fA-F0-9]{40}$/.test(address))
                 return address as `0x${string}`;
-            }
 
             try {
                 const resolvedAddress = (await publicClient.getEnsAddress({
@@ -39,9 +64,7 @@ export function viem(client: ViemWalletClient): EVMWalletClient {
                 account: client.account,
             });
 
-            return {
-                signature,
-            };
+            return { signature };
         },
         async signTypedData(data: EVMTypedData) {
             if (!client.account) throw new Error("No account connected");
@@ -54,68 +77,82 @@ export function viem(client: ViemWalletClient): EVMWalletClient {
                 account: client.account,
             });
 
-            return {
-                signature,
-            };
+            return { signature };
         },
         async sendTransaction(transaction: EVMTransaction) {
-            const { to, abi, functionName, args, value } = transaction;
+            const { to, abi, functionName, args, value, options } = transaction;
+            if (!client.account) throw new Error("No account connected");
 
             const toAddress = await this.resolveAddress(to);
-            if (!client.account) throw new Error("No account connected");
 
+            const paymaster = options?.paymaster ?? defaultPaymaster;
+            const paymasterInput =
+                options?.paymasterInput ?? defaultPaymasterInput;
+            const isPaymasterTx = !!paymaster || !!paymasterInput;
+
+            // If paymaster params exist, extend with EIP712 actions
+            const sendingClient = isPaymasterTx
+                ? client.extend(eip712WalletActions())
+                : client;
+
+            // Simple ETH transfer (no ABI)
             if (!abi) {
-                const tx = await client.sendTransaction({
+                const txParams = {
                     account: client.account,
                     to: toAddress,
                     chain: client.chain,
                     value,
-                });
-
-                const transaction = await publicClient.waitForTransactionReceipt({
-                    hash: tx,
-                });
-
-                return {
-                    hash: transaction.transactionHash,
-                    status: transaction.status,
+                    ...(isPaymasterTx ? { paymaster, paymasterInput } : {}),
                 };
+
+                const txHash = await sendingClient.sendTransaction(txParams);
+                return waitForReceipt(txHash);
             }
 
+            // Contract call
             if (!functionName) {
-                throw new Error("Function name is required");
+                throw new Error("Function name is required for contract calls");
             }
 
-            await publicClient.simulateContract({
+            const { request } = await publicClient.simulateContract({
                 account: client.account,
                 address: toAddress,
-                abi,
+                abi: abi,
                 functionName,
                 args,
                 chain: client.chain,
             });
-            const hash = await client.writeContract({
-                account: client.account,
-                address: toAddress,
-                abi,
+
+            // Encode the call data ourselves
+            const data = encodeFunctionData({
+                abi: abi,
                 functionName,
                 args,
-                chain: client.chain,
-                value,
             });
 
-            const t = await publicClient.waitForTransactionReceipt({
-                hash: hash,
-            });
+            if (isPaymasterTx) {
+                // With paymaster, we must use sendTransaction() directly
+                const txParams = {
+                    account: client.account,
+                    chain: client.chain,
+                    to: request.address,
+                    data,
+                    value: request.value,
+                    paymaster,
+                    paymasterInput,
+                };
+                const txHash = await sendingClient.sendTransaction(txParams);
+                return waitForReceipt(txHash);
+            }
 
-            return {
-                hash: t.transactionHash,
-                status: t.status,
-            };
+            // Without paymaster, use writeContract which handles encoding too,
+            // but since we already have request, let's let writeContract do its thing.
+            // However, writeContract expects the original request format (with abi, functionName, args).
+            const txHash = await client.writeContract(request);
+            return waitForReceipt(txHash);
         },
         async read(request: EVMReadRequest) {
             const { address, abi, functionName, args } = request;
-
             if (!abi) throw new Error("Read request must include ABI for EVM");
 
             const result = await publicClient.readContract({
@@ -125,22 +162,21 @@ export function viem(client: ViemWalletClient): EVMWalletClient {
                 args,
             });
 
-            return {
-                value: result,
-            };
+            return { value: result };
         },
         async balanceOf(address: string) {
             const resolvedAddress = await this.resolveAddress(address);
-
             const balance = await publicClient.getBalance({
                 address: resolvedAddress,
             });
 
+            const chain = client.chain ?? mainnet;
+
             return {
                 value: balance,
-                decimals: 18,
-                symbol: "ETH",
-                name: "Ether",
+                decimals: chain.nativeCurrency.decimals,
+                symbol: chain.nativeCurrency.symbol,
+                name: chain.nativeCurrency.name,
             };
         },
     };