feat(common,world): improvements for smart accounts (#2578)

Co-authored-by: mous <[email protected]> Co-authored-by: alvrs <[email protected]> Co-authored-by: karooolis <[email protected]>
latticexyz · Apr 2, 2024 · d2e4d0f · d2e4d0f
1 parent 66865c5
commit d2e4d0f
Show file tree

Hide file tree

Showing 10 changed files with 76 additions and 34 deletions.
diff --git a/.changeset/smooth-ducks-sell.md b/.changeset/smooth-ducks-sell.md
@@ -0,0 +1,9 @@
+---
+"@latticexyz/common": patch
+---
+
+`transactionQueue` decorator now accepts an optional `publicClient` argument, which will be used in place of the extended viem client for making public action calls (`getChainId`, `getTransactionCount`, `simulateContract`, `call`). This helps in cases where the extended viem client is a smart account client, like in [permissionless.js](https://github.com/pimlicolabs/permissionless.js), where the transport is the bundler, not an RPC.
+
+`writeObserver` decorator now accepts any `Client`, not just a `WalletClient`.
+
+`createBurnerAccount` now returns a `PrivateKeyAccount`, the more specific `Account` type.
diff --git a/.changeset/smooth-laws-float.md b/.changeset/smooth-laws-float.md
@@ -0,0 +1,5 @@
+---
+"@latticexyz/world": patch
+---
+
+`callFrom` decorator now accepts any `Client`, not just a `WalletClient`. It also no longer attempts to wrap/redirect calls to `call`, `callFrom`, and `registerDelegationWithSignature`.
diff --git a/packages/common/src/actions/transactionQueue.ts b/packages/common/src/actions/transactionQueue.ts
@@ -1,14 +1,20 @@
-import type { Transport, Chain, Account, WalletActions, WalletClient } from "viem";
+import type { Transport, Chain, Account, WalletActions, Client, PublicClient } from "viem";
 import { writeContract as mud_writeContract } from "../writeContract";
 import { sendTransaction as mud_sendTransaction } from "../sendTransaction";
 
-export function transactionQueue<TChain extends Chain, TAccount extends Account>(): (
-  client: WalletClient<Transport, TChain, TAccount>,
-) => Pick<WalletActions<TChain, TAccount>, "writeContract" | "sendTransaction"> {
+export type TransactionQueueOptions<chain extends Chain> = {
+  publicClient?: PublicClient<Transport, chain>;
+};
+
+export function transactionQueue<chain extends Chain, account extends Account>({
+  publicClient,
+}: TransactionQueueOptions<chain> = {}): (
+  client: Client<Transport, chain, account>,
+) => Pick<WalletActions<chain, account>, "writeContract" | "sendTransaction"> {
   return (client) => ({
     // Applies to: `client.writeContract`, `getContract(client, ...).write`
-    writeContract: (args) => mud_writeContract(client, args),
+    writeContract: (args) => mud_writeContract(client, args, publicClient),
     // Applies to: `client.sendTransaction`
-    sendTransaction: (args) => mud_sendTransaction(client, args),
+    sendTransaction: (args) => mud_sendTransaction(client, args, publicClient),
   });
 }
diff --git a/packages/common/src/actions/writeObserver.ts b/packages/common/src/actions/writeObserver.ts
@@ -4,8 +4,8 @@ import type {
   Chain,
   Account,
   WalletActions,
-  WalletClient,
   WriteContractReturnType,
+  Client,
 } from "viem";
 import { getAction } from "viem/utils";
 import { writeContract } from "viem/actions";
@@ -16,7 +16,7 @@ type WriteObserverParameters = { onWrite: (write: ContractWrite) => void };
 export function writeObserver<TChain extends Chain, TAccount extends Account>({
   onWrite,
 }: WriteObserverParameters): (
-  client: WalletClient<Transport, TChain, TAccount>,
+  client: Client<Transport, TChain, TAccount>,
 ) => Pick<WalletActions<TChain, TAccount>, "writeContract"> {
   let nextWriteId = 0;
 

diff --git a/packages/common/src/createBurnerAccount.ts b/packages/common/src/createBurnerAccount.ts
@@ -1,7 +1,7 @@
-import { Account, Hex } from "viem";
+import { Hex, PrivateKeyAccount } from "viem";
 import { privateKeyToAccount } from "viem/accounts";
 
-export function createBurnerAccount(privateKey: Hex): Account {
+export function createBurnerAccount(privateKey: Hex): PrivateKeyAccount {
   const account = privateKeyToAccount(privateKey);
   // We may override account features here
   return {

diff --git a/packages/common/src/getContract.ts b/packages/common/src/getContract.ts
@@ -118,10 +118,14 @@ export function getContract<
               TChain,
               TAccount
             >;
-            const result = writeContract(walletClient, request);
+            const result = writeContract(walletClient, request, publicClient);
 
             const id = `${walletClient.chain.id}:${walletClient.account.address}:${nextWriteId++}`;
-            onWrite?.({ id, request: request as WriteContractParameters, result });
+            onWrite?.({
+              id,
+              request: request as WriteContractParameters,
+              result,
+            });
 
             return result;
           };

diff --git a/packages/common/src/sendTransaction.ts b/packages/common/src/sendTransaction.ts
@@ -6,6 +6,7 @@ import {
   SendTransactionParameters,
   Transport,
   SendTransactionReturnType,
+  PublicClient,
 } from "viem";
 import { call, sendTransaction as viem_sendTransaction } from "viem/actions";
 import pRetry from "p-retry";
@@ -19,12 +20,13 @@ const debug = parentDebug.extend("sendTransaction");
 
 /** @deprecated Use `walletClient.extend(transactionQueue())` instead. */
 export async function sendTransaction<
-  TChain extends Chain | undefined,
-  TAccount extends Account | undefined,
-  TChainOverride extends Chain | undefined,
+  chain extends Chain | undefined,
+  account extends Account | undefined,
+  chainOverride extends Chain | undefined,
 >(
-  client: Client<Transport, TChain, TAccount>,
-  request: SendTransactionParameters<TChain, TAccount, TChainOverride>,
+  client: Client<Transport, chain, account>,
+  request: SendTransactionParameters<chain, account, chainOverride>,
+  publicClient?: PublicClient<Transport, chain>,
 ): Promise<SendTransactionReturnType> {
   const rawAccount = request.account ?? client.account;
   if (!rawAccount) {
@@ -34,23 +36,23 @@ export async function sendTransaction<
   const account = parseAccount(rawAccount);
 
   const nonceManager = await getNonceManager({
-    client,
+    client: publicClient ?? client,
     address: account.address,
     blockTag: "pending",
   });
 
-  async function prepare(): Promise<SendTransactionParameters<TChain, TAccount, TChainOverride>> {
+  async function prepare(): Promise<SendTransactionParameters<chain, account, chainOverride>> {
     if (request.gas) {
       debug("gas provided, skipping simulate", request.to);
       return request;
     }
 
     debug("simulating tx to", request.to);
-    await call(client, {
+    await call(publicClient ?? client, {
       ...request,
       blockTag: "pending",
       account,
-    } as CallParameters<TChain>);
+    } as CallParameters<chain>);
 
     return request;
   }
@@ -67,7 +69,7 @@ export async function sendTransaction<
 
           const nonce = nonceManager.nextNonce();
           debug("sending tx with nonce", nonce, "to", preparedRequest.to);
-          const parameters: SendTransactionParameters<TChain, TAccount, TChainOverride> = { nonce, ...preparedRequest };
+          const parameters: SendTransactionParameters<chain, account, chainOverride> = { nonce, ...preparedRequest };
           return await viem_sendTransaction(client, parameters);
         },
         {

diff --git a/packages/common/src/writeContract.ts b/packages/common/src/writeContract.ts
@@ -9,6 +9,7 @@ import {
   WriteContractReturnType,
   ContractFunctionName,
   ContractFunctionArgs,
+  PublicClient,
 } from "viem";
 import { simulateContract, writeContract as viem_writeContract } from "viem/actions";
 import pRetry from "p-retry";
@@ -31,6 +32,7 @@ export async function writeContract<
 >(
   client: Client<Transport, chain, account>,
   request: WriteContractParameters<abi, functionName, args, chain, account, chainOverride>,
+  publicClient?: PublicClient<Transport, chain>,
 ): Promise<WriteContractReturnType> {
   const rawAccount = request.account ?? client.account;
   if (!rawAccount) {
@@ -40,7 +42,7 @@ export async function writeContract<
   const account = parseAccount(rawAccount);
 
   const nonceManager = await getNonceManager({
-    client,
+    client: publicClient ?? client,
     address: account.address,
     blockTag: "pending",
   });
@@ -54,11 +56,14 @@ export async function writeContract<
     }
 
     debug("simulating", request.functionName, "at", request.address);
-    const result = await simulateContract<chain, account, abi, functionName, args, chainOverride>(client, {
-      ...request,
-      blockTag: "pending",
-      account,
-    } as unknown as SimulateContractParameters<abi, functionName, args, chain, chainOverride>);
+    const result = await simulateContract<chain, account | undefined, abi, functionName, args, chainOverride>(
+      publicClient ?? client,
+      {
+        ...request,
+        blockTag: "pending",
+        account,
+      } as unknown as SimulateContractParameters<abi, functionName, args, chain, chainOverride>,
+    );
 
     return result.request as unknown as WriteContractParameters<abi, functionName, args, chain, account, chainOverride>;
   }
@@ -75,7 +80,10 @@ export async function writeContract<
 
           const nonce = nonceManager.nextNonce();
           debug("calling", preparedWrite.functionName, "with nonce", nonce, "at", preparedWrite.address);
-          return await viem_writeContract(client, { nonce, ...preparedWrite } as typeof preparedWrite);
+          return await viem_writeContract(client, {
+            nonce,
+            ...preparedWrite,
+          } as typeof preparedWrite);
         },
         {
           retries: 3,

diff --git a/packages/dev-tools/src/actions/getTransactionResult.ts b/packages/dev-tools/src/actions/getTransactionResult.ts
@@ -19,7 +19,10 @@ export function getTransactionResult(
     const transaction = getTransaction(publicClient, write);
     const transactionReceipt = getTransactionReceipt(publicClient, write);
     cache[write.id] = Promise.all([transaction, transactionReceipt]).then(([tx, receipt]) => {
-      const { functionName, args } = decodeFunctionData({ abi: worldAbi, data: tx.input });
+      const { functionName, args } = decodeFunctionData({
+        abi: worldAbi,
+        data: tx.input,
+      });
       return publicClient.simulateContract({
         account: tx.from,
         address: tx.to!,

diff --git a/packages/world/ts/actions/callFrom.ts b/packages/world/ts/actions/callFrom.ts
@@ -1,7 +1,6 @@
 import {
   slice,
   concat,
-  type WalletClient,
   type Transport,
   type Chain,
   type Account,
@@ -10,6 +9,7 @@ import {
   type WriteContractReturnType,
   type EncodeFunctionDataParameters,
   type PublicClient,
+  Client,
 } from "viem";
 import { getAction, encodeFunctionData } from "viem/utils";
 import { writeContract } from "viem/actions";
@@ -52,12 +52,17 @@ type SystemFunction = { systemId: Hex; systemFunctionSelector: Hex };
 // The function mapping is cached to avoid redundant retrievals for the same World function.
 export function callFrom<TChain extends Chain, TAccount extends Account>(
   params: CallFromParameters,
-): (client: WalletClient<Transport, TChain, TAccount>) => Pick<WalletActions<TChain, TAccount>, "writeContract"> {
+): (client: Client<Transport, TChain, TAccount>) => Pick<WalletActions<TChain, TAccount>, "writeContract"> {
   return (client) => ({
     // Applies to: `client.writeContract`, `getContract(client, ...).write`
     writeContract: async (writeArgs): Promise<WriteContractReturnType> => {
-      // Skip if the contract isn't the World.
-      if (writeArgs.address !== params.worldAddress) {
+      // Skip if the contract isn't the World or the function called should not be redirected through `callFrom`.
+      if (
+        writeArgs.address !== params.worldAddress ||
+        writeArgs.functionName === "call" ||
+        writeArgs.functionName === "callFrom" ||
+        writeArgs.functionName === "registerDelegationWithSignature"
+      ) {
         return getAction(client, writeContract, "writeContract")(writeArgs);
       }