Merge pull request #1874 from cprussin/move-staking-client-to-context

chore(staking): initialize pyth staking client only once

Author: cprussin

apps/staking/src/api.ts

@@ -1,8 +1,9 @@
+import type { PythStakingClient } from "@pythnetwork/staking-sdk";
+import { PublicKey } from "@solana/web3.js";
 import clsx from "clsx";
 import { useMemo, useCallback } from "react";
 
 import {
-  type Context,
   delegateIntegrityStaking,
   cancelWarmupIntegrityStaking,
   unstakeIntegrityStaking,
@@ -127,7 +128,7 @@ export const OracleIntegrityStaking = ({
             <tbody className="bg-white/5">
               {otherPublishers.map((publisher) => (
                 <Publisher
-                  key={publisher.publicKey}
+                  key={publisher.publicKey.toBase58()}
                   availableToStake={availableToStake}
                   publisher={publisher}
                   totalStaked={staked}
@@ -152,7 +153,7 @@ type PublisherProps = {
   isSelf?: boolean;
   publisher: {
     name: string;
-    publicKey: string;
+    publicKey: PublicKey;
     isSelf: boolean;
     selfStake: bigint;
     poolCapacity: bigint;
@@ -372,7 +373,7 @@ const PublisherTableCell = Styled("td", "py-4 px-5 whitespace-nowrap");
 
 type StakeToPublisherButtonProps = {
   publisherName: string;
-  publisherKey: string;
+  publisherKey: PublicKey;
   availableToStake: bigint;
   poolCapacity: bigint;
   poolUtilization: bigint;
@@ -423,13 +424,15 @@ const StakeToPublisherButton = ({
 
 const useTransferActionForPublisher = (
   action: (
-    context: Context,
-    publicKey: string,
+    client: PythStakingClient,
+    stakingAccount: PublicKey,
+    publisher: PublicKey,
     amount: bigint,
   ) => Promise<void>,
-  publicKey: string,
+  publisher: PublicKey,
 ) =>
   useCallback(
-    (context: Context, amount: bigint) => action(context, publicKey, amount),
-    [action, publicKey],
+    (client: PythStakingClient, stakingAccount: PublicKey, amount: bigint) =>
+      action(client, stakingAccount, publisher, amount),
+    [action, publisher],
   );

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

@@ -1,4 +1,6 @@
 import { Field, Input, Label } from "@headlessui/react";
+import type { PythStakingClient } from "@pythnetwork/staking-sdk";
+import type { PublicKey } from "@solana/web3.js";
 import {
   type ChangeEvent,
   type ComponentProps,
@@ -8,7 +10,6 @@ import {
   useState,
 } from "react";
 
-import type { Context } from "../../api";
 import { useLogger } from "../../hooks/use-logger";
 import { StateType, useTransfer } from "../../hooks/use-transfer";
 import { stringToTokens, tokensToString } from "../../tokens";
@@ -28,7 +29,11 @@ type Props = {
     | ReactNode
     | ReactNode[]
     | undefined;
-  transfer: (context: Context, amount: bigint) => Promise<void>;
+  transfer: (
+    client: PythStakingClient,
+    stakingAccount: PublicKey,
+    amount: bigint,
+  ) => Promise<void>;
   className?: string | undefined;
   secondary?: boolean | undefined;
   small?: boolean | undefined;
@@ -51,9 +56,9 @@ export const TransferButton = ({
   const { amountInput, setAmount, updateAmount, resetAmount, amount } =
     useAmountInput(max);
   const doTransfer = useCallback(
-    (context: Context) =>
+    (client: PythStakingClient, stakingAccount: PublicKey) =>
       amount.type === AmountType.Valid
-        ? transfer(context, amount.amount)
+        ? transfer(client, stakingAccount, amount.amount)
         : Promise.reject(new InvalidAmountError()),
     [amount, transfer],
   );

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

@@ -12,7 +12,7 @@ const demand = (key: string): string => {
   if (value && value !== "") {
     return value;
   } else {
-    throw new Error(`Missing environment variable ${key}!`);
+    throw new MissingEnvironmentError(key);
   }
 };
 
@@ -36,3 +36,10 @@ export const WALLETCONNECT_PROJECT_ID = demandInProduction(
 );
 export const RPC = process.env.RPC;
 export const IS_MAINNET = process.env.IS_MAINNET !== undefined;
+
+class MissingEnvironmentError extends Error {
+  constructor(name: string) {
+    super(`Missing environment variable: ${name}!`);
+    this.name = "MissingEnvironmentError";
+  }
+}

apps/staking/src/config/server.ts

@@ -1,23 +1,22 @@
+import { PublicKey } from "@solana/web3.js";
 import useSWR from "swr";
 
-import { useApiContext } from "./use-api-context";
+import { useSelectedStakeAccount } from "./use-stake-account";
 import { loadAccountHistory } from "../api";
 
 const ONE_SECOND_IN_MS = 1000;
 const ONE_MINUTE_IN_MS = 60 * ONE_SECOND_IN_MS;
 const REFRESH_INTERVAL = 1 * ONE_MINUTE_IN_MS;
 
-export const getCacheKey = ({
-  stakeAccount,
-}: ReturnType<typeof useApiContext>) =>
-  `${stakeAccount.address.toBase58()}/history`;
+export const getCacheKey = (stakeAccount: PublicKey) =>
+  `${stakeAccount.toBase58()}/history`;
 
 export const useAccountHistory = () => {
-  const apiContext = useApiContext();
+  const { client, account } = useSelectedStakeAccount();
 
   const { data, isLoading, ...rest } = useSWR(
-    getCacheKey(apiContext),
-    () => loadAccountHistory(apiContext),
+    getCacheKey(account.address),
+    () => loadAccountHistory(client, account.address),
     {
       refreshInterval: REFRESH_INTERVAL,
     },

apps/staking/src/hooks/use-account-history.ts

@@ -1,23 +1,22 @@
+import { PublicKey } from "@solana/web3.js";
 import { useCallback } from "react";
 import useSWR from "swr";
 
-import { useApiContext } from "./use-api-context";
+import { useSelectedStakeAccount } from "./use-stake-account";
 import { loadData } from "../api";
 
 const ONE_SECOND_IN_MS = 1000;
 const ONE_MINUTE_IN_MS = 60 * ONE_SECOND_IN_MS;
 const REFRESH_INTERVAL = 1 * ONE_MINUTE_IN_MS;
 
-export const getCacheKey = ({
-  stakeAccount,
-}: ReturnType<typeof useApiContext>) => stakeAccount.address.toBase58();
+export const getCacheKey = (stakeAccount: PublicKey) => stakeAccount.toBase58();
 
 export const useDashboardData = () => {
-  const apiContext = useApiContext();
+  const { client, account } = useSelectedStakeAccount();
 
   const { data, isLoading, mutate, ...rest } = useSWR(
-    getCacheKey(apiContext),
-    () => loadData(apiContext),
+    getCacheKey(account.address),
+    () => loadData(client, account),
     {
       refreshInterval: REFRESH_INTERVAL,
     },

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

@@ -1,6 +1,7 @@
 "use client";
 
 import type { StakeAccountPositions } from "@pythnetwork/staking-sdk";
+import { PythStakingClient } from "@pythnetwork/staking-sdk";
 import { useConnection, useWallet } from "@solana/wallet-adapter-react";
 import {
   type ComponentProps,
@@ -25,21 +26,36 @@ export enum StateType {
 
 const State = {
   Initialized: () => ({ type: StateType.Initialized as const }),
+
   NoWallet: () => ({ type: StateType.NoWallet as const }),
+
   Loading: () => ({ type: StateType.Loading as const }),
-  NoAccounts: () => ({ type: StateType.NoAccounts as const }),
+
+  NoAccounts: (client: PythStakingClient) => ({
+    type: StateType.NoAccounts as const,
+    client,
+  }),
+
   Loaded: (
+    client: PythStakingClient,
     account: StakeAccountPositions,
     allAccounts: [StakeAccountPositions, ...StakeAccountPositions[]],
     selectAccount: (account: StakeAccountPositions) => void,
   ) => ({
     type: StateType.Loaded as const,
+    client,
     account,
     allAccounts,
     selectAccount,
   }),
-  ErrorState: (error: LoadStakeAccountError, reset: () => void) => ({
+
+  ErrorState: (
+    client: PythStakingClient,
+    error: LoadStakeAccountError,
+    reset: () => void,
+  ) => ({
     type: StateType.Error as const,
+    client,
     error,
     reset,
   }),
@@ -67,7 +83,7 @@ const useStakeAccountState = () => {
     (account: StakeAccountPositions) => {
       setState((cur) =>
         cur.type === StateType.Loaded
-          ? State.Loaded(account, cur.allAccounts, setAccount)
+          ? State.Loaded(cur.client, account, cur.allAccounts, setAccount)
           : cur,
       );
     },
@@ -85,27 +101,34 @@ const useStakeAccountState = () => {
       ) {
         throw new WalletConnectedButInvalidError();
       }
-      getStakeAccounts(connection, {
-        publicKey: wallet.publicKey,
-        signAllTransactions: wallet.signAllTransactions,
-        signTransaction: wallet.signTransaction,
-      })
+      const client = new PythStakingClient({
+        connection,
+        wallet: {
+          publicKey: wallet.publicKey,
+          signAllTransactions: wallet.signAllTransactions,
+          signTransaction: wallet.signTransaction,
+        },
+      });
+      getStakeAccounts(client)
         .then((accounts) => {
           const [firstAccount, ...otherAccounts] = accounts;
           if (firstAccount) {
             setState(
               State.Loaded(
+                client,
                 firstAccount,
                 [firstAccount, ...otherAccounts],
                 setAccount,
               ),
             );
           } else {
-            setState(State.NoAccounts());
+            setState(State.NoAccounts(client));
           }
         })
         .catch((error: unknown) => {
-          setState(State.ErrorState(new LoadStakeAccountError(error), reset));
+          setState(
+            State.ErrorState(client, new LoadStakeAccountError(error), reset),
+          );
         })
         .finally(() => {
           loading.current = false;
@@ -129,6 +152,15 @@ export const useStakeAccount = () => {
   }
 };
 
+export const useSelectedStakeAccount = () => {
+  const state = useStakeAccount();
+  if (state.type === StateType.Loaded) {
+    return state;
+  } else {
+    throw new InvalidStateError();
+  }
+};
+
 class LoadStakeAccountError extends Error {
   constructor(cause: unknown) {
     super(cause instanceof Error ? cause.message : "");
@@ -152,3 +184,12 @@ class WalletConnectedButInvalidError extends Error {
     );
   }
 }
+
+class InvalidStateError extends Error {
+  constructor() {
+    super(
+      "Cannot use `useSelectedStakeAccount` when stake accounts aren't loaded or a stake account isn't selected!  Ensure this hook is only called when a stake account is selected.",
+    );
+    this.name = "InvalidStateError";
+  }
+}