Merge pull request #18 from OriginTrail/fix/oauth-stale-code

Fixed a bug with Oauth flow with a stale login code

Author: Jurij89

apps/agent/src/app/(protected)/_layout.tsx

@@ -1,30 +1,96 @@
-import { useCallback, useEffect, useState } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
 import { router, Slot, useGlobalSearchParams, usePathname } from "expo-router";
 import * as SplashScreen from "expo-splash-screen";
+import AsyncStorage from "@react-native-async-storage/async-storage";
 
 import { McpContextProvider } from "@/client";
 import { useAlerts } from "@/components/Alerts";
 import { SettingsProvider } from "@/hooks/useSettings";
 
+const REDIRECT_COUNT_KEY = "oauth_redirect_count";
+const MAX_REDIRECTS = 3;
+
+const getRedirectCount = async (): Promise<number> => {
+  const count = await AsyncStorage.getItem(REDIRECT_COUNT_KEY);
+  return count ? parseInt(count, 10) : 0;
+};
+
+const incrementRedirectCount = async (): Promise<number> => {
+  const newCount = (await getRedirectCount()) + 1;
+  await AsyncStorage.setItem(REDIRECT_COUNT_KEY, newCount.toString());
+  return newCount;
+};
+
+const resetRedirectCount = async (): Promise<void> => {
+  await AsyncStorage.removeItem(REDIRECT_COUNT_KEY);
+};
+
 export default function ProtectedLayout() {
   const params = useGlobalSearchParams<{ code?: string; error?: string }>();
   const { showAlert } = useAlerts();
 
   const [idleMcp, setIdleMcp] = useState(false);
   const [idleSettings, setIdleSettings] = useState(false);
+  const isHandlingError = useRef(false);
 
   const mcpCallback = useCallback(
     (error?: Error) => {
       setIdleMcp(true);
+      router.setParams({ code: undefined }); // Always clear code param
 
-      if (!error) router.setParams({ code: undefined });
-      else
-        showAlert({
-          type: "error",
-          title: "MCP Error",
-          message: error.message,
-          timeout: 5000,
-        });
+      if (!error) {
+        // Success - reset redirect count
+        resetRedirectCount();
+        return;
+      }
+
+      // Prevent concurrent error handling
+      if (isHandlingError.current) return;
+      isHandlingError.current = true;
+
+      // Check if this is an auth-related error requiring re-authentication
+      const errorMessage = error.message.toLowerCase();
+      const isAuthError =
+        errorMessage.includes("invalid_grant") ||
+        errorMessage.includes("invalid grant") ||
+        errorMessage.includes("invalid token") ||
+        errorMessage.includes("authorization code") ||
+        errorMessage.includes("oauth client") ||
+        errorMessage.includes("unauthorized");
+
+      if (isAuthError) {
+        // Auto-redirect with loop protection
+        (async () => {
+          try {
+            const redirectCount = await incrementRedirectCount();
+            if (redirectCount > MAX_REDIRECTS) {
+              await resetRedirectCount();
+              showAlert({
+                type: "error",
+                title: "Authentication Failed",
+                message:
+                  "Unable to authenticate. Please try again later or contact support.",
+                timeout: 10000,
+              });
+              isHandlingError.current = false;
+              return;
+            }
+            router.replace("/authorize");
+          } catch {
+            isHandlingError.current = false;
+          }
+        })();
+        return;
+      }
+
+      // For other errors, show alert
+      showAlert({
+        type: "error",
+        title: "Connection Error",
+        message: error.message,
+        timeout: 5000,
+      });
+      isHandlingError.current = false;
     },
     [showAlert],
   );

apps/agent/src/app/(protected)/login.tsx

@@ -1,5 +1,5 @@
 import { Redirect, router, useLocalSearchParams } from "expo-router";
-import { useCallback } from "react";
+import { useCallback, useState } from "react";
 import { View } from "react-native";
 import * as Linking from "expo-linking";
 import * as SplashScreen from "expo-splash-screen";
@@ -14,8 +14,12 @@ import Footer from "@/components/layout/Footer";
 import FormTitle from "@/components/forms/FormTitle";
 import LoginForm from "@/components/forms/LoginForm";
 
-const getErrorMessage = (err: any) => {
-  if (!(err instanceof AuthError)) return "Unknown error occurred!";
+const getErrorMessage = (err: any): string => {
+  if (!(err instanceof AuthError)) {
+    // Check for OAuth-related errors in raw error messages
+    const rawMessage = err?.message || String(err);
+    return rawMessage;
+  }
   switch (err.code) {
     case AuthError.Code.INVALID_CREDENTIALS:
       return "Invalid username or password";
@@ -28,9 +32,22 @@ const getErrorMessage = (err: any) => {
   }
 };
 
+// Check if error indicates a stale/invalid authorization code
+const isStaleCodeError = (errorMessage: string): boolean => {
+  const lowerMessage = errorMessage.toLowerCase();
+  return (
+    /invalid.*(grant|code|authorization)/i.test(errorMessage) ||
+    /authorization.*code/i.test(errorMessage) ||
+    /oauth.*client/i.test(errorMessage) ||
+    lowerMessage.includes("invalid_grant") ||
+    lowerMessage.includes("expired")
+  );
+};
+
 export default function Login() {
   SplashScreen.hide();
   const { code } = useLocalSearchParams<{ code?: string }>();
+  const [validCode, setValidCode] = useState<string | undefined>(code);
 
   const tryLogin = useCallback(
     async ({
@@ -44,7 +61,7 @@ export default function Login() {
     }) => {
       try {
         const url = await login({
-          code: code ?? "",
+          code: validCode ?? "",
           credentials: { email, password },
           rememberMe,
           fetch: (url, opts) => fetch(url.toString(), opts as any),
@@ -56,10 +73,19 @@ export default function Login() {
         else Linking.openURL(url);
       } catch (error) {
         const errorMessage = getErrorMessage(error);
+
+        // Detect stale authorization code errors
+        if (isStaleCodeError(errorMessage)) {
+          router.setParams({ code: undefined });
+          setValidCode(undefined);
+          // Generic message - don't expose OAuth internals
+          throw new Error("Your session has expired. Please log in again.");
+        }
+
         throw new Error(errorMessage);
       }
     },
-    [code],
+    [validCode],
   );
 
   const { connected } = useMcpClient();