diff --git a/packages/entrykit/src/config/input.ts b/packages/entrykit/src/config/input.ts index 8b6eace21e..8a21e95cfd 100644 --- a/packages/entrykit/src/config/input.ts +++ b/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"]; }; diff --git a/packages/entrykit/src/config/output.ts b/packages/entrykit/src/config/output.ts index 16495e35f7..c9229dc0f1 100644 --- a/packages/entrykit/src/config/output.ts +++ b/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"]; }; diff --git a/packages/entrykit/src/createBundlerClient.ts b/packages/entrykit/src/createBundlerClient.ts index 494b4ba2ee..e0899fd88e 100644 --- a/packages/entrykit/src/createBundlerClient.ts +++ b/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), diff --git a/packages/entrykit/src/getPaymaster.ts b/packages/entrykit/src/getPaymaster.ts index 90103ad3f0..4d265a2bfb 100644 --- a/packages/entrykit/src/getPaymaster.ts +++ b/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 { diff --git a/packages/entrykit/src/getSessionClient.ts b/packages/entrykit/src/getSessionClient.ts index bb5cccf346..4a800f6795 100644 --- a/packages/entrykit/src/getSessionClient.ts +++ b/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 { 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 diff --git a/packages/entrykit/src/onboarding/ConnectedSteps.tsx b/packages/entrykit/src/onboarding/ConnectedSteps.tsx index d098312e36..e5dc83741c 100644 --- a/packages/entrykit/src/onboarding/ConnectedSteps.tsx +++ b/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(null); const userAddress = userClient.account.address; diff --git a/packages/entrykit/src/onboarding/deposit/DepositViaRelayForm.tsx b/packages/entrykit/src/onboarding/deposit/DepositViaRelayForm.tsx index 6ff26255c1..90e387e49b 100644 --- a/packages/entrykit/src/onboarding/deposit/DepositViaRelayForm.tsx +++ b/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(); diff --git a/packages/entrykit/src/onboarding/deposit/DepositViaTransferForm.tsx b/packages/entrykit/src/onboarding/deposit/DepositViaTransferForm.tsx index 0e7de4d46f..5cbe817fed 100644 --- a/packages/entrykit/src/onboarding/deposit/DepositViaTransferForm.tsx +++ b/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"); diff --git a/packages/entrykit/src/onboarding/quarry/WithdrawGasBalanceButton.tsx b/packages/entrykit/src/onboarding/quarry/WithdrawGasBalanceButton.tsx index 23df974a6d..54b61abd7a 100644 --- a/packages/entrykit/src/onboarding/quarry/WithdrawGasBalanceButton.tsx +++ b/packages/entrykit/src/onboarding/quarry/WithdrawGasBalanceButton.tsx @@ -23,7 +23,7 @@ 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; @@ -31,7 +31,7 @@ export function WithdrawGasBalanceButton({ userAddress }: Props) { 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 { diff --git a/packages/entrykit/src/onboarding/quarry/getSpender.ts b/packages/entrykit/src/onboarding/quarry/getSpender.ts index f1722a9962..2707af5bdb 100644 --- a/packages/entrykit/src/onboarding/quarry/getSpender.ts +++ b/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, { diff --git a/packages/entrykit/src/onboarding/useSetupSession.ts b/packages/entrykit/src/onboarding/useSetupSession.ts index 969bde3c49..7a15e7d283 100644 --- a/packages/entrykit/src/onboarding/useSetupSession.ts +++ b/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 => { 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); diff --git a/packages/entrykit/src/quarry/getAllowance.ts b/packages/entrykit/src/quarry/getAllowance.ts index eafd6503f6..e36d4ee624 100644 --- a/packages/entrykit/src/quarry/getAllowance.ts +++ b/packages/entrykit/src/quarry/getAllowance.ts @@ -10,7 +10,7 @@ export type GetAllowanceParams = { }; export async function getAllowance({ client, userAddress }: GetAllowanceParams): Promise { - const paymaster = getPaymaster(client.chain); + const paymaster = getPaymaster(client.chain, undefined); if (paymaster?.type !== "quarry") return null; return await getAction( diff --git a/packages/entrykit/src/quarry/getBalance.ts b/packages/entrykit/src/quarry/getBalance.ts index 62bd587bfd..54fd6956f8 100644 --- a/packages/entrykit/src/quarry/getBalance.ts +++ b/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, { diff --git a/packages/entrykit/src/useSessionClient.ts b/packages/entrykit/src/useSessionClient.ts index 2ff40bc2dd..f2dc23b8eb 100644 --- a/packages/entrykit/src/useSessionClient.ts +++ b/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 | undefined; userAddress: Address | undefined; worldAddress: Address; + paymasterOverride: EntryKitConfig["paymasterOverride"]; }): UndefinedInitialDataOptions { return queryOptions({ queryKey: ["getSessionClient", client?.uid, userAddress, worldAddress], @@ -37,6 +40,7 @@ export function getSessionClientQueryOptions({ sessionSigner, userAddress, worldAddress, + paymasterOverride, }); }, staleTime: Infinity, @@ -47,7 +51,7 @@ export function getSessionClientQueryOptions({ export function useSessionClient(userAddress: Address | undefined): UseQueryResult { const queryClient = useQueryClient(); - const { chainId, worldAddress } = useEntryKitConfig(); + const { chainId, worldAddress, paymasterOverride } = useEntryKitConfig(); const client = useClient({ chainId }); return useQuery( getSessionClientQueryOptions({ @@ -55,6 +59,7 @@ export function useSessionClient(userAddress: Address | undefined): UseQueryResu client, userAddress, worldAddress, + paymasterOverride, }), ); }