feat: add batch operations support for transactions and enhance wallet details retrieval

d4mr · d4mr · commit f24e84b0ca7a · 2025-01-13T13:15:41.000+05:30
diff --git a/src/server/routes/backend-wallet/send-transactions-atomic.ts b/src/server/routes/backend-wallet/send-transactions-atomic.ts
@@ -0,0 +1,151 @@
+import { Type, type Static } from "@sinclair/typebox";
+import type { FastifyInstance } from "fastify";
+import { StatusCodes } from "http-status-codes";
+import type { Address, Hex } from "thirdweb";
+import { insertTransaction } from "../../../shared/utils/transaction/insert-transaction";
+import { AddressSchema } from "../../schemas/address";
+import {
+  requestQuerystringSchema,
+  standardResponseSchema,
+  transactionWritesResponseSchema,
+} from "../../schemas/shared-api-schemas";
+import {
+  maybeAddress,
+  walletChainParamSchema,
+  walletWithAAHeaderSchema,
+} from "../../schemas/wallet";
+import { getChainIdFromChain } from "../../utils/chain";
+import {
+  getWalletDetails,
+  isSmartBackendWallet,
+  type ParsedWalletDetails,
+  WalletDetailsError,
+} from "../../../shared/db/wallets/get-wallet-details";
+import { createCustomError } from "../../middleware/error";
+
+const requestBodySchema = Type.Object({
+  transactions: Type.Array(
+    Type.Object({
+      toAddress: Type.Optional(AddressSchema),
+      data: Type.String({
+        examples: ["0x..."],
+      }),
+      value: Type.String({
+        examples: ["10000000"],
+      }),
+    }),
+  ),
+});
+
+export async function sendTransactionsAtomicRoute(fastify: FastifyInstance) {
+  fastify.route<{
+    Params: Static<typeof walletChainParamSchema>;
+    Body: Static<typeof requestBodySchema>;
+    Reply: Static<typeof transactionWritesResponseSchema>;
+    Querystring: Static<typeof requestQuerystringSchema>;
+  }>({
+    method: "POST",
+    url: "/backend-wallet/:chain/send-transactions-atomic",
+    schema: {
+      summary: "Send a batch of raw transactions atomically",
+      description:
+        "Send a batch of raw transactions in a single UserOp. Can only be used with smart wallets.",
+      tags: ["Backend Wallet"],
+      operationId: "sendTransactionsAtomic",
+      params: walletChainParamSchema,
+      body: requestBodySchema,
+      headers: Type.Omit(walletWithAAHeaderSchema, ["x-transaction-mode"]),
+      querystring: requestQuerystringSchema,
+      response: {
+        ...standardResponseSchema,
+        [StatusCodes.OK]: transactionWritesResponseSchema,
+      },
+    },
+    handler: async (request, reply) => {
+      const { chain } = request.params;
+      const {
+        "x-backend-wallet-address": fromAddress,
+        "x-idempotency-key": idempotencyKey,
+        "x-account-address": accountAddress,
+        "x-account-factory-address": accountFactoryAddress,
+        "x-account-salt": accountSalt,
+      } = request.headers as Static<typeof walletWithAAHeaderSchema>;
+      const chainId = await getChainIdFromChain(chain);
+      const shouldSimulate = request.query.simulateTx ?? false;
+      const transactionRequests = request.body.transactions;
+
+      const hasSmartHeaders = !!accountAddress;
+
+      // check that we either use SBW, or send using EOA with smart wallet headers
+      if (!hasSmartHeaders) {
+        let backendWallet: ParsedWalletDetails | undefined;
+
+        try {
+          backendWallet = await getWalletDetails({
+            address: fromAddress,
+          });
+        } catch (e: unknown) {
+          if (e instanceof WalletDetailsError) {
+            throw createCustomError(
+              `Failed to get wallet details for backend wallet ${fromAddress}. ${e.message}`,
+              StatusCodes.BAD_REQUEST,
+              "WALLET_DETAILS_ERROR",
+            );
+          }
+        }
+
+        if (!backendWallet) {
+          throw createCustomError(
+            "Failed to get wallet details for backend wallet. See: https://portal.thirdweb.com/engine/troubleshooting",
+            StatusCodes.INTERNAL_SERVER_ERROR,
+            "WALLET_DETAILS_ERROR",
+          );
+        }
+
+        if (!isSmartBackendWallet(backendWallet)) {
+          throw createCustomError(
+            "Backend wallet is not a smart wallet, and x-account-address is not provided. Either use a smart backend wallet or provide x-account-address. This endpoint can only be used with smart wallets.",
+            StatusCodes.BAD_REQUEST,
+            "BACKEND_WALLET_NOT_SMART",
+          );
+        }
+      }
+
+      if (transactionRequests.length === 0) {
+        throw createCustomError(
+          "No transactions provided",
+          StatusCodes.BAD_REQUEST,
+          "NO_TRANSACTIONS_PROVIDED",
+        );
+      }
+
+      const queueId = await insertTransaction({
+        insertedTransaction: {
+          transactionMode: undefined,
+          isUserOp: false,
+          chainId,
+          from: fromAddress as Address,
+          accountAddress: maybeAddress(accountAddress, "x-account-address"),
+          accountFactoryAddress: maybeAddress(
+            accountFactoryAddress,
+            "x-account-factory-address",
+          ),
+          accountSalt: accountSalt,
+          batchOperations: transactionRequests.map((transactionRequest) => ({
+            to: transactionRequest.toAddress as Address | undefined,
+            data: transactionRequest.data as Hex,
+            value: BigInt(transactionRequest.value),
+          })),
+        },
+        shouldSimulate,
+        idempotencyKey,
+      });
+
+      reply.status(StatusCodes.OK).send({
+        result: {
+          queueId,
+        },
+      });
+    },
+  });
+}
diff --git a/src/server/routes/index.ts b/src/server/routes/index.ts
@@ -112,6 +112,7 @@ import { getAllWebhooksData } from "./webhooks/get-all";
 import { revokeWebhook } from "./webhooks/revoke";
 import { testWebhookRoute } from "./webhooks/test";
 import { readMulticallRoute } from "./contract/read/read-batch";
