fix: ensure posts and comments display correct score after user votes

Author: sunaurus

src/app/(ui)/vote/VoteButtons.tsx

@@ -12,26 +12,46 @@ type BaseProps = {
   voteConfig: VoteConfig;
   loggedInUserScore?: number;
   score?: number;
+  onScoreChange?(newScore: number): void;
 };
 
 type CommentProps = BaseProps & { commentId: number; postId: number };
 type PostProps = BaseProps & { postId: number };
 
 type Props = CommentProps | PostProps;
+
+const getNewScore = (
+  prevScore: number | undefined,
+  prevUserScore: number | undefined,
+  newUserScore: number,
+) => {
+  let newScore = prevScore ?? 0;
+  if (newScore) {
+    newScore = newScore - (prevUserScore ?? 0) + newUserScore;
+  }
+
+  return newScore;
+};
+
 export const VoteButtons = memo(
   (props: Props) => {
-    const totalScore = !isComment(props) ? props.score : undefined; // We only show score between vote buttons for posts, not comments
-    const [userScore, setUserScore] = useState(props.loggedInUserScore);
+    const [score, setScore] = useState({
+      userScore: props.loggedInUserScore,
+      totalScore: props.score,
+    });
 
     const [{ optimisticUserScore, optimisticScore }, setOptimisticUserScore] =
       useOptimistic(
-        { optimisticUserScore: userScore, optimisticScore: totalScore },
+        {
+          optimisticUserScore: score.userScore ?? 0,
+          optimisticScore: score.totalScore ?? 0,
+        },
         (prevState, newUserScore: number) => {
-          let newScore = prevState.optimisticScore;
-          if (newScore) {
-            newScore =
-              newScore - (prevState.optimisticUserScore ?? 0) + newUserScore;
-          }
+          let newScore = getNewScore(
+            prevState.optimisticScore,
+            prevState.optimisticUserScore,
+            newUserScore,
+          );
 
           return {
             optimisticUserScore: newUserScore,
@@ -43,6 +63,29 @@ export const VoteButtons = memo(
     const upvoteResultScore = (optimisticUserScore ?? 0) <= 0 ? 1 : 0;
     const downvoteResultScore = (optimisticUserScore ?? 0) >= 0 ? -1 : 0;
 
+    const voteAction = async (userScoreAfterVoting: number) => {
+      setOptimisticUserScore(userScoreAfterVoting);
+      const action = isComment(props)
+        ? voteCommentAction.bind(null, props.commentId, userScoreAfterVoting)
+        : votePostAction.bind(null, props.postId, userScoreAfterVoting);
+      await action();
+      setScore((prev) => {
+        const newScore = getNewScore(
+          prev.totalScore,
+          prev.userScore,
+          userScoreAfterVoting,
+        );
+
+        if (props.onScoreChange) {
+          props.onScoreChange(newScore ?? 0);
+        }
+        return {
+          totalScore: newScore,
+          userScore: userScoreAfterVoting,
+        };
+      });
+    };
+
     return (
       <div
         className={classNames(
@@ -52,13 +95,7 @@ export const VoteButtons = memo(
       >
         <form
           action={async () => {
-            setOptimisticUserScore(upvoteResultScore);
-
-            const action = isComment(props)
-              ? voteCommentAction.bind(null, props.commentId, upvoteResultScore)
-              : votePostAction.bind(null, props.postId, upvoteResultScore);
-            await action();
-            setUserScore(upvoteResultScore);
+            await voteAction(upvoteResultScore);
           }}
         >
           <button type={"submit"}>
@@ -74,24 +111,15 @@ export const VoteButtons = memo(
             />
           </button>
         </form>
-        {optimisticScore !== undefined && props.voteConfig.scoresVisible && (
+        {!isComment(props) && props.voteConfig.scoresVisible && (
           <div className={"mb-0.5 w-8 text-center font-semibold"}>
             {formatCompactNumber(optimisticScore)}
           </div>
         )}
         {props.voteConfig.downvotesEnabled && (
           <form
             action={async () => {
-              setOptimisticUserScore(downvoteResultScore);
-              const action = isComment(props)
-                ? voteCommentAction.bind(
-                    null,
-                    props.commentId,
-                    downvoteResultScore,
-                  )
-                : votePostAction.bind(null, props.postId, downvoteResultScore);
-              await action();
-              setUserScore(downvoteResultScore);
+              await voteAction(downvoteResultScore);
             }}
           >
             <button type={"submit"}>

src/app/comment/Comment.tsx

@@ -47,6 +47,7 @@ export const Comment = (props: {
   const [deleted, setDeleted] = useState(props.commentView.comment.deleted);
   const [addedReplies, setAddedReplies] = useState<CommentNode[]>([]);
   const [saved, setSaved] = useState(props.commentView.saved);
+  const [score, setScore] = useState(props.commentView.counts.score);
 
   const markdownProps =
     editedContent !== null ? { content: editedContent } : props.markdown;
@@ -109,6 +110,7 @@ export const Comment = (props: {
         className={"peer-checked:collapse peer-checked:max-h-0"}
         commentId={props.commentView.comment.id}
         loggedInUserScore={props.commentView.my_vote}
+        onScoreChange={setScore}
         postId={props.commentView.post.id}
         score={props.commentView.counts.score}
         voteConfig={props.voteConfig}
@@ -148,8 +150,7 @@ export const Comment = (props: {
           <div className={"flex flex-wrap items-center gap-1"}>
             {props.voteConfig.scoresVisible && (
               <div className={""}>
-                {props.commentView.counts.score}{" "}
-                {props.commentView.counts.score === 1 ? "point" : "points"}
+                {score} {score === 1 ? "point" : "points"}
               </div>
             )}
             <FormattedTimestamp