Redirect with query param (#2811)

* validated

* k

* k

* k

* minor update

web/src/app/auth/login/page.tsx

@@ -11,7 +11,7 @@ import { SignInButton } from "./SignInButton";
 import { EmailPasswordForm } from "./EmailPasswordForm";
 import { Card, Title, Text } from "@tremor/react";
 import Link from "next/link";
-import { Logo } from "@/components/Logo";
+
 import { LoginText } from "./LoginText";
 import { getSecondsUntilExpiration } from "@/lib/time";
 import AuthFlowContainer from "@/components/auth/AuthFlowContainer";
@@ -37,6 +37,10 @@ const Page = async ({
     console.log(`Some fetch failed for the login page - ${e}`);
   }
 
+  const nextUrl = Array.isArray(searchParams?.next)
+    ? searchParams?.next[0]
+    : searchParams?.next || null;
+
   // simply take the user to the home page if Auth is disabled
   if (authTypeMetadata?.authType === "disabled") {
     return redirect("/");
@@ -59,7 +63,7 @@ const Page = async ({
   let authUrl: string | null = null;
   if (authTypeMetadata) {
     try {
-      authUrl = await getAuthUrlSS(authTypeMetadata.authType);
+      authUrl = await getAuthUrlSS(authTypeMetadata.authType, nextUrl!);
     } catch (e) {
       console.log(`Some fetch failed for the login page - ${e}`);
     }
@@ -88,6 +92,7 @@ const Page = async ({
             />
           </>
         )}
+
         {authTypeMetadata?.authType === "basic" && (
           <Card className="mt-4 w-96">
             <div className="flex">

web/src/app/auth/oauth/callback/route.ts

@@ -8,7 +8,8 @@ export const GET = async (request: NextRequest) => {
   const url = new URL(buildUrl("/auth/oauth/callback"));
   url.search = request.nextUrl.search;
 
-  const response = await fetch(url.toString());
+  // Set 'redirect' to 'manual' to prevent automatic redirection
+  const response = await fetch(url.toString(), { redirect: "manual" });
   const setCookieHeader = response.headers.get("set-cookie");
 
   if (response.status === 401) {
@@ -21,9 +22,13 @@ export const GET = async (request: NextRequest) => {
     return NextResponse.redirect(new URL("/auth/error", getDomain(request)));
   }
 
+  // Get the redirect URL from the backend's 'Location' header, or default to '/'
+  const redirectUrl = response.headers.get("location") || "/";
+
   const redirectResponse = NextResponse.redirect(
-    new URL("/", getDomain(request))
+    new URL(redirectUrl, getDomain(request))
   );
+
   redirectResponse.headers.set("set-cookie", setCookieHeader);
   return redirectResponse;
 };

web/src/app/auth/oidc/callback/route.ts

@@ -7,17 +7,27 @@ export const GET = async (request: NextRequest) => {
   // which adds back a redirect to the main app.
   const url = new URL(buildUrl("/auth/oidc/callback"));
   url.search = request.nextUrl.search;
-
-  const response = await fetch(url.toString());
+  // Set 'redirect' to 'manual' to prevent automatic redirection
+  const response = await fetch(url.toString(), { redirect: "manual" });
   const setCookieHeader = response.headers.get("set-cookie");
 
+  if (response.status === 401) {
+    return NextResponse.redirect(
+      new URL("/auth/create-account", getDomain(request))
+    );
+  }
+
   if (!setCookieHeader) {
     return NextResponse.redirect(new URL("/auth/error", getDomain(request)));
   }
 
+  // Get the redirect URL from the backend's 'Location' header, or default to '/'
+  const redirectUrl = response.headers.get("location") || "/";
+
   const redirectResponse = NextResponse.redirect(
-    new URL("/", getDomain(request))
+    new URL(redirectUrl, getDomain(request))
   );
+
   redirectResponse.headers.set("set-cookie", setCookieHeader);
   return redirectResponse;
 };

web/src/app/page.tsx

@@ -3,7 +3,6 @@ import { redirect } from "next/navigation";
 
 export default async function Page() {
   const settings = await fetchSettingsSS();
-
   if (!settings) {
     redirect("/search");
   }

web/src/app/search/page.tsx

@@ -36,8 +36,13 @@ import WrappedSearch from "./WrappedSearch";
 import { SearchProvider } from "@/components/context/SearchContext";
 import { fetchLLMProvidersSS } from "@/lib/llm/fetchLLMs";
 import { LLMProviderDescriptor } from "../admin/configuration/llm/interfaces";
+import { headers } from "next/headers";
 
-export default async function Home() {
+export default async function Home({
+  searchParams,
+}: {
+  searchParams: { [key: string]: string | string[] | undefined };
+}) {
   // Disable caching so we always get the up to date connector / document set / persona info
   // importantly, this prevents users from adding a connector, going back to the main page,
   // and then getting hit with a "No Connectors" popup
@@ -82,8 +87,17 @@ export default async function Home() {
   const llmProviders = (results[7] || []) as LLMProviderDescriptor[];
 
   const authDisabled = authTypeMetadata?.authType === "disabled";
+
   if (!authDisabled && !user) {
-    return redirect("/auth/login");
+    const headersList = headers();
+    const fullUrl = headersList.get("x-url") || "/search";
+    const searchParamsString = new URLSearchParams(
+      searchParams as unknown as Record<string, string>
+    ).toString();
+    const redirectUrl = searchParamsString
+      ? `${fullUrl}?${searchParamsString}`
+      : fullUrl;
+    return redirect(`/auth/login?next=${encodeURIComponent(redirectUrl)}`);
   }
 
   if (user && !user.is_verified && authTypeMetadata?.requiresVerification) {

web/src/components/UserDropdown.tsx

@@ -3,7 +3,7 @@
 import { useState, useRef, useContext, useEffect, useMemo } from "react";
 import { FiLogOut } from "react-icons/fi";
 import Link from "next/link";
-import { useRouter } from "next/navigation";
+import { useRouter, usePathname, useSearchParams } from "next/navigation";
 import { User, UserRole } from "@/lib/types";
 import { checkUserIsNoAuthUser, logout } from "@/lib/user";
 import { Popover } from "./popover/Popover";
@@ -65,6 +65,8 @@ export function UserDropdown({
   const [userInfoVisible, setUserInfoVisible] = useState(false);
   const userInfoRef = useRef<HTMLDivElement>(null);
   const router = useRouter();
+  const pathname = usePathname();
+  const searchParams = useSearchParams();
 
   const combinedSettings = useContext(SettingsContext);
   const customNavItems: NavigationItem[] = useMemo(
@@ -87,8 +89,17 @@ export function UserDropdown({
     logout().then((isSuccess) => {
       if (!isSuccess) {
         alert("Failed to logout");
+        return;
       }
-      router.push("/auth/login");
+
+      // Construct the current URL
+      const currentUrl = `${pathname}${searchParams.toString() ? `?${searchParams.toString()}` : ""}`;
+
+      // Encode the current URL to use as a redirect parameter
+      const encodedRedirect = encodeURIComponent(currentUrl);
+
+      // Redirect to login page with the current page as a redirect parameter
+      router.push(`/auth/login?next=${encodedRedirect}`);
     });
   };
 

web/src/lib/chat/fetchChatData.ts

@@ -20,7 +20,7 @@ import { fetchLLMProvidersSS } from "@/lib/llm/fetchLLMs";
 import { LLMProviderDescriptor } from "@/app/admin/configuration/llm/interfaces";
 import { Folder } from "@/app/chat/folders/interfaces";
 import { personaComparator } from "@/app/admin/assistants/lib";
-import { cookies } from "next/headers";
+import { cookies, headers } from "next/headers";
 import {
   SIDEBAR_TOGGLED_COOKIE_NAME,
   DOCUMENT_SIDEBAR_WIDTH_COOKIE_NAME,
@@ -29,6 +29,7 @@ import { hasCompletedWelcomeFlowSS } from "@/components/initialSetup/welcome/Wel
 import { fetchAssistantsSS } from "../assistants/fetchAssistantsSS";
 import { NEXT_PUBLIC_DEFAULT_SIDEBAR_OPEN } from "../constants";
 import { checkLLMSupportsImageInput } from "../llm/utils";
+import { redirect } from "next/navigation";
 
 interface FetchChatDataResult {
   user: User | null;
@@ -98,7 +99,15 @@ export async function fetchChatData(searchParams: {
 
   const authDisabled = authTypeMetadata?.authType === "disabled";
   if (!authDisabled && !user) {
-    return { redirect: "/auth/login" };
+    const headersList = headers();
+    const fullUrl = headersList.get("x-url") || "/chat";
+    const searchParamsString = new URLSearchParams(
+      searchParams as unknown as Record<string, string>
+    ).toString();
+    const redirectUrl = searchParamsString
+      ? `${fullUrl}?${searchParamsString}`
+      : fullUrl;
+    return redirect(`/auth/login?next=${encodeURIComponent(redirectUrl)}`);
   }
 
   if (user && !user.is_verified && authTypeMetadata?.requiresVerification) {

web/src/lib/userSS.ts

@@ -40,8 +40,12 @@ export const getAuthDisabledSS = async (): Promise<boolean> => {
   return (await getAuthTypeMetadataSS()).authType === "disabled";
 };
 
-const geOIDCAuthUrlSS = async (): Promise<string> => {
-  const res = await fetch(buildUrl("/auth/oidc/authorize"));
+const getOIDCAuthUrlSS = async (nextUrl: string | null): Promise<string> => {
+  const res = await fetch(
+    buildUrl(
+      `/auth/oidc/authorize${nextUrl ? `?next=${encodeURIComponent(nextUrl)}` : ""}`
+    )
+  );
   if (!res.ok) {
     throw new Error("Failed to fetch data");
   }
@@ -51,7 +55,7 @@ const geOIDCAuthUrlSS = async (): Promise<string> => {
 };
 
 const getGoogleOAuthUrlSS = async (): Promise<string> => {
-  const res = await fetch(buildUrl("/auth/oauth/authorize"));
+  const res = await fetch(buildUrl(`/auth/oauth/authorize`));
   if (!res.ok) {
     throw new Error("Failed to fetch data");
   }
@@ -70,7 +74,10 @@ const getSAMLAuthUrlSS = async (): Promise<string> => {
   return data.authorization_url;
 };
 
-export const getAuthUrlSS = async (authType: AuthType): Promise<string> => {
+export const getAuthUrlSS = async (
+  authType: AuthType,
+  nextUrl: string | null
+): Promise<string> => {
   // Returns the auth url for the given auth type
   switch (authType) {
     case "disabled":
@@ -84,7 +91,7 @@ export const getAuthUrlSS = async (authType: AuthType): Promise<string> => {
       return await getSAMLAuthUrlSS();
     }
     case "oidc": {
-      return await geOIDCAuthUrlSS();
+      return await getOIDCAuthUrlSS(nextUrl);
     }
   }
 };