+import { sendTransactionsAtomicRoute } from "./backend-wallet/send-transactions-atomic";
 
 export async function withRoutes(fastify: FastifyInstance) {
   // Backend Wallets
@@ -125,6 +126,7 @@ export async function withRoutes(fastify: FastifyInstance) {
   await fastify.register(withdraw);
   await fastify.register(sendTransaction);
   await fastify.register(sendTransactionBatch);
+  await fastify.register(sendTransactionsAtomicRoute);
   await fastify.register(signTransaction);
   await fastify.register(signMessageRoute);
   await fastify.register(signTypedData);
diff --git a/src/server/schemas/transaction/index.ts b/src/server/schemas/transaction/index.ts
@@ -210,6 +210,16 @@ export const TransactionSchema = Type.Object({
     }),
     Type.Null(),
   ]),
+  batchOperations: Type.Union([
+    Type.Array(
+      Type.Object({
+        to: Type.Union([Type.String(), Type.Null()]),
+        data: Type.Union([Type.String(), Type.Null()]),
+        value: Type.String(),
+      }),
+    ),
+    Type.Null(),
+  ]),
 });
 
 export const toTransactionSchema = (
@@ -255,6 +265,17 @@ export const toTransactionSchema = (
     return null;
   };
 
+  const resolveBatchOperations = (): Static<
+    typeof TransactionSchema
+  >["batchOperations"] => {
+    if (!transaction.batchOperations) return null;
+    return transaction.batchOperations.map((op) => ({
+      to: op.to ?? null,
+      data: op.data ?? null,
+      value: op.value.toString(),
+    }));
+  };
+
   const resolveGas = (): string | null => {
     if (transaction.status === "sent") {
       return transaction.gas.toString();
@@ -351,6 +372,8 @@ export const toTransactionSchema = (
     userOpHash:
       "userOpHash" in transaction ? (transaction.userOpHash as Hex) : null,
 
+    batchOperations: resolveBatchOperations(),
+
     // Deprecated
     retryGasValues: null,
     retryMaxFeePerGas: null,
diff --git a/src/shared/db/wallets/get-wallet-details.ts b/src/shared/db/wallets/get-wallet-details.ts
@@ -1,3 +1,4 @@
+import LruMap from "mnemonist/lru-map";
 import { getAddress } from "thirdweb";
 import { z } from "zod";
 import type { PrismaTransaction } from "../../schemas/prisma";
@@ -130,6 +131,7 @@ export type SmartBackendWalletType = (typeof SmartBackendWalletTypes)[number];
 export type BackendWalletType = (typeof BackendWalletTypes)[number];
 export type ParsedWalletDetails = z.infer<typeof walletDetailsSchema>;
 
+export const walletDetailsCache = new LruMap<string, ParsedWalletDetails>(2048);
 /**
  * Return the wallet details for the given address.
  *
@@ -143,20 +145,26 @@ export type ParsedWalletDetails = z.infer<typeof walletDetailsSchema>;
  */
 export const getWalletDetails = async ({
   pgtx,
-  address,
+  address: _walletAddress,
 }: GetWalletDetailsParams) => {
+  const walletAddress = _walletAddress.toLowerCase();
+  const cachedDetails = walletDetailsCache.get(walletAddress);
+  if (cachedDetails) {
+    return cachedDetails;
+  }
+
   const prisma = getPrismaWithPostgresTx(pgtx);
   const config = await getConfig();
 
   const walletDetails = await prisma.walletDetails.findUnique({
     where: {
-      address: address.toLowerCase(),
+      address: walletAddress.toLowerCase(),
     },
   });
 
   if (!walletDetails) {
     throw new WalletDetailsError(
-      `No wallet details found for address ${address}`,
+      `No wallet details found for address ${walletAddress}`,
     );
   }
 
@@ -167,7 +175,7 @@ export const getWalletDetails = async ({
   ) {
     if (!walletDetails.awsKmsArn) {
       throw new WalletDetailsError(
-        `AWS KMS ARN is missing for the wallet with address ${address}`,
+        `AWS KMS ARN is missing for the wallet with address ${walletAddress}`,
       );
     }
 
@@ -188,7 +196,7 @@ export const getWalletDetails = async ({
   ) {
     if (!walletDetails.gcpKmsResourcePath) {
       throw new WalletDetailsError(
-        `GCP KMS resource path is missing for the wallet with address ${address}`,
+        `GCP KMS resource path is missing for the wallet with address ${walletAddress}`,
       );
     }
 
@@ -209,14 +217,16 @@ export const getWalletDetails = async ({
 
   // zod schema can validate all necessary fields are populated after decryption
   try {
-    return walletDetailsSchema.parse(walletDetails, {
+    const result = walletDetailsSchema.parse(walletDetails, {
       errorMap: (issue) => {
         const fieldName = issue.path.join(".");
         return {
-          message: `${fieldName} is necessary for wallet ${address} of type ${walletDetails.type}, but not found in wallet details or configuration`,
+          message: `${fieldName} is necessary for wallet ${walletAddress} of type ${walletDetails.type}, but not found in wallet details or configuration`,
         };
       },
     });
+    walletDetailsCache.set(walletAddress, result);
+    return result;
   } catch (e) {
     if (e instanceof z.ZodError) {
       throw new WalletDetailsError(
diff --git a/src/shared/utils/transaction/insert-transaction.ts b/src/shared/utils/transaction/insert-transaction.ts
@@ -4,6 +4,7 @@ import { TransactionDB } from "../../../shared/db/transactions/db";
 import {
   getWalletDetails,
   isSmartBackendWallet,
+  WalletDetailsError,
   type ParsedWalletDetails,
 } from "../../../shared/db/wallets/get-wallet-details";
 import { doesChainSupportService } from "../../lib/chain/chain-capabilities";
@@ -105,8 +106,12 @@ export const insertTransaction = async (
         entrypointAddress: walletDetails.entrypointAddress ?? undefined,
       };
     }
-  } catch {
-    // if wallet details are not found, this is a smart backend wallet using a v4 endpoint
+  } catch (e) {
+    if (e instanceof WalletDetailsError) {
+      // do nothing. The this is a smart backend wallet using a v4 endpoint
+    }
+    // if other type of error, rethrow
+    throw e;
   }
 
   if (!walletDetails && queuedTransaction.accountAddress) {
@@ -139,13 +144,16 @@ export const insertTransaction = async (
             walletDetails.accountFactoryAddress ?? undefined,
         };
       }
-    } catch {
+    } catch (e: unknown) {
       // if wallet details are not found for this either, this backend wallet does not exist at all
-      throw createCustomError(
-        "Account not found",
-        StatusCodes.BAD_REQUEST,
-        "ACCOUNT_NOT_FOUND",
-      );
+      if (e instanceof WalletDetailsError) {
+        throw createCustomError(
+          "Account not found",
+          StatusCodes.BAD_REQUEST,
+          "ACCOUNT_NOT_FOUND",
+        );
+      }
+      throw e;
     }
   }
 
diff --git a/src/shared/utils/transaction/types.ts b/src/shared/utils/transaction/types.ts
@@ -13,6 +13,14 @@ export type AnyTransaction =
   | CancelledTransaction
   | ErroredTransaction;
 
+export type BatchOperation = {
+  to?: Address;
+  value: bigint;
+  data?: Hex;
+  functionName?: string;
+  functionArgs?: unknown[];
+};
+
 // InsertedTransaction is the raw input from the caller.
 export type InsertedTransaction = {
   isUserOp: boolean;
@@ -49,6 +57,7 @@ export type InsertedTransaction = {
   accountSalt?: string;
   accountFactoryAddress?: Address;
   entrypointAddress?: Address;
+  batchOperations?: BatchOperation[];
   target?: Address;
   sender?: Address;
 };
diff --git a/src/worker/tasks/send-transaction-worker.ts b/src/worker/tasks/send-transaction-worker.ts
