feat(staking): reload stake account positions properly (#1914)

Author: keyvankhademi

apps/staking/src/api.ts

@@ -151,15 +151,21 @@ export type AccountHistory = {
   locked: bigint;
 }[];
 
-export const getStakeAccounts = async (
+export const getAllStakeAccountAddresses = async (
   client: PythStakingClient,
-): Promise<StakeAccountPositions[]> =>
+): Promise<PublicKey[]> =>
   client.getAllStakeAccountPositions(client.wallet.publicKey);
 
+export const getStakeAccount = async (
+  client: PythStakingClient,
+  stakeAccountAddress: PublicKey,
+): Promise<StakeAccountPositions> =>
+  client.getStakeAccountPositions(stakeAccountAddress);
+
 export const loadData = async (
   client: PythStakingClient,
   hermesClient: HermesClient,
-  stakeAccount?: StakeAccountPositions | undefined,
+  stakeAccount?: PublicKey | undefined,
 ): Promise<Data> =>
   stakeAccount === undefined
     ? loadDataNoStakeAccount(client, hermesClient)
@@ -191,23 +197,25 @@ const loadDataNoStakeAccount = async (
 const loadDataForStakeAccount = async (
   client: PythStakingClient,
   hermesClient: HermesClient,
-  stakeAccount: StakeAccountPositions,
+  stakeAccount: PublicKey,
 ) => {
   const [
     { publishers, ...baseInfo },
     stakeAccountCustody,
     unlockSchedule,
     claimableRewards,
+    stakeAccountPositions,
   ] = await Promise.all([
     loadBaseInfo(client, hermesClient),
-    client.getStakeAccountCustody(stakeAccount.address),
-    client.getUnlockSchedule(stakeAccount.address),
-    client.getClaimableRewards(stakeAccount.address),
+    client.getStakeAccountCustody(stakeAccount),
+    client.getUnlockSchedule(stakeAccount),
+    client.getClaimableRewards(stakeAccount),
+    client.getStakeAccountPositions(stakeAccount),
   ]);
 
   const filterGovernancePositions = (positionState: PositionState) =>
     getAmountByTargetAndState({
-      stakeAccountPositions: stakeAccount,
+      stakeAccountPositions,
       targetWithParameters: { voting: {} },
       positionState,
       epoch: baseInfo.currentEpoch,
@@ -218,7 +226,7 @@ const loadDataForStakeAccount = async (
     positionState: PositionState,
   ) =>
     getAmountByTargetAndState({
-      stakeAccountPositions: stakeAccount,
+      stakeAccountPositions,
       targetWithParameters: { integrityPool: { publisher } },
       positionState,
       epoch: baseInfo.currentEpoch,
@@ -330,7 +338,7 @@ export const loadAccountHistory = async (
 export const createStakeAccountAndDeposit = async (
   _client: PythStakingClient,
   _amount: bigint,
-): Promise<StakeAccountPositions> => {
+): Promise<PublicKey> => {
   await new Promise((resolve) => setTimeout(resolve, MOCK_DELAY));
   throw new NotImplementedError();
 };

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

@@ -13,7 +13,7 @@ import { Tokens } from "../Tokens";
 type Props = { api: States[ApiStateType.Loaded] };
 
 export const AccountHistory = ({ api }: Props) => {
-  const history = useData(api.accountHisoryCacheKey, api.loadAccountHistory);
+  const history = useData(api.accountHistoryCacheKey, api.loadAccountHistory);
 
   switch (history.type) {
     case StateType.NotLoaded:

apps/staking/src/components/Header/current-stake-account.tsx

@@ -18,10 +18,10 @@ export const CurrentStakeAccount = ({
       <div className="flex flex-col items-end text-xs md:flex-row md:items-baseline md:text-sm">
         <div className="font-semibold">Stake account:</div>
         <CopyButton
-          text={api.account.address.toBase58()}
+          text={api.account.toBase58()}
           className="-mr-2 ml-2 text-pythpurple-400 md:mr-0"
         >
-          <TruncatedKey>{api.account.address}</TruncatedKey>
+          <TruncatedKey>{api.account}</TruncatedKey>
         </CopyButton>
       </div>
     </div>

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

@@ -76,7 +76,7 @@ export const OracleIntegrityStaking = ({
     () =>
       api.type === ApiStateType.Loaded &&
       publishers.find((publisher) =>
-        publisher.stakeAccount?.equals(api.account.address),
+        publisher.stakeAccount?.equals(api.account),
       ),
     [publishers, api],
   );

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

@@ -136,7 +136,7 @@ const ConnectedButton = ({
                   <StyledMenu
                     items={api.allAccounts.map((account) => ({
                       account,
-                      id: account.address.toBase58(),
+                      id: account.toBase58(),
                     }))}
                   >
                     {(item) => (
@@ -154,7 +154,7 @@ const ConnectedButton = ({
                             invisible: item.account !== api.account,
                           })}
                         />
-                        <TruncatedKey>{item.account.address}</TruncatedKey>
+                        <TruncatedKey>{item.account}</TruncatedKey>
                       </WalletMenuItem>
                     )}
                   </StyledMenu>

apps/staking/src/hooks/use-api.tsx

@@ -1,10 +1,7 @@
 "use client";
 
 import { HermesClient } from "@pythnetwork/hermes-client";
-import {
-  type StakeAccountPositions,
-  PythStakingClient,
-} from "@pythnetwork/staking-sdk";
+import { PythStakingClient } from "@pythnetwork/staking-sdk";
 import { useConnection, useWallet } from "@solana/wallet-adapter-react";
 import type { PublicKey } from "@solana/web3.js";
 import {
@@ -46,7 +43,7 @@ const State = {
   [StateType.LoadedNoStakeAccount]: (
     client: PythStakingClient,
     hermesClient: HermesClient,
-    onCreateAccount: (newAccount: StakeAccountPositions) => void,
+    onCreateAccount: (newAccount: PublicKey) => void,
   ) => ({
     type: StateType.LoadedNoStakeAccount as const,
     client,
@@ -62,18 +59,18 @@ const State = {
   [StateType.Loaded]: (
     client: PythStakingClient,
     hermesClient: HermesClient,
-    account: StakeAccountPositions,
-    allAccounts: [StakeAccountPositions, ...StakeAccountPositions[]],
-    selectAccount: (account: StakeAccountPositions) => void,
+    account: PublicKey,
+    allAccounts: [PublicKey, ...PublicKey[]],
+    selectAccount: (account: PublicKey) => void,
     mutate: ReturnType<typeof useSWRConfig>["mutate"],
   ) => {
-    const dashboardDataCacheKey = account.address.toBase58();
-    const accountHisoryCacheKey = `${account.address.toBase58()}/history`;
+    const dashboardDataCacheKey = account.toBase58();
+    const accountHistoryCacheKey = `${account.toBase58()}/history`;
 
     const reload = async () => {
       await Promise.all([
         mutate(dashboardDataCacheKey),
-        mutate(accountHisoryCacheKey),
+        mutate(accountHistoryCacheKey),
       ]);
     };
 
@@ -86,7 +83,7 @@ const State = {
         ) => Promise<void>,
       ) =>
       async (...args: T) => {
-        await fn(client, account.address, ...args);
+        await fn(client, account, ...args);
         await reload();
       };
 
@@ -98,10 +95,10 @@ const State = {
       allAccounts,
       selectAccount,
       dashboardDataCacheKey,
-      accountHisoryCacheKey,
+      accountHistoryCacheKey,
 
       loadData: () => api.loadData(client, hermesClient, account),
-      loadAccountHistory: () => api.loadAccountHistory(client, account.address),
+      loadAccountHistory: () => api.loadAccountHistory(client, account),
 
       claim: bindApi(api.claim),
       deposit: bindApi(api.deposit),
@@ -157,7 +154,7 @@ const useApiContext = (hermesUrl: string) => {
   const { mutate } = useSWRConfig();
 
   const setAccount = useCallback(
-    (account: StakeAccountPositions) => {
+    (account: PublicKey) => {
       setState((cur) =>
         cur.type === StateType.Loaded
           ? State[StateType.Loaded](
@@ -195,7 +192,7 @@ const useApiContext = (hermesUrl: string) => {
       const hermesClient = new HermesClient(hermesUrl);
       setState(State[StateType.LoadingStakeAccounts](client, hermesClient));
       api
-        .getStakeAccounts(client)
+        .getAllStakeAccountAddresses(client)
         .then((accounts) => {
           const [firstAccount, ...otherAccounts] = accounts;
           if (firstAccount) {

governance/pyth_staking_sdk/src/pyth-staking-client.ts

@@ -92,7 +92,7 @@ export class PythStakingClient {
   /** Gets a users stake accounts */
   public async getAllStakeAccountPositions(
     user: PublicKey,
-  ): Promise<StakeAccountPositions[]> {
+  ): Promise<PublicKey[]> {
     const positionDataMemcmp = this.stakingProgram.coder.accounts.memcmp(
       "positionData",
     ) as {
@@ -117,13 +117,7 @@ export class PythStakingClient {
           ],
         },
       );
-    return res.map((account) =>
-      deserializeStakeAccountPositions(
-        account.pubkey,
-        account.account.data,
-        this.stakingProgram.idl,
-      ),
-    );
+    return res.map((account) => account.pubkey);
   }
 
   public async getStakeAccountPositions(
@@ -458,37 +452,41 @@ export class PythStakingClient {
             ),
           })
           .instruction(),
-      ));
+      ),
+    );
 
-      const mergePositionsInstruction = await Promise.all(
-        publishers.map(({ pubkey }) =>
-          this.integrityPoolProgram.methods
-            .mergeDelegationPositions()
-            .accounts({
-              owner: this.wallet.publicKey,
-              publisher: pubkey,
-              stakeAccountPositions,
-            })
-            .instruction(),
-        ),
-      );
+    const mergePositionsInstruction = await Promise.all(
+      publishers.map(({ pubkey }) =>
+        this.integrityPoolProgram.methods
+          .mergeDelegationPositions()
+          .accounts({
+            owner: this.wallet.publicKey,
+            publisher: pubkey,
+            stakeAccountPositions,
+          })
+          .instruction(),
+      ),
+    );
 
-      return {
-        advanceDelegationRecordInstructions,
-        mergePositionsInstruction,
-      };
+    return {
+      advanceDelegationRecordInstructions,
+      mergePositionsInstruction,
+    };
   }
 
   public async advanceDelegationRecord(stakeAccountPositions: PublicKey) {
     const instructions = await this.getAdvanceDelegationRecordInstructions(
       stakeAccountPositions,
     );
 
-
-    return sendTransaction([
-      ...instructions.advanceDelegationRecordInstructions,
-      ...instructions.mergePositionsInstruction,
-    ], this.connection, this.wallet);
+    return sendTransaction(
+      [
+        ...instructions.advanceDelegationRecordInstructions,
+        ...instructions.mergePositionsInstruction,
+      ],
+      this.connection,
+      this.wallet,
+    );
   }
 
   public async getClaimableRewards(stakeAccountPositions: PublicKey) {