feat: add multi-chain support to account configs (#666)

packages/alchemy/src/config/actions/getAccount.ts

@@ -1,4 +1,3 @@
-import { ClientOnlyPropertyError } from "../errors.js";
 import type { AccountState } from "../store/types.js";
 import type { AlchemyAccountsConfig, SupportedAccountTypes } from "../types";
 import { type CreateAccountParams } from "./createAccount.js";
@@ -15,7 +14,10 @@ export const getAccount = <TAccount extends SupportedAccountTypes>(
 ): GetAccountResult<TAccount> => {
   const accounts = config.clientStore.getState().accounts;
   if (!accounts) {
-    throw new ClientOnlyPropertyError("account");
+    return {
+      account: undefined,
+      status: "DISCONNECTED",
+    };
   }
 
   return accounts[type];

packages/alchemy/src/config/actions/getChain.ts

@@ -0,0 +1,12 @@
+import type { Chain } from "viem";
+import type { AlchemyAccountsConfig } from "../types";
+
+/**
+ * Gets the currently active chain
+ *
+ * @param config the account config object
+ * @returns the currently active chain
+ */
+export function getChain(config: AlchemyAccountsConfig): Chain {
+  return config.coreStore.getState().chain;
+}

packages/alchemy/src/config/actions/setChain.ts

@@ -0,0 +1,25 @@
+import type { Chain } from "viem";
+import { createAlchemyPublicRpcClient } from "../../client/rpcClient.js";
+import { ChainNotFoundError } from "../errors.js";
+import type { AlchemyAccountsConfig } from "../types";
+
+/**
+ * Allows you to change the current chain in the core store
+ *
+ * @param config the accounts config object
+ * @param chain the chain to change to. It must be present in the connections config object
+ */
+export async function setChain(config: AlchemyAccountsConfig, chain: Chain) {
+  const connection = config.coreStore.getState().connections.get(chain.id);
+  if (connection == null) {
+    throw new ChainNotFoundError(chain);
+  }
+
+  config.coreStore.setState(() => ({
+    chain,
+    bundlerClient: createAlchemyPublicRpcClient({
+      chain,
+      connectionConfig: connection,
+    }),
+  }));
+}

packages/alchemy/src/config/actions/watchChain.ts

@@ -0,0 +1,14 @@
+import type { Chain } from "viem";
+import type { AlchemyAccountsConfig } from "../types";
+
+/**
+ * Allows you to subscribe to changes of the chain in the client store.
+ *
+ * @param config the account config object
+ * @returns a function which accepts an onChange callback that will be fired when the chain changes
+ */
+export function watchChain(config: AlchemyAccountsConfig) {
+  return (onChange: (chain: Chain) => void) => {
+    return config.coreStore.subscribe(({ chain }) => chain, onChange);
+  };
+}

packages/alchemy/src/config/createConfig.ts

@@ -2,7 +2,11 @@ import { ConnectionConfigSchema } from "@alchemy/aa-core";
 import { DEFAULT_SESSION_MS } from "../signer/session/manager.js";
 import { createClientStore } from "./store/client.js";
 import { createCoreStore } from "./store/core.js";
-import type { AlchemyAccountsConfig, CreateConfigProps } from "./types";
+import type {
+  AlchemyAccountsConfig,
+  Connection,
+  CreateConfigProps,
+} from "./types";
 
 export const DEFAULT_IFRAME_CONTAINER_ID = "alchemy-signer-iframe-container";
 
@@ -17,13 +21,31 @@ export const createConfig = ({
   storage,
   ...connectionConfig
 }: CreateConfigProps): AlchemyAccountsConfig => {
-  const connection = ConnectionConfigSchema.parse(connectionConfig);
+  const connections: Connection[] = [];
+  if (connectionConfig.connections != null) {
+    connectionConfig.connections.forEach(({ chain, ...config }) => {
+      connections.push({
+        ...ConnectionConfigSchema.parse(config),
+        chain,
+      });
+    });
+  } else {
+    connections.push({
+      ...ConnectionConfigSchema.parse(connectionConfig),
+      chain,
+    });
+  }
 
   const config: AlchemyAccountsConfig = {
-    coreStore: createCoreStore({ connection, chain }),
+    coreStore: createCoreStore({
+      connections,
+      chain,
+      storage: storage?.(),
+      ssr,
+    }),
     clientStore: createClientStore({
       client: {
-        connection: signerConnection ?? connection,
+        connection: signerConnection ?? connections[0],
         iframeConfig,
         rootOrgId,
         rpId,

packages/alchemy/src/config/errors.ts

@@ -1,3 +1,4 @@
+import type { Chain } from "viem";
 import { BaseError } from "../errors/base.js";
 
 export class ClientOnlyPropertyError extends BaseError {
@@ -7,3 +8,13 @@ export class ClientOnlyPropertyError extends BaseError {
     super(`${property} is only available on the client`);
   }
 }
+
+export class ChainNotFoundError extends BaseError {
+  name: string = "ChainNotFoundError";
+
+  constructor(chain: Chain) {
+    super(`Chain (${chain.name}) not found in connections config object`, {
+      docsPath: "https://accountkit.alchemy.com/react/createConfig",
+    });
+  }
+}

packages/alchemy/src/config/hydrate.ts

@@ -62,6 +62,7 @@ export function hydrate(
     async onMount() {
       if (config._internal.ssr) {
         await config.clientStore.persist.rehydrate();
+        await config.coreStore.persist.rehydrate();
       }
 
       await reconnect(config);

packages/alchemy/src/config/store/core.ts

@@ -1,30 +1,104 @@
-import { type ConnectionConfig } from "@alchemy/aa-core";
 import type { Chain } from "viem";
-import { subscribeWithSelector } from "zustand/middleware";
+import {
+  createJSONStorage,
+  persist,
+  subscribeWithSelector,
+} from "zustand/middleware";
 import { createStore } from "zustand/vanilla";
 import { createAlchemyPublicRpcClient } from "../../client/rpcClient.js";
+import type { Connection } from "../types.js";
+import { bigintMapReplacer } from "../utils/replacer.js";
+import { bigintMapReviver } from "../utils/reviver.js";
+import { DEFAULT_STORAGE_KEY } from "./client.js";
+import type { CoreState, CoreStore } from "./types.js";
 
 export type CreateCoreStoreParams = {
-  connection: ConnectionConfig;
+  connections: Connection[];
   chain: Chain;
+  storage?: Storage;
+  ssr?: boolean;
 };
 
-export const createCoreStore = ({
-  connection,
-  chain,
-}: CreateCoreStoreParams) => {
-  const bundlerClient = createAlchemyPublicRpcClient({
+/**
+ * Create the core store for alchemy accounts. This store contains the bundler client
+ * as well as the chain configs (including the initial chain to use)
+ *
+ * @param params connections configs
+ * @param params.connections a collection of chains and their connection configs
+ * @param params.chain the initial chain to use
+ * @param params.storage the storage to use for persisting the state
+ * @param params.ssr whether the store is being created on the server
+ * @returns the core store
+ */
+export const createCoreStore = (params: CreateCoreStoreParams): CoreStore => {
+  const {
+    connections,
     chain,
-    connectionConfig: connection,
-  });
+    storage = typeof window !== "undefined" ? localStorage : undefined,
+    ssr,
+  } = params;
 
   // State defined in here should work either on the server or on the client
   // bundler client for example can be used in either setting to make RPC calls
   const coreStore = createStore(
-    subscribeWithSelector(() => ({
-      bundlerClient,
-    }))
+    subscribeWithSelector(
+      storage
+        ? persist(() => createInitialCoreState(connections, chain), {
+            name: `${DEFAULT_STORAGE_KEY}:core`,
+            storage: createJSONStorage<CoreState>(() => storage, {
+              replacer: (key, value) => {
+                if (key === "bundlerClient") {
+                  const client = value as CoreState["bundlerClient"];
+                  return {
+                    connection: connections.find(
+                      (x) => x.chain.id === client.chain.id
+                    ),
+                  };
+                }
+                return bigintMapReplacer(key, value);
+              },
+              reviver: (key, value) => {
+                if (key === "bundlerClient") {
+                  const connection = value as Connection;
+                  return createAlchemyPublicRpcClient({
+                    chain: connection.chain,
+                    connectionConfig: connection,
+                  });
+                }
+
+                return bigintMapReviver(key, value);
+              },
+            }),
+            skipHydration: ssr,
+          })
+        : () => createInitialCoreState(connections, chain)
+    )
   );
 
   return coreStore;
 };
+
+const createInitialCoreState = (
+  connections: Connection[],
+  chain: Chain
+): CoreState => {
+  const connectionMap = connections.reduce((acc, connection) => {
+    acc.set(connection.chain.id, connection);
+    return acc;
+  }, new Map<number, Connection>());
+
+  if (!connectionMap.has(chain.id)) {
+    throw new Error("Chain not found in connections");
+  }
+
+  const bundlerClient = createAlchemyPublicRpcClient({
+    chain,
+    connectionConfig: connectionMap.get(chain.id)!,
+  });
+
+  return {
+    bundlerClient,
+    chain,
+    connections: connectionMap,
+  };
+};

packages/alchemy/src/config/store/types.ts

@@ -1,4 +1,4 @@
-import type { Address } from "viem";
+import type { Address, Chain } from "viem";
 import type { PartialBy } from "viem/chains";
 import type { Mutate, StoreApi } from "zustand/vanilla";
 import type { ClientWithAlchemyMethods } from "../../client/types";
@@ -10,7 +10,11 @@ import type {
   User,
 } from "../../signer";
 import type { AccountConfig } from "../actions/createAccount";
-import type { SupportedAccount, SupportedAccountTypes } from "../types";
+import type {
+  Connection,
+  SupportedAccount,
+  SupportedAccountTypes,
+} from "../types";
 
 export type AccountState<TAccount extends SupportedAccountTypes> =
   | {
@@ -71,9 +75,13 @@ export type ClientStore = Mutate<
   [["zustand/subscribeWithSelector", never], ["zustand/persist", ClientState]]
 >;
 
-export type CoreState = { bundlerClient: ClientWithAlchemyMethods };
+export type CoreState = {
+  bundlerClient: ClientWithAlchemyMethods;
+  chain: Chain;
+  connections: Map<number, Connection>;
+};
 
 export type CoreStore = Mutate<
   StoreApi<CoreState>,
-  [["zustand/subscribeWithSelector", never]]
+  [["zustand/subscribeWithSelector", never], ["zustand/persist", CoreState]]
 >;

packages/alchemy/src/config/types.ts

@@ -41,16 +41,32 @@ export type AlchemyAccountsConfig = {
 };
 
 // #region CreateConfigProps
-export type CreateConfigProps = ConnectionConfig & {
+export type Connection = ConnectionConfig & { chain: Chain };
+
+type RpcConnectionConfig =
+  | (Connection & {
+      /**
+       * Optional parameter that allows you to specify a different RPC Url
+       * or connection to be used specifically by the signer.
+       * This is useful if you have a different backend proxy for the signer
+       * than for your Bundler or Node RPC calls.
+       */
+      signerConnection?: ConnectionConfig;
+      connections?: never;
+    })
+  | {
+      connections: Connection[];
+      chain: Chain;
+      /**
+       * When providing multiple connections, you must specify the signer connection config
+       * to use since the signer is chain agnostic and has a different RPC url.
+       */
+      signerConnection: ConnectionConfig;
+    };
+
+export type CreateConfigProps = RpcConnectionConfig & {
   chain: Chain;
   sessionConfig?: AlchemySignerParams["sessionConfig"];
-  /**
-   * Optional parameter that allows you to specify a different RPC Url
-   * or connection to be used specifically by the signer.
-   * This is useful if you have a different backend proxy for the signer
-   * than for your Bundler or Node RPC calls.
-   */
-  signerConnection?: ConnectionConfig;
   /**
    * Enable this parameter if you are using the config in an SSR setting (eg. NextJS)
    * Turing this setting on will disable automatic hydration of the client store

packages/alchemy/src/react/hooks/useChain.ts

@@ -0,0 +1,44 @@
+import { useMutation } from "@tanstack/react-query";
+import { useSyncExternalStore } from "react";
+import type { Chain } from "viem";
+import { getChain } from "../../config/actions/getChain.js";
+import { setChain as setChainInternal } from "../../config/actions/setChain.js";
+import { watchChain } from "../../config/actions/watchChain.js";
+import { useAlchemyAccountContext } from "../context.js";
+import type { BaseHookMutationArgs } from "../types.js";
+
+export type UseChainParams = BaseHookMutationArgs<void, { chain: Chain }>;
+
+export interface UseChainResult {
+  chain: Chain;
+  setChain: (chain: Chain) => void;
+  isSettingChain: boolean;
+}
+
+/**
+ * A hook that returns the current chain as well as a function to set the chain
+ *
+ * @param mutationArgs the mutation arguments
+ * @returns an object containing the current chain and a function to set the chain as well as loading state of setting the chain
+ */
+export function useChain({ ...mutationArgs }: UseChainParams) {
+  const { config } = useAlchemyAccountContext();
+
+  const chain = useSyncExternalStore(
+    watchChain(config),
+    () => getChain(config),
+    () => getChain(config)
+  );
+
+  const { mutate: setChain, isPending } = useMutation({
+    mutationFn: ({ chain }: { chain: Chain }) =>
+      setChainInternal(config, chain),
+    ...mutationArgs,
+  });
+
+  return {
+    chain,
+    setChain,
+    isSettingChain: isPending,
+  };
+}

packages/alchemy/src/react/index.ts

@@ -13,6 +13,8 @@ export type * from "./hooks/useAuthenticate.js";
 export { useAuthenticate } from "./hooks/useAuthenticate.js";
 export type * from "./hooks/useBundlerClient.js";
 export { useBundlerClient } from "./hooks/useBundlerClient.js";
+export type * from "./hooks/useChain.js";
+export { useChain } from "./hooks/useChain.js";
 export type * from "./hooks/useClientActions.js";
 export { useClientActions } from "./hooks/useClientActions.js";
 export type * from "./hooks/useDropAndReplaceUserOperation.js";