diff --git a/apps/client/src/app/room/[code]/game/[gameId]/game.tsx b/apps/client/src/app/room/[code]/game/[gameId]/game.tsx
index 3cf51ac..bcfea39 100644
--- a/apps/client/src/app/room/[code]/game/[gameId]/game.tsx
+++ b/apps/client/src/app/room/[code]/game/[gameId]/game.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { useCallback, useContext, useEffect, useMemo, useState } from "react";
+import { useCallback, useContext, useEffect, useState } from "react";
 import { EventFrom, State } from "xstate";
 import { useMachine } from "@xstate/react";
 import { AnimatePresence } from "framer-motion";
@@ -15,8 +15,8 @@ import {
   getLeaderboardById,
 } from "@ai/utils/queries";
 import {
+  CurrentGameComponent,
   gameMachine,
-  getCurrentComponent,
 } from "@ai/components/game/game-machine";
 import { SocketContext } from "@ai/utils/socket-provider";
 import { cn } from "@ai/utils/cn";
@@ -37,6 +37,7 @@ export default function Game({ gameInfo, session }: GameProps) {
   // Socket for real-time communication
   const socket = useContext(SocketContext);
 
+  // Next.js router
   const router = useRouter();
 
   // Wait until the client mounts to avoid hydration errors
@@ -168,31 +169,6 @@ export default function Game({ gameInfo, session }: GameProps) {
     };
   }, [handlePlayAnotherGame, handleServerEvent, socket]);
 
