feat(entrykit): add custom paymaster client override

packages/entrykit/src/config/input.ts

@@ -1,3 +1,4 @@
+import { BundlerClientConfig } from "viem/account-abstraction";
 import { Address } from "viem/accounts";
 
 export type EntryKitConfigInput = {
@@ -28,4 +29,9 @@ export type EntryKitConfigInput = {
    * Icon should be 1:1 aspect ratio, at least 200x200.
    */
   readonly appIcon?: string;
+
+  /**
+   * Custom paymaster client which overrides any paymaster address passed through the chain config.
+   */
+  readonly paymasterOverride?: BundlerClientConfig["paymaster"];
 };

packages/entrykit/src/config/output.ts

@@ -1,3 +1,4 @@
+import { BundlerClientConfig } from "viem/account-abstraction";
 import { Address } from "viem/accounts";
 
 export type EntryKitConfig = {
@@ -28,4 +29,9 @@ export type EntryKitConfig = {
    * Icon should be 1:1 aspect ratio, at least 200x200.
    */
   readonly appIcon: string;
+
+  /**
+   * Custom paymaster client which overrides any paymaster address passed through the chain config.
+   */
+  readonly paymasterOverride?: BundlerClientConfig["paymaster"];
 };

packages/entrykit/src/createBundlerClient.ts

@@ -23,18 +23,20 @@ export function createBundlerClient<
   if (!client) throw new Error("No `client` provided to `createBundlerClient`.");
 
   const chain = config.chain ?? client.chain;
-  const paymaster = chain ? getPaymaster(chain) : undefined;
+  const paymaster = chain ? getPaymaster(chain, config.paymaster) : undefined;
 
   // TODO: lift this out to make `createBundlerClient` configurable?
   return viem_createBundlerClient({
     ...defaultClientConfig,
     paymaster: paymaster
-      ? {
-          getPaymasterData: async () => ({
-            paymaster: paymaster.address,
-            paymasterData: "0x",
-          }),
-        }
+      ? paymaster.type === "custom"
+        ? paymaster.paymasterClient
+        : {
+            getPaymasterData: async () => ({
+              paymaster: paymaster.address,
+              paymasterData: "0x",
+            }),
+          }
       : undefined,
     userOperation: {
       estimateFeesPerGas: createFeeEstimator(client),

packages/entrykit/src/getPaymaster.ts

@@ -1,14 +1,33 @@
 import { Chain, Hex } from "viem";
+import { BundlerClientConfig } from "viem/account-abstraction";
+import { EntryKitConfig } from "./config/output";
 
-export type Paymaster = {
-  readonly type: "simple" | "quarry";
-  readonly address: Hex;
-  readonly canSponsor?: boolean;
-};
+export type Paymaster =
+  | {
+      readonly type: "simple" | "quarry";
+      readonly address: Hex;
+      readonly canSponsor?: boolean;
+    }
+  | {
+      readonly type: "custom";
+      readonly address?: Hex;
+      readonly canSponsor?: undefined;
+      readonly paymasterClient: BundlerClientConfig["paymaster"];
+    };
 
-export function getPaymaster(chain: Chain): Paymaster | undefined {
+export function getPaymaster(
+  chain: Chain,
+  paymasterOverride: EntryKitConfig["paymasterOverride"],
+): Paymaster | undefined {
   const contracts = chain.contracts ?? {};
 
+  if (paymasterOverride) {
+    return {
+      type: "custom",
+      paymasterClient: paymasterOverride,
+    };
+  }
+
   if ("quarryPaymaster" in contracts && contracts.quarryPaymaster != null) {
     if ("address" in contracts.quarryPaymaster) {
       return {

packages/entrykit/src/getSessionClient.ts

@@ -1,6 +1,7 @@
 import { Account, Address, Chain, Client, LocalAccount, RpcSchema, Transport } from "viem";
 import { smartAccountActions } from "permissionless";
 import { callFrom, sendUserOperationFrom } from "@latticexyz/world/internal";
+import { EntryKitConfig } from "./config/output";
 import { createBundlerClient } from "./createBundlerClient";
 import { SessionClient } from "./common";
 import { SmartAccount } from "viem/account-abstraction";
@@ -11,11 +12,13 @@ export async function getSessionClient({
   sessionAccount,
   sessionSigner,
   worldAddress,
+  paymasterOverride,
 }: {
   userAddress: Address;
   sessionAccount: SmartAccount;
   sessionSigner: LocalAccount;
   worldAddress: Address;
+  paymasterOverride: EntryKitConfig["paymasterOverride"];
 }): Promise<SessionClient> {
   const client = sessionAccount.client;
   if (!clientHasChain(client)) {
@@ -26,6 +29,7 @@ export async function getSessionClient({
     transport: getBundlerTransport(client.chain),
     client,
     account: sessionAccount,
+    paymaster: paymasterOverride,
   });
 
   const sessionClient = bundlerClient

packages/entrykit/src/onboarding/ConnectedSteps.tsx

@@ -20,8 +20,8 @@ export type Props = {
 };
 
 export function ConnectedSteps({ connector, userClient, initialUserAddress }: Props) {
-  const { chain } = useEntryKitConfig();
-  const paymaster = getPaymaster(chain);
+  const { chain, paymasterOverride } = useEntryKitConfig();
+  const paymaster = getPaymaster(chain, paymasterOverride);
   const [focusedId, setFocusedId] = useState<string | null>(null);
 
   const userAddress = userClient.account.address;

packages/entrykit/src/onboarding/deposit/DepositViaRelayForm.tsx

@@ -20,8 +20,8 @@ type Props = {
 };
 
 export function DepositViaRelayForm({ amount, setAmount, sourceChain, setSourceChainId }: Props) {
-  const { chain, chainId: destinationChainId } = useEntryKitConfig();
-  const paymaster = getPaymaster(chain);
+  const { chain, chainId: destinationChainId, paymasterOverride } = useEntryKitConfig();
+  const paymaster = getPaymaster(chain, paymasterOverride);
   const { data: wallet } = useWalletClient();
   const { address: userAddress } = useAccount();
   const { addDeposit } = useDeposits();

packages/entrykit/src/onboarding/deposit/DepositViaTransferForm.tsx

@@ -16,8 +16,8 @@ type Props = {
 };
 
 export function DepositViaTransferForm({ amount, setAmount, sourceChain, setSourceChainId }: Props) {
-  const { chain } = useEntryKitConfig();
-  const paymaster = getPaymaster(chain);
+  const { chain, paymasterOverride } = useEntryKitConfig();
+  const paymaster = getPaymaster(chain, paymasterOverride);
   const publicClient = usePublicClient();
   const { address: userAddress } = useAccount();
   const { writeContractAsync } = useWriteContract();
@@ -50,7 +50,7 @@ export function DepositViaTransferForm({ amount, setAmount, sourceChain, setSour
   const deposit = useMutation({
     mutationKey: ["depositViaTransfer", amount?.toString()],
     mutationFn: async () => {
-      if (!paymaster) throw new Error("Paymaster not found");
+      if (!paymaster?.address) throw new Error("Paymaster not found");
       if (!publicClient) throw new Error("Public client not found");
       if (!amount) throw new Error("Amount cannot be 0");
 

packages/entrykit/src/onboarding/quarry/WithdrawGasBalanceButton.tsx

@@ -23,15 +23,15 @@ export function WithdrawGasBalanceButton({ userAddress }: Props) {
   const { chainId: userChainId } = useAccount();
   const queryClient = useQueryClient();
   const client = useClient({ chainId });
-  const paymaster = getPaymaster(chain);
+  const paymaster = getPaymaster(chain, undefined);
   const balance = useShowQueryError(useBalance(userAddress));
   const shouldSwitchChain = chainId != null && chainId !== userChainId;
 
   const withdraw = useMutation({
     mutationKey: ["withdraw", userAddress],
     mutationFn: async () => {
       if (!client) throw new Error("Client not ready.");
-      if (!paymaster) throw new Error("Paymaster not found");
+      if (!paymaster?.address) throw new Error("Paymaster not found");
       if (!balance.data) throw new Error("No gas balance to withdraw.");
 
       try {

packages/entrykit/src/onboarding/quarry/getSpender.ts

@@ -10,7 +10,7 @@ export type GetSpenderParams = {
 };
 
 export async function getSpender({ client, userAddress, sessionAddress }: GetSpenderParams) {
-  const paymaster = getPaymaster(client.chain);
+  const paymaster = getPaymaster(client.chain, undefined);
   if (paymaster?.type !== "quarry") return null;
 
   const record = await getRecord(client, {

packages/entrykit/src/onboarding/useSetupSession.ts

@@ -20,7 +20,7 @@ import { storeEventsAbi } from "@latticexyz/store";
 
 export function useSetupSession({ connector, userClient }: { connector: Connector; userClient: ConnectedClient }) {
   const queryClient = useQueryClient();
-  const { chainId, worldAddress } = useEntryKitConfig();
+  const { chainId, worldAddress, paymasterOverride } = useEntryKitConfig();
   const client = useClient({ chainId });
 
   const mutationKey = ["setupSession", client?.chain.id, userClient.account.address];
@@ -37,7 +37,7 @@ export function useSetupSession({ connector, userClient }: { connector: Connecto
       registerDelegation: boolean;
     }): Promise<void> => {
       if (!client) throw new Error("Client not ready.");
-      const paymaster = getPaymaster(client.chain);
+      const paymaster = getPaymaster(client.chain, paymasterOverride);
       const sessionAddress = sessionClient.account.address;
 
       console.log("setting up session", userClient);

packages/entrykit/src/quarry/getAllowance.ts

@@ -10,7 +10,7 @@ export type GetAllowanceParams = {
 };
 
 export async function getAllowance({ client, userAddress }: GetAllowanceParams): Promise<null | bigint> {
-  const paymaster = getPaymaster(client.chain);
+  const paymaster = getPaymaster(client.chain, undefined);
   if (paymaster?.type !== "quarry") return null;
 
   return await getAction(

packages/entrykit/src/quarry/getBalance.ts

@@ -9,7 +9,7 @@ export type GetBalanceParams = {
 };
 
 export async function getBalance({ client, userAddress }: GetBalanceParams) {
-  const paymaster = getPaymaster(client.chain);
+  const paymaster = getPaymaster(client.chain, undefined);
   if (paymaster?.type !== "quarry") return null;
 
   const record = await getRecord(client, {

packages/entrykit/src/useSessionClient.ts

@@ -9,6 +9,7 @@ import {
   useQuery,
   useQueryClient,
 } from "@tanstack/react-query";
+import { EntryKitConfig } from "./config/output";
 import { getSessionClient } from "./getSessionClient";
 import { SessionClient } from "./common";
 import { getSessionAccountQueryOptions } from "./useSessionAccount";
@@ -18,11 +19,13 @@ export function getSessionClientQueryOptions({
   client,
   userAddress,
   worldAddress,
+  paymasterOverride,
 }: {
   queryClient: QueryClient;
   client: Client<Transport, Chain> | undefined;
   userAddress: Address | undefined;
   worldAddress: Address;
+  paymasterOverride: EntryKitConfig["paymasterOverride"];
 }): UndefinedInitialDataOptions<SessionClient> {
   return queryOptions<SessionClient>({
     queryKey: ["getSessionClient", client?.uid, userAddress, worldAddress],
@@ -37,6 +40,7 @@ export function getSessionClientQueryOptions({
         sessionSigner,
         userAddress,
         worldAddress,
+        paymasterOverride,
       });
     },
     staleTime: Infinity,
@@ -47,14 +51,15 @@ export function getSessionClientQueryOptions({
 
 export function useSessionClient(userAddress: Address | undefined): UseQueryResult<SessionClient> {
   const queryClient = useQueryClient();
-  const { chainId, worldAddress } = useEntryKitConfig();
+  const { chainId, worldAddress, paymasterOverride } = useEntryKitConfig();
   const client = useClient({ chainId });
   return useQuery(
     getSessionClientQueryOptions({
       queryClient,
       client,
       userAddress,
       worldAddress,
+      paymasterOverride,
     }),
   );
 }