fix(staking): use anchor wallet for API interfaces (#1869)

Author: cprussin

apps/staking/src/api.ts

@@ -1,7 +1,7 @@
 // TODO remove these disables when moving off the mock APIs
 /* eslint-disable @typescript-eslint/no-unused-vars, @typescript-eslint/no-non-null-assertion */
 
-import type { WalletContextState } from "@solana/wallet-adapter-react";
+import type { AnchorWallet } from "@solana/wallet-adapter-react";
 import type { Connection } from "@solana/web3.js";
 
 export type StakeAccount = {
@@ -10,7 +10,7 @@ export type StakeAccount = {
 
 export type Context = {
   connection: Connection;
-  wallet: WalletContextState;
+  wallet: AnchorWallet;
   stakeAccount: StakeAccount;
 };
 
@@ -141,7 +141,7 @@ type AccountHistory = {
 
 export const getStakeAccounts = async (
   _connection: Connection,
-  _wallet: WalletContextState,
+  _wallet: AnchorWallet,
 ): Promise<StakeAccount[]> => {
   await new Promise((resolve) => setTimeout(resolve, MOCK_DELAY));
   return MOCK_STAKE_ACCOUNTS;

apps/staking/src/components/WalletButton/index.tsx

@@ -31,9 +31,11 @@ import {
   type SVGAttributes,
   type ReactNode,
   type ElementType,
+  type Ref,
   useCallback,
   useMemo,
   useState,
+  forwardRef,
 } from "react";
 
 import { usePrimaryDomain } from "../../hooks/use-primary-domain";
@@ -191,27 +193,26 @@ type WalletMenuItemProps<T extends ElementType> = Omit<
   icon?: ComponentType<SVGAttributes<SVGSVGElement>>;
 };
 
-const WalletMenuItem = <T extends ElementType>({
-  as,
-  children,
-  icon: Icon,
-  className,
-  ...props
-}: WalletMenuItemProps<T>) => {
+const WalletMenuItemImpl = <T extends ElementType>(
+  { as, children, icon: Icon, className, ...props }: WalletMenuItemProps<T>,
+  ref: Ref<HTMLButtonElement>,
+) => {
   const Component = as ?? "button";
   return (
     <Component
       className={clsx(
         "flex items-center gap-2 whitespace-nowrap px-4 py-2 text-left hover:bg-pythpurple-800/20 data-[focus]:bg-pythpurple-800/20",
         className,
       )}
+      ref={ref}
       {...props}
     >
       {Icon && <Icon className="size-4 text-pythpurple-600" />}
       {children}
     </Component>
   );
 };
+const WalletMenuItem = forwardRef(WalletMenuItemImpl);
 
 const DisconnectedButton = (props: Props) => {
   const modal = useWalletModal();

apps/staking/src/hooks/use-api-context.ts

@@ -1,14 +1,18 @@
-import { useWallet, useConnection } from "@solana/wallet-adapter-react";
+import { useAnchorWallet, useConnection } from "@solana/wallet-adapter-react";
 import { useMemo } from "react";
 
 import { StateType, useStakeAccount } from "./use-stake-account";
 import type { Context } from "../api";
 
 export const useApiContext = (): Context => {
-  const wallet = useWallet();
+  const wallet = useAnchorWallet();
   const { connection } = useConnection();
   const stakeAccount = useStakeAccount();
 
+  if (wallet === undefined) {
+    throw new NoWalletConnectedError();
+  }
+
   if (stakeAccount.type !== StateType.Loaded) {
     throw new NoStakeAccountSelectedError();
   }
@@ -19,6 +23,14 @@ export const useApiContext = (): Context => {
   );
 };
 
+class NoWalletConnectedError extends Error {
+  constructor() {
+    super(
+      "The `useApiContext` hook cannot be called if a wallet isn't connected!  Ensure all components that use this hook are only rendered if a wallet is connected!",
+    );
+  }
+}
+
 class NoStakeAccountSelectedError extends Error {
   constructor() {
     super(

apps/staking/src/hooks/use-stake-account.tsx

@@ -75,7 +75,18 @@ const useStakeAccountState = () => {
     if (wallet.connected && !wallet.disconnecting && !loading.current) {
       loading.current = true;
       setState(State.Loading());
-      getStakeAccounts(connection, wallet)
+      if (
+        !wallet.publicKey ||
+        !wallet.signAllTransactions ||
+        !wallet.signTransaction
+      ) {
+        throw new WalletConnectedButInvalidError();
+      }
+      getStakeAccounts(connection, {
+        publicKey: wallet.publicKey,
+        signAllTransactions: wallet.signAllTransactions,
+        signTransaction: wallet.signTransaction,
+      })
         .then((accounts) => {
           const [firstAccount, ...otherAccounts] = accounts;
           if (firstAccount) {
@@ -96,12 +107,10 @@ const useStakeAccountState = () => {
         .finally(() => {
           loading.current = false;
         });
-    } else if (!wallet.connected) {
-      setState(State.NoWallet());
     }
   }, [connection, setAccount, wallet]);
 
-  return state;
+  return wallet.connected && !wallet.disconnecting ? state : State.NoWallet();
 };
 
 export const useStakeAccount = () => {
@@ -120,3 +129,11 @@ class NotInitializedError extends Error {
     );
   }
 }
+
+class WalletConnectedButInvalidError extends Error {
+  constructor() {
+    super(
+      "The wallet is connected but is missing a public key or methods to sign transactions!",
+    );
+  }
+}