-  // Game component shown based off state
-  const currentComponent = useMemo(() => {
-    return getCurrentComponent(
-      gameInfo,
-      state,
-      send,
-      hostId,
-      submittedPlayerIds,
-      currFaceOffQuestion,
-      votedPlayers,
-      leaderboard,
-      session,
-    );
-  }, [
-    gameInfo,
-    state,
-    send,
-    hostId,
-    submittedPlayerIds,
-    currFaceOffQuestion,
-    votedPlayers,
-    leaderboard,
-    session,
-  ]);
-
   return (
     <main
       className={cn(
@@ -205,7 +181,20 @@ export default function Game({ gameInfo, session }: GameProps) {
           <Menu session={session} roomCode={gameInfo.game.roomCode} />
         </div>
         <AnimatePresence mode="wait">
-          {isMounted ? currentComponent : null}
+          {isMounted ? (
+            <CurrentGameComponent
+              key={state.value as string}
+              gameInfo={gameInfo}
+              state={state}
+              send={send}
+              hostId={hostId}
+              submittedPlayerIds={submittedPlayerIds}
+              currFaceOffQuestion={currFaceOffQuestion}
+              votedPlayers={votedPlayers}
+              leaderboard={leaderboard}
+              session={session}
+            />
+          ) : null}
         </AnimatePresence>
       </section>
     </main>
diff --git a/apps/client/src/components/game/game-machine.tsx b/apps/client/src/components/game/game-machine.tsx
index 0e92a1b..7b4e2e5 100644
--- a/apps/client/src/components/game/game-machine.tsx
+++ b/apps/client/src/components/game/game-machine.tsx
@@ -174,102 +174,127 @@ export const gameMachine = createMachine(
   },
 );
 
-export const getCurrentComponent = (
-  gameInfo: GameInfo,
-  state: StateFrom<typeof gameMachine>,
-  send: (event: EventFrom<typeof gameMachine>) => StateFrom<typeof gameMachine>,
-  hostId: string | null,
-  submittedPlayerIds: Set<string>,
-  currFaceOffQuestion: QuestionGenerations | undefined,
-  votedPlayers: UserVote[],
-  leaderboard: GetGameLeaderboardResponse | undefined,
-  session: Session,
-) => {
-  const stateMachineComponentMap: Record<
-    StateValueFrom<typeof gameMachine>,
-    ReactNode
-  > = {
-    connectingToMainframe: (
-      <TransitionWrapper key="connectingToMainframe">
-        <ConnectToMainframe />
-      </TransitionWrapper>
-    ),
-    connectionEstablished: (
-      <TransitionWrapper key="connectionEstablished">
-        <ConnectionEstablished />
-      </TransitionWrapper>
-    ),
-    prompt: (
-      <TransitionWrapper key="prompt">
-        <Prompt
-          gameInfo={gameInfo}
-          state={state}
-          send={send}
-          session={session}
-        />
-      </TransitionWrapper>
-    ),
-    promptSubmitted: (
-      <TransitionWrapper key="promptSubmitted">
-        <PromptSubmitted
-          gameInfo={gameInfo}
-          submittedPlayerIds={submittedPlayerIds}
-          session={session}
-        />
-      </TransitionWrapper>
-    ),
-    promptDone: (
-      <TransitionWrapper key="promptDone">
-        <AnnouncementText text="All participants have submitted!" />
-      </TransitionWrapper>
-    ),
-    faceOff: (
-      <TransitionWrapper key="faceOff">
-        <FaceOff
-          gameInfo={gameInfo}
-          currQuestionGenerations={currFaceOffQuestion}
-          session={session}
-        />
-      </TransitionWrapper>
-    ),
-    faceOffResults: (
-      <TransitionWrapper key="faceOffResults">
-        <FaceOffResult
-          currQuestionGenerations={currFaceOffQuestion}
-          votedPlayers={votedPlayers}
-        />
-      </TransitionWrapper>
-    ),
-    nextRound: (
-      <TransitionWrapper key="nextRound">
-        <NextRound nextQueryNum={state.context.round} totalQueries={3} />
-      </TransitionWrapper>
-    ),
-    winnerLeadUp: (
-      <TransitionWrapper key="winnerLeadUp">
-        <WinnerLeadUp />
-      </TransitionWrapper>
-    ),
-    winner: (
-      <TransitionWrapper key="winner">
-        <Winner leaderboard={leaderboard} />
-      </TransitionWrapper>
-    ),
-    leaderboard: (
-      <TransitionWrapper key="leaderboard">
-        <Leaderboard
-          gameInfo={gameInfo}
-          state={state}
-          send={send}
-          leaderboard={leaderboard}
-          hostId={hostId}
-          session={session}
-        />
-      </TransitionWrapper>
-    ),
-  };
-
-  return stateMachineComponentMap[
-    state.value as StateValueFrom<typeof gameMachine>
-  ];
+export const CurrentGameComponent = ({
+  gameInfo,
+  state,
+  send,
+  hostId,
+  submittedPlayerIds,
+  currFaceOffQuestion,
+  votedPlayers,
+  leaderboard,
+  session,
+}: {
+  gameInfo: GameInfo;
+  state: StateFrom<typeof gameMachine>;
+  send: (event: EventFrom<typeof gameMachine>) => StateFrom<typeof gameMachine>;
+  hostId: string | null;
+  submittedPlayerIds: Set<string>;
+  currFaceOffQuestion: QuestionGenerations | undefined;
+  votedPlayers: UserVote[];
+  leaderboard: GetGameLeaderboardResponse | undefined;
+  session: Session;
+}) => {
+  switch (state.value as StateValueFrom<typeof gameMachine>) {
+    case "connectingToMainframe": {
+      return (
+        <TransitionWrapper key="connectingToMainframe">
+          <ConnectToMainframe />
+        </TransitionWrapper>
+      );
+    }
+    case "connectionEstablished": {
+      return (
+        <TransitionWrapper key="connectionEstablished">
+          <ConnectionEstablished />
+        </TransitionWrapper>
+      );
+    }
+    case "prompt": {
+      return (
+        <TransitionWrapper key="prompt">
+          <Prompt
+            gameInfo={gameInfo}
+            state={state}
+            send={send}
+            session={session}
+          />
+        </TransitionWrapper>
+      );
+    }
+    case "promptSubmitted": {
+      return (
+        <TransitionWrapper key="promptSubmitted">
+          <PromptSubmitted
+            gameInfo={gameInfo}
+            submittedPlayerIds={submittedPlayerIds}
+            session={session}
+          />
+        </TransitionWrapper>
+      );
+    }
+    case "promptDone": {
+      return (
+        <TransitionWrapper key="promptDone">
+          <AnnouncementText text="All participants have submitted!" />
+        </TransitionWrapper>
+      );
+    }
+    case "faceOff": {
+      return (
+        <TransitionWrapper key="faceOff">
+          <FaceOff
+            gameInfo={gameInfo}
+            currQuestionGenerations={currFaceOffQuestion}
+            session={session}
+          />
+        </TransitionWrapper>
+      );
+    }
+    case "faceOffResults": {
+      return (
+        <TransitionWrapper key="faceOffResults">
+          <FaceOffResult
+            currQuestionGenerations={currFaceOffQuestion}
+            votedPlayers={votedPlayers}
+          />
+        </TransitionWrapper>
+      );
+    }
+    case "nextRound": {
+      return (
+        <TransitionWrapper key="nextRound">
+          <NextRound nextQueryNum={state.context.round} totalQueries={3} />
+        </TransitionWrapper>
+      );
+    }
+    case "winnerLeadUp": {
+      return (
+        <TransitionWrapper key="winnerLeadUp">
+          <WinnerLeadUp />
+        </TransitionWrapper>
+      );
+    }
+    case "winner": {
+      return (
+        <TransitionWrapper key="winner">
+          <Winner leaderboard={leaderboard} />
+        </TransitionWrapper>
+      );
+    }
+    case "leaderboard": {
+      return (
+        <TransitionWrapper key="leaderboard">
+          <Leaderboard
+            gameInfo={gameInfo}
+            state={state}
+            send={send}
+            leaderboard={leaderboard}
+            hostId={hostId}
+            session={session}
+          />
+        </TransitionWrapper>
+      );
+    }
+  }
 };