From 1b37e0869726dd271743e049bedfb40e432c27eb Mon Sep 17 00:00:00 2001
From: Connor Prussin <connor@prussin.net>
Date: Wed, 11 Sep 2024 12:47:19 -0700
Subject: [PATCH] fix(staking): use the sdk to drive epoch timestamps

---
 apps/staking/src/api.ts                       | 39 ++++++-------------
 .../src/components/Dashboard/index.tsx        |  4 ++
 .../src/components/Governance/index.tsx       |  3 ++
 .../OracleIntegrityStaking/index.tsx          | 20 +++++++++-
 .../src/components/ProgramSection/index.tsx   | 14 ++++---
 .../src/components/StakingTimeline/index.tsx  | 12 +++---
 6 files changed, 52 insertions(+), 40 deletions(-)

diff --git a/apps/staking/src/api.ts b/apps/staking/src/api.ts
index ee5e717c8..3e4d4084f 100644
--- a/apps/staking/src/api.ts
+++ b/apps/staking/src/api.ts
@@ -26,6 +26,7 @@ const publishersRankingSchema = z
 type Data = {
   total: bigint;
   availableRewards: bigint;
+  currentEpoch: bigint;
   lastSlash:
     | {
         amount: bigint;
@@ -202,13 +203,11 @@ const loadDataForStakeAccount = async (
     stakeAccountCustody,
     unlockSchedule,
     claimableRewards,
-    currentEpoch,
   ] = await Promise.all([
     loadBaseInfo(client, hermesClient),
     client.getStakeAccountCustody(stakeAccount.address),
     client.getUnlockSchedule(stakeAccount.address),
     client.getClaimableRewards(stakeAccount.address),
-    getCurrentEpoch(client.connection),
   ]);
 
   const filterGovernancePositions = (positionState: PositionState) =>
@@ -216,7 +215,7 @@ const loadDataForStakeAccount = async (
       stakeAccountPositions: stakeAccount,
       targetWithParameters: { voting: {} },
       positionState,
-      epoch: currentEpoch,
+      epoch: baseInfo.currentEpoch,
     });
 
   const filterOISPositions = (
@@ -227,7 +226,7 @@ const loadDataForStakeAccount = async (
       stakeAccountPositions: stakeAccount,
       targetWithParameters: { integrityPool: { publisher } },
       positionState,
-      epoch: currentEpoch,
+      epoch: baseInfo.currentEpoch,
     });
 
   return {
@@ -271,13 +270,15 @@ const loadBaseInfo = async (
   client: PythStakingClient,
   hermesClient: HermesClient,
 ) => {
-  const [publishers, walletAmount, poolConfig] = await Promise.all([
-    loadPublisherData(client, hermesClient),
-    client.getOwnerPythBalance(),
-    client.getPoolConfigAccount(),
-  ]);
-
-  return { yieldRate: poolConfig.y, walletAmount, publishers };
+  const [publishers, walletAmount, poolConfig, currentEpoch] =
+    await Promise.all([
+      loadPublisherData(client, hermesClient),
+      client.getOwnerPythBalance(),
+      client.getPoolConfigAccount(),
+      getCurrentEpoch(client.connection),
+    ]);
+
+  return { yieldRate: poolConfig.y, walletAmount, publishers, currentEpoch };
 };
 
 const loadPublisherData = async (
@@ -454,22 +455,6 @@ export const optPublisherOut = async (
   throw new NotImplementedError();
 };
 
-export const getUpcomingEpoch = (): Date => {
-  const d = new Date();
-  d.setUTCDate(d.getUTCDate() + ((5 + 7 - d.getUTCDay()) % 7 || 7));
-  d.setUTCHours(0);
-  d.setUTCMinutes(0);
-  d.setUTCSeconds(0);
-  d.setUTCMilliseconds(0);
-  return d;
-};
-
-export const getNextFullEpoch = (): Date => {
-  const d = getUpcomingEpoch();
-  d.setUTCDate(d.getUTCDate() + 7);
-  return d;
-};
-
 const MOCK_DELAY = 500;
 
 const mkMockHistory = (): AccountHistory => [
diff --git a/apps/staking/src/components/Dashboard/index.tsx b/apps/staking/src/components/Dashboard/index.tsx
index 677c5fb18..169f3ecd7 100644
--- a/apps/staking/src/components/Dashboard/index.tsx
+++ b/apps/staking/src/components/Dashboard/index.tsx
@@ -9,6 +9,7 @@ import { Styled } from "../Styled";
 
 type Props = {
   api: States[ApiStateType.Loaded] | States[ApiStateType.LoadedNoStakeAccount];
+  currentEpoch: bigint;
   total: bigint;
   lastSlash:
     | {
@@ -42,6 +43,7 @@ type Props = {
 
 export const Dashboard = ({
   api,
+  currentEpoch,
   total,
   lastSlash,
   walletAmount,
@@ -150,6 +152,7 @@ export const Dashboard = ({
         <DashboardTabPanel id={TabIds.Governance}>
           <Governance
             api={api}
+            currentEpoch={currentEpoch}
             availableToStake={availableToStakeGovernance}
             warmup={governance.warmup}
             staked={governance.staked}
@@ -160,6 +163,7 @@ export const Dashboard = ({
         <DashboardTabPanel id={TabIds.IntegrityStaking}>
           <OracleIntegrityStaking
             api={api}
+            currentEpoch={currentEpoch}
             availableToStake={availableToStakeIntegrity}
             locked={locked}
             warmup={integrityStakingWarmup}
diff --git a/apps/staking/src/components/Governance/index.tsx b/apps/staking/src/components/Governance/index.tsx
index 27f0c6bcf..8783bca2e 100644
--- a/apps/staking/src/components/Governance/index.tsx
+++ b/apps/staking/src/components/Governance/index.tsx
@@ -3,6 +3,7 @@ import { ProgramSection } from "../ProgramSection";
 
 type Props = {
   api: States[ApiStateType.Loaded] | States[ApiStateType.LoadedNoStakeAccount];
+  currentEpoch: bigint;
   availableToStake: bigint;
   warmup: bigint;
   staked: bigint;
@@ -12,6 +13,7 @@ type Props = {
 
 export const Governance = ({
   api,
+  currentEpoch,
   availableToStake,
   warmup,
   staked,
@@ -21,6 +23,7 @@ export const Governance = ({
   <ProgramSection
     name="Governance"
     description="Vote and Influence the Network"
+    currentEpoch={currentEpoch}
     available={availableToStake}
     warmup={warmup}
     staked={staked}
diff --git a/apps/staking/src/components/OracleIntegrityStaking/index.tsx b/apps/staking/src/components/OracleIntegrityStaking/index.tsx
index 8a8ce3abb..ba422aeeb 100644
--- a/apps/staking/src/components/OracleIntegrityStaking/index.tsx
+++ b/apps/staking/src/components/OracleIntegrityStaking/index.tsx
@@ -47,6 +47,7 @@ const PAGE_SIZE = 10;
 
 type Props = {
   api: States[ApiStateType.Loaded] | States[ApiStateType.LoadedNoStakeAccount];
+  currentEpoch: bigint;
   availableToStake: bigint;
   locked: bigint;
   warmup: bigint;
@@ -59,6 +60,7 @@ type Props = {
 
 export const OracleIntegrityStaking = ({
   api,
+  currentEpoch,
   availableToStake,
   locked,
   warmup,
@@ -88,6 +90,7 @@ export const OracleIntegrityStaking = ({
       name="Oracle Integrity Staking (OIS)"
       description="Protect DeFi"
       className="pb-0 sm:pb-0"
+      currentEpoch={currentEpoch}
       available={availableToStake}
       warmup={warmup}
       staked={staked}
@@ -129,6 +132,7 @@ export const OracleIntegrityStaking = ({
                 <Publisher
                   api={api}
                   isSelf
+                  currentEpoch={currentEpoch}
                   availableToStake={availableToStake}
                   publisher={self}
                   totalStaked={staked}
@@ -147,6 +151,7 @@ export const OracleIntegrityStaking = ({
       >
         <PublisherList
           api={api}
+          currentEpoch={currentEpoch}
           title={self ? "Other Publishers" : "Publishers"}
           availableToStake={availableToStake}
           publishers={otherPublishers}
@@ -401,6 +406,7 @@ const OptOutButton = ({ api, self }: OptOutButtonProps) => {
 
 type PublisherListProps = {
   api: States[ApiStateType.Loaded] | States[ApiStateType.LoadedNoStakeAccount];
+  currentEpoch: bigint;
   title: string;
   availableToStake: bigint;
   totalStaked: bigint;
@@ -410,6 +416,7 @@ type PublisherListProps = {
 
 const PublisherList = ({
   api,
+  currentEpoch,
   title,
   availableToStake,
   publishers,
@@ -583,6 +590,7 @@ const PublisherList = ({
             {paginatedPublishers.map((publisher) => (
               <Publisher
                 api={api}
+                currentEpoch={currentEpoch}
                 key={publisher.publicKey.toBase58()}
                 availableToStake={availableToStake}
                 publisher={publisher}
@@ -686,6 +694,7 @@ const PublisherTableHeader = Styled(
 
 type PublisherProps = {
   api: States[ApiStateType.Loaded] | States[ApiStateType.LoadedNoStakeAccount];
+  currentEpoch: bigint;
   availableToStake: bigint;
   totalStaked: bigint;
   isSelf?: boolean;
@@ -713,6 +722,7 @@ type PublisherProps = {
 
 const Publisher = ({
   api,
+  currentEpoch,
   publisher,
   availableToStake,
   totalStaked,
@@ -837,6 +847,7 @@ const Publisher = ({
         >
           <StakeToPublisherButton
             api={api}
+            currentEpoch={currentEpoch}
             availableToStake={availableToStake}
             publisher={publisher}
             yieldRate={yieldRate}
@@ -913,7 +924,10 @@ const Publisher = ({
                           max={staked}
                           transfer={unstake}
                         >
-                          <StakingTimeline cooldownOnly />
+                          <StakingTimeline
+                            cooldownOnly
+                            currentEpoch={currentEpoch}
+                          />
                         </TransferButton>
                       </td>
                     </tr>
@@ -933,12 +947,14 @@ const PublisherTableCell = Styled("td", "py-4 px-5 whitespace-nowrap");
 type StakeToPublisherButtonProps = {
   api: States[ApiStateType.Loaded] | States[ApiStateType.LoadedNoStakeAccount];
   publisher: PublisherProps["publisher"];
+  currentEpoch: bigint;
   availableToStake: bigint;
   yieldRate: bigint;
 };
 
 const StakeToPublisherButton = ({
   api,
+  currentEpoch,
   availableToStake,
   publisher,
   yieldRate,
@@ -987,7 +1003,7 @@ const StakeToPublisherButton = ({
               %
             </div>
           </div>
-          <StakingTimeline />
+          <StakingTimeline currentEpoch={currentEpoch} />
         </>
       )}
     </TransferButton>
diff --git a/apps/staking/src/components/ProgramSection/index.tsx b/apps/staking/src/components/ProgramSection/index.tsx
index 7eed64f37..8e8d49237 100644
--- a/apps/staking/src/components/ProgramSection/index.tsx
+++ b/apps/staking/src/components/ProgramSection/index.tsx
@@ -1,8 +1,8 @@
 import { ArrowLongDownIcon } from "@heroicons/react/24/outline";
+import { epochToDate } from "@pythnetwork/staking-sdk";
 import clsx from "clsx";
 import type { HTMLAttributes, ReactNode, ComponentProps } from "react";
 
-import { getUpcomingEpoch, getNextFullEpoch } from "../../api";
 import { StakingTimeline } from "../StakingTimeline";
 import { Tokens } from "../Tokens";
 import { TransferButton } from "../TransferButton";
@@ -10,6 +10,7 @@ import { TransferButton } from "../TransferButton";
 type Props = HTMLAttributes<HTMLDivElement> & {
   name: string;
   description: string;
+  currentEpoch: bigint;
   warmup: bigint;
   staked: bigint;
   cooldown: bigint;
@@ -45,6 +46,7 @@ export const ProgramSection = ({
   name,
   description,
   children,
+  currentEpoch,
   warmup,
   staked,
   cooldown,
@@ -82,7 +84,7 @@ export const ProgramSection = ({
               max={available}
               transfer={stake}
             >
-              <StakingTimeline />
+              <StakingTimeline currentEpoch={currentEpoch} />
             </TransferButton>
           ),
         })}
@@ -96,7 +98,7 @@ export const ProgramSection = ({
         {...(warmup > 0n && {
           details: (
             <div className="mt-2 text-xs text-neutral-500">
-              Staking {getUpcomingEpoch().toLocaleString()}
+              Staking {epochToDate(currentEpoch + 1n).toLocaleString()}
             </div>
           ),
         })}
@@ -131,7 +133,7 @@ export const ProgramSection = ({
               max={staked}
               transfer={unstake}
             >
-              <StakingTimeline cooldownOnly />
+              <StakingTimeline cooldownOnly currentEpoch={currentEpoch} />
             </TransferButton>
           ),
         })}
@@ -147,13 +149,13 @@ export const ProgramSection = ({
             {cooldown > 0n && (
               <div className="mt-2 text-xs text-neutral-500">
                 <Tokens>{cooldown}</Tokens> end{" "}
-                {getUpcomingEpoch().toLocaleString()}
+                {epochToDate(currentEpoch + 1n).toLocaleString()}
               </div>
             )}
             {cooldown2 > 0n && (
               <div className="mt-2 text-xs text-neutral-500">
                 <Tokens>{cooldown2}</Tokens> end{" "}
-                {getNextFullEpoch().toLocaleString()}
+                {epochToDate(currentEpoch + 2n).toLocaleString()}
               </div>
             )}
           </>
diff --git a/apps/staking/src/components/StakingTimeline/index.tsx b/apps/staking/src/components/StakingTimeline/index.tsx
index 306a8423f..932313ec4 100644
--- a/apps/staking/src/components/StakingTimeline/index.tsx
+++ b/apps/staking/src/components/StakingTimeline/index.tsx
@@ -1,12 +1,12 @@
 import { ArrowLongDownIcon } from "@heroicons/react/24/outline";
-
-import { getUpcomingEpoch, getNextFullEpoch } from "../../api";
+import { epochToDate } from "@pythnetwork/staking-sdk";
 
 type Props = {
   cooldownOnly?: boolean | undefined;
+  currentEpoch: bigint;
 };
 
-export const StakingTimeline = ({ cooldownOnly }: Props) => (
+export const StakingTimeline = ({ cooldownOnly, currentEpoch }: Props) => (
   <div className="mb-2 flex flex-col gap-1">
     <div className="text-sm">Timeline</div>
     <div className="grid grid-cols-[max-content_1fr_max-content] items-center gap-x-4 gap-y-3 border border-neutral-600/50 bg-pythpurple-100/10 px-8 py-6 text-sm font-light">
@@ -15,7 +15,7 @@ export const StakingTimeline = ({ cooldownOnly }: Props) => (
           <div className="size-4 rounded-full border border-dashed border-pythpurple-100" />
           <div>Warmup</div>
           <div className="text-right">
-            {getUpcomingEpoch().toLocaleString()}
+            {epochToDate(currentEpoch + 1n).toLocaleString()}
           </div>
           <ArrowLongDownIcon className="size-4 scale-y-[200%] [mask-image:linear-gradient(to_bottom,_transparent,_black_125%)]" />
           <div>Staking</div>
@@ -25,7 +25,9 @@ export const StakingTimeline = ({ cooldownOnly }: Props) => (
       <div className="size-4 rounded-full border border-pythpurple-100 bg-pythpurple-600" />
       <div>Cooldown</div>
       <div className="text-right">
-        {cooldownOnly ? getNextFullEpoch().toLocaleString() : "One full epoch"}
+        {cooldownOnly
+          ? epochToDate(currentEpoch + 2n).toLocaleString()
+          : "One full epoch"}
       </div>
     </div>
   </div>