From 4af84a29dc15ee2732a71e4552260eba4e4a9c3a Mon Sep 17 00:00:00 2001
From: Rohit Malhotra <rohitvinodmalhotra@gmail.com>
Date: Fri, 20 Dec 2024 14:37:04 -0500
Subject: [PATCH] Adding more resilience to refresh token logic (#5704)

---
 frontend/src/api/open-hands.ts              |  9 +++++++--
 frontend/src/context/auth-context.tsx       | 17 ++++++++++++++++-
 frontend/src/hooks/query/use-github-user.ts |  3 ++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/frontend/src/api/open-hands.ts b/frontend/src/api/open-hands.ts
index 5c9522f07e..23694a09e7 100644
--- a/frontend/src/api/open-hands.ts
+++ b/frontend/src/api/open-hands.ts
@@ -156,11 +156,16 @@ class OpenHands {
    */
   static async refreshToken(
     appMode: GetConfigResponse["APP_MODE"],
+    userId: string,
   ): Promise<string> {
     if (appMode === "oss") return "";
 
-    const response =
-      await openHands.post<GitHubAccessTokenResponse>("/api/refresh-token");
+    const response = await openHands.post<GitHubAccessTokenResponse>(
+      "/api/refresh-token",
+      {
+        userId,
+      },
+    );
     return response.data.access_token;
   }
 
diff --git a/frontend/src/context/auth-context.tsx b/frontend/src/context/auth-context.tsx
index ddce62a2ce..6ae9450c44 100644
--- a/frontend/src/context/auth-context.tsx
+++ b/frontend/src/context/auth-context.tsx
@@ -18,6 +18,7 @@ interface AuthContextType {
   gitHubToken: string | null;
   setToken: (token: string | null) => void;
   setGitHubToken: (token: string | null) => void;
+  setUserId: (userId: string) => void;
   clearToken: () => void;
   clearGitHubToken: () => void;
   refreshToken: () => Promise<boolean>;
@@ -34,6 +35,10 @@ function AuthProvider({ children }: React.PropsWithChildren) {
     () => localStorage.getItem("ghToken"),
   );
 
+  const [userIdState, setUserIdState] = React.useState<string>(
+    () => localStorage.getItem("userId") || "",
+  );
+
   const clearToken = () => {
     setTokenState(null);
     localStorage.removeItem("token");
@@ -43,7 +48,9 @@ function AuthProvider({ children }: React.PropsWithChildren) {
 
   const clearGitHubToken = () => {
     setGitHubTokenState(null);
+    setUserIdState("");
     localStorage.removeItem("ghToken");
+    localStorage.removeItem("userId");
 
     removeOpenHandsGitHubTokenHeader();
     removeGitHubAuthTokenHeader();
@@ -72,6 +79,11 @@ function AuthProvider({ children }: React.PropsWithChildren) {
     }
   };
 
+  const setUserId = (userId: string) => {
+    setUserIdState(userIdState);
+    localStorage.setItem("userId", userId);
+  };
+
   const logout = () => {
     clearGitHubToken();
     posthog.reset();
@@ -84,7 +96,7 @@ function AuthProvider({ children }: React.PropsWithChildren) {
       return false;
     }
 
-    const newToken = await OpenHands.refreshToken(config.APP_MODE);
+    const newToken = await OpenHands.refreshToken(config.APP_MODE, userIdState);
     if (newToken) {
       setGitHubToken(newToken);
       return true;
@@ -97,9 +109,11 @@ function AuthProvider({ children }: React.PropsWithChildren) {
   React.useEffect(() => {
     const storedToken = localStorage.getItem("token");
     const storedGitHubToken = localStorage.getItem("ghToken");
+    const userId = localStorage.getItem("userId") || "";
 
     setToken(storedToken);
     setGitHubToken(storedGitHubToken);
+    setUserId(userId);
     setupGithubAxiosInterceptors(refreshToken, logout);
   }, []);
 
@@ -109,6 +123,7 @@ function AuthProvider({ children }: React.PropsWithChildren) {
       gitHubToken: gitHubTokenState,
       setToken,
       setGitHubToken,
+      setUserId,
       clearToken,
       clearGitHubToken,
       refreshToken,
diff --git a/frontend/src/hooks/query/use-github-user.ts b/frontend/src/hooks/query/use-github-user.ts
index 6b9d7143b3..ff7cd18457 100644
--- a/frontend/src/hooks/query/use-github-user.ts
+++ b/frontend/src/hooks/query/use-github-user.ts
@@ -6,7 +6,7 @@ import { useAuth } from "#/context/auth-context";
 import { useConfig } from "./use-config";
 
 export const useGitHubUser = () => {
-  const { gitHubToken } = useAuth();
+  const { gitHubToken, setUserId } = useAuth();
   const { data: config } = useConfig();
 
   const user = useQuery({
@@ -18,6 +18,7 @@ export const useGitHubUser = () => {
 
   React.useEffect(() => {
     if (user.data) {
+      setUserId(user.data.id.toString());
       posthog.identify(user.data.login, {
         company: user.data.company,
         name: user.data.name,