Improve auth handling

All-Hands-AI · Nov 22, 2024 · f0135ec · f0135ec
1 parent 5e2a9a4
commit f0135ec
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 24 deletions.
diff --git a/frontend/src/api/open-hands.ts b/frontend/src/api/open-hands.ts
@@ -239,13 +239,19 @@ class OpenHands {
   static async getGitHubAccessToken(
     code: string,
   ): Promise<GitHubAccessTokenResponse> {
-    return request(`/api/github/callback`, {
+    const response = await fetch("/api/github/callback", {
       method: "POST",
       body: JSON.stringify({ code }),
       headers: {
         "Content-Type": "application/json",
       },
     });
+
+    if (!response.ok) {
+      throw new Error("Failed to get GitHub access token");
+    }
+
+    return response.json();
   }
 
   /**

diff --git a/frontend/src/routes/_oh.tsx b/frontend/src/routes/_oh.tsx
@@ -84,7 +84,11 @@ export default function MainApp() {
 
   const config = useConfig();
   const user = useGitHubUser();
-  const { data: isAuthed, isFetched } = useIsAuthed();
+  const {
+    data: isAuthed,
+    isFetched,
+    isFetching: isFetchingAuth,
+  } = useIsAuthed();
   const aiConfigOptions = useAIConfigOptions();
 
   const gitHubAuthUrl = useGitHubAuthUrl({
@@ -251,7 +255,7 @@ export default function MainApp() {
           />
         </ModalBackdrop>
       )}
-      {!isAuthed && config.data?.APP_MODE === "saas" && (
+      {!isFetchingAuth && !isAuthed && config.data?.APP_MODE === "saas" && (
         <WaitlistModal ghToken={gitHubToken} githubAuthUrl={gitHubAuthUrl} />
       )}
       {consentFormIsOpen && (

diff --git a/frontend/src/routes/oauth.github.callback.tsx b/frontend/src/routes/oauth.github.callback.tsx
@@ -1,34 +1,34 @@
-import {
-  ClientLoaderFunctionArgs,
-  json,
-  redirect,
-  useLoaderData,
-} from "@remix-run/react";
+import { useNavigate, useSearchParams } from "@remix-run/react";
+import { useQuery } from "@tanstack/react-query";
+import React from "react";
 import OpenHands from "#/api/open-hands";
+import { useAuth } from "#/context/auth-context";
 
-export const clientLoader = async ({ request }: ClientLoaderFunctionArgs) => {
-  const url = new URL(request.url);
-  const code = url.searchParams.get("code");
-
-  if (code) {
-    const { access_token: accessToken } =
-      await OpenHands.getGitHubAccessToken(code);
+function OAuthGitHubCallback() {
+  const navigate = useNavigate();
+  const [searchParams] = useSearchParams();
+  const { setGitHubToken } = useAuth();
 
-    localStorage.setItem("ghToken", accessToken);
+  const code = searchParams.get("code");
 
-    return redirect("/");
-  }
+  const { data, isSuccess, error } = useQuery({
+    queryKey: ["access_token", code],
+    queryFn: () => OpenHands.getGitHubAccessToken(code!),
+    enabled: !!code,
+  });
 
-  return json({ error: "No code provided" }, { status: 400 });
-};
+  React.useEffect(() => {
+    if (isSuccess) {
+      setGitHubToken(data.access_token);
+      navigate("/");
+    }
+  }, [isSuccess]);
 
-function OAuthGitHubCallback() {
-  const { error } = useLoaderData<typeof clientLoader>();
   if (error) {
     return (
       <div>
         <h1>Error</h1>
-        <p>{error}</p>
+        <p>{error.message}</p>
       </div>
     );
   }