Skip to content

Commit

Permalink
refactor: switch alchemy client configs to use alchemy transport
Browse files Browse the repository at this point in the history
  • Loading branch information
moldy530 committed Sep 27, 2024
1 parent bacf9bd commit affecef
Show file tree
Hide file tree
Showing 36 changed files with 712 additions and 479 deletions.
7 changes: 2 additions & 5 deletions aa-sdk/core/src/client/smartAccountClient.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { custom, type Transaction } from "viem";
import { polygonMumbai } from "viem/chains";
import type { SpyInstance } from "vitest";
import type { MockInstance } from "vitest";
import * as receiptActions from "../actions/bundler/getUserOperationReceipt.js";
import type { UserOperationReceipt } from "../types.js";
import {
Expand Down Expand Up @@ -110,10 +110,7 @@ describe("SmartAccountClient Tests", async () => {
const thenExpectRetriesToBe = async (
expectedRetryMsDelays: number[],
expectedMockCalls: number,
getUserOperationReceiptMock: SpyInstance<
any,
Promise<UserOperationReceipt | null>
>
getUserOperationReceiptMock: MockInstance
) => {
expect(retryMsDelays).toEqual(expectedRetryMsDelays);
expect(getUserOperationReceiptMock).toHaveBeenCalledTimes(
Expand Down
8 changes: 8 additions & 0 deletions account-kit/core/src/actions/getAlchemyTransport.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import type { AlchemyTransport } from "@account-kit/infra";
import type { AlchemyAccountsConfig } from "../types";

export function getAlchemyTransport(
config: AlchemyAccountsConfig
): AlchemyTransport {
return config.store.getState().transport;
}
42 changes: 15 additions & 27 deletions account-kit/core/src/actions/getSmartAccountClient.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import {
createAlchemySmartAccountClientFromExisting,
createAlchemySmartAccountClient,
type AlchemySmartAccountClient,
type AlchemySmartAccountClientConfig,
} from "@account-kit/infra";
Expand All @@ -19,7 +19,7 @@ import {
type MultiOwnerPluginActions,
type PluginManagerActions,
} from "@account-kit/smart-contracts";
import type { Address, Chain, Transport } from "viem";
import type { Address, Chain } from "viem";
import type {
AlchemyAccountsConfig,
Connection,
Expand All @@ -29,21 +29,16 @@ import type {
} from "../types";
import { createAccount } from "./createAccount.js";
import { getAccount, type GetAccountParams } from "./getAccount.js";
import { getBundlerClient } from "./getBundlerClient.js";
import { getAlchemyTransport } from "./getAlchemyTransport.js";
import { getConnection } from "./getConnection.js";
import { getSignerStatus } from "./getSignerStatus.js";

export type GetSmartAccountClientParams<
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined,
TAccount extends SupportedAccountTypes = SupportedAccountTypes
> = Omit<
AlchemySmartAccountClientConfig<
TTransport,
TChain,
SupportedAccount<TAccount>
>,
"rpcUrl" | "chain" | "apiKey" | "jwt" | "account"
AlchemySmartAccountClientConfig<TChain, SupportedAccount<TAccount>>,
"transport" | "account"
> &
GetAccountParams<TAccount>;

Expand All @@ -60,29 +55,22 @@ export type ClientActions<
: never;

export type GetSmartAccountClientResult<
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined,
TAccount extends SupportedAccounts = SupportedAccounts
> = {
client?: AlchemySmartAccountClient<
TTransport,
TChain,
TAccount,
ClientActions<TAccount>
>;
client?: AlchemySmartAccountClient<TChain, TAccount, ClientActions<TAccount>>;
address?: Address;
isLoadingClient: boolean;
error?: Error;
};

export function getSmartAccountClient<
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined,
TAccount extends SupportedAccountTypes = SupportedAccountTypes
>(
params: GetSmartAccountClientParams<TTransport, TChain, TAccount>,
params: GetSmartAccountClientParams<TChain, TAccount>,
config: AlchemyAccountsConfig
): GetSmartAccountClientResult<TTransport, TChain, SupportedAccount<TAccount>>;
): GetSmartAccountClientResult<TChain, SupportedAccount<TAccount>>;

/**
* Obtains a smart account client based on the provided parameters and configuration. Supports creating any of the SupportAccountTypes in Account Kit.
Expand Down Expand Up @@ -116,7 +104,7 @@ export function getSmartAccountClient(
config
);
const signerStatus = getSignerStatus(config);
const bundlerClient = getBundlerClient(config);
const transport = getAlchemyTransport(config);
const connection = getConnection(config);
const clientState =
config.store.getState().smartAccountClients[connection.chain.id]?.[type];
Expand Down Expand Up @@ -189,8 +177,8 @@ export function getSmartAccountClient(
switch (account.source) {
case "LightAccount":
return {
client: createAlchemySmartAccountClientFromExisting({
client: bundlerClient,
client: createAlchemySmartAccountClient({
transport,
account: account,
policyId: connection.policyId,
...clientParams,
Expand All @@ -200,8 +188,8 @@ export function getSmartAccountClient(
};
case "MultiOwnerLightAccount":
return {
client: createAlchemySmartAccountClientFromExisting({
client: bundlerClient,
client: createAlchemySmartAccountClient({
transport,
account: account,
policyId: connection.policyId,
...clientParams,
Expand All @@ -211,8 +199,8 @@ export function getSmartAccountClient(
};
case "MultiOwnerModularAccount":
return {
client: createAlchemySmartAccountClientFromExisting({
client: bundlerClient,
client: createAlchemySmartAccountClient({
transport,
account: account,
policyId: connection.policyId,
...clientParams,
Expand Down
6 changes: 4 additions & 2 deletions account-kit/core/src/actions/setChain.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { createAlchemyPublicRpcClient } from "@account-kit/infra";
import { alchemy, createAlchemyPublicRpcClient } from "@account-kit/infra";
import { switchChain } from "@wagmi/core";
import type { Chain } from "viem";
import { ChainNotFoundError } from "../errors.js";
Expand Down Expand Up @@ -27,12 +27,14 @@ export async function setChain(config: AlchemyAccountsConfig, chain: Chain) {
}

await switchChain(config._internal.wagmiConfig, { chainId: chain.id });
const transport = alchemy(connection);

config.store.setState(() => ({
chain,
transport,
bundlerClient: createAlchemyPublicRpcClient({
transport,
chain,
connectionConfig: connection,
}),
}));
}
15 changes: 5 additions & 10 deletions account-kit/core/src/actions/watchSmartAccountClient.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import type { Chain, Transport } from "viem";
import type { Chain } from "viem";
import { ClientOnlyPropertyError } from "../errors.js";
import type {
AlchemyAccountsConfig,
Expand Down Expand Up @@ -27,24 +27,19 @@ import {
* @template TTransport extends Transport = Transport
* @template TChain extends Chain | undefined = Chain | undefined
*
* @param {GetSmartAccountClientParams<TTransport, TChain, TAccount>} params the parameters needed to get the smart account client
* @param {GetSmartAccountClientParams<TChain, TAccount>} params the parameters needed to get the smart account client
* @param {AlchemyAccountsConfig} config the configuration containing the client store and other settings
* @returns {(onChange: (client: GetSmartAccountClientResult<TTransport, TChain, SupportedAccount<TAccount>>) => void) => (() => void)} a function that accepts a callback to be called when the client changes and returns a function to unsubscribe from the store
* @returns {(onChange: (client: GetSmartAccountClientResult<TChain, SupportedAccount<TAccount>>) => void) => (() => void)} a function that accepts a callback to be called when the client changes and returns a function to unsubscribe from the store
*/ export function watchSmartAccountClient<
TAccount extends SupportedAccountTypes,
TTransport extends Transport = Transport,
TChain extends Chain | undefined = Chain | undefined
>(
params: GetSmartAccountClientParams<TTransport, TChain, TAccount>,
params: GetSmartAccountClientParams<TChain, TAccount>,
config: AlchemyAccountsConfig
) {
return (
onChange: (
client: GetSmartAccountClientResult<
TTransport,
TChain,
SupportedAccount<TAccount>
>
client: GetSmartAccountClientResult<TChain, SupportedAccount<TAccount>>
) => void
) => {
const accounts = config.store.getState().accounts;
Expand Down
1 change: 0 additions & 1 deletion account-kit/core/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,6 @@ export {
export type * from "@account-kit/infra";
export {
createAlchemySmartAccountClient,
createAlchemySmartAccountClientFromExisting,
type AlchemySmartAccountClient,
type AlchemySmartAccountClientConfig,
} from "@account-kit/infra";
23 changes: 20 additions & 3 deletions account-kit/core/src/store/store.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import type { NoUndefined } from "@aa-sdk/core";
import { createAlchemyPublicRpcClient } from "@account-kit/infra";
import { alchemy, createAlchemyPublicRpcClient } from "@account-kit/infra";
import { AlchemySignerStatus, AlchemyWebSigner } from "@account-kit/signer";
import type { Chain } from "viem";
import {
Expand Down Expand Up @@ -48,17 +48,32 @@ export const createAccountKitStore = (
),
};
}

if (key === "transport") {
const transport = value as StoreState["transport"];
return {
connection: connections.find(
(x) => x.chain.id === transport({}).value?.chain.id
),
};
}

return bigintMapReplacer(key, value);
},
reviver: (key, value) => {
if (key === "bundlerClient") {
const { connection } = value as { connection: Connection };
return createAlchemyPublicRpcClient({
transport: alchemy(connection),
chain: connection.chain,
connectionConfig: connection,
});
}

if (key === "transport") {
const { connection } = value as { connection: Connection };
return alchemy(connection);
}

return bigintMapReviver(key, value);
},
}),
Expand Down Expand Up @@ -124,13 +139,15 @@ const createInitialStoreState = (
throw new Error("Chain not found in connections");
}

const transport = alchemy(connectionMap.get(chain.id)!);
const bundlerClient = createAlchemyPublicRpcClient({
transport,
chain,
connectionConfig: connectionMap.get(chain.id)!,
});
const chains = connections.map((c) => c.chain);
const accountConfigs = createEmptyAccountConfigState(chains);
const baseState: StoreState = {
transport,
bundlerClient,
chain,
connections: connectionMap,
Expand Down
9 changes: 6 additions & 3 deletions account-kit/core/src/store/types.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
import type { ClientWithAlchemyMethods } from "@account-kit/infra";
import type {
AlchemyTransport,
ClientWithAlchemyMethods,
} from "@account-kit/infra";
import type {
AlchemySignerParams,
AlchemySignerStatus,
Expand All @@ -7,7 +10,7 @@ import type {
User,
} from "@account-kit/signer";
import type { State as WagmiState } from "@wagmi/core";
import type { Address, Chain, Transport } from "viem";
import type { Address, Chain } from "viem";
import type { PartialBy } from "viem/chains";
import type { Mutate, StoreApi } from "zustand/vanilla";
import type { AccountConfig } from "../actions/createAccount";
Expand Down Expand Up @@ -87,7 +90,6 @@ export type StoreState = {
smartAccountClients: {
[chain: number]: Partial<{
[key in SupportedAccountTypes]: GetSmartAccountClientResult<
Transport,
Chain,
SupportedAccount<key>
>;
Expand All @@ -96,6 +98,7 @@ export type StoreState = {
// serializable state
// NOTE: in some cases this can be serialized to cookie storage
// be mindful of how big this gets. cookie limit 4KB
transport: AlchemyTransport;
bundlerClient: ClientWithAlchemyMethods;
config: ClientStoreConfig;
accountConfigs: {
Expand Down
93 changes: 93 additions & 0 deletions account-kit/infra/src/__snapshots__/alchemyTransport.test.ts.snap
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html

exports[`Alchemy Transport Tests > should correctly create a split transport 1`] = `
[
{
"fallback": [Function],
"overrides": [
{
"methods": [
"eth_sendUserOperation",
"eth_estimateUserOperationGas",
"eth_getUserOperationReceipt",
"eth_getUserOperationByHash",
"eth_supportedEntryPoints",
"rundler_maxPriorityFeePerGas",
"pm_getPaymasterData",
"pm_getPaymasterStubData",
],
"transport": [Function],
},
],
},
]
`;

exports[`Alchemy Transport Tests > should correctly create a split transport 2`] = `
[
{
"fallback": [Function],
"overrides": [
{
"methods": [
"eth_sendUserOperation",
"eth_estimateUserOperationGas",
"eth_getUserOperationReceipt",
"eth_getUserOperationByHash",
"eth_supportedEntryPoints",
"rundler_maxPriorityFeePerGas",
"pm_getPaymasterData",
"pm_getPaymasterStubData",
],
"transport": [Function],
},
],
},
]
`;

exports[`Alchemy Transport Tests > should correctly create a split transport 3`] = `
[
{
"fallback": [Function],
"overrides": [
{
"methods": [
"eth_sendUserOperation",
"eth_estimateUserOperationGas",
"eth_getUserOperationReceipt",
"eth_getUserOperationByHash",
"eth_supportedEntryPoints",
"rundler_maxPriorityFeePerGas",
"pm_getPaymasterData",
"pm_getPaymasterStubData",
],
"transport": [Function],
},
],
},
]
`;

exports[`Alchemy Transport Tests > should correctly create a split transport 4`] = `
[
{
"fallback": [Function],
"overrides": [
{
"methods": [
"eth_sendUserOperation",
"eth_estimateUserOperationGas",
"eth_getUserOperationReceipt",
"eth_getUserOperationByHash",
"eth_supportedEntryPoints",
"rundler_maxPriorityFeePerGas",
"pm_getPaymasterData",
"pm_getPaymasterStubData",
],
"transport": [Function],
},
],
},
]
`;
Loading

0 comments on commit affecef

Please sign in to comment.