cleanp

backend/danswer/auth/schemas.py

@@ -35,7 +35,6 @@ class UserCreate(schemas.BaseUserCreate):
     role: UserRole = UserRole.BASIC
     has_web_login: bool | None = True
     tenant_id: str | None = None
-    referral_source: str | None = None
 
 
 class UserUpdate(schemas.BaseUserUpdate):

backend/danswer/auth/users.py

@@ -1,5 +1,3 @@
-import base64
-import json
 import smtplib
 import uuid
 from collections.abc import AsyncGenerator
@@ -224,14 +222,11 @@ async def create(
         request: Optional[Request] = None,
     ) -> User:
         # Print the cookies from the request
+        referral_source = None
         if request:
-            cookies = request.cookies
-            print("Cookies from request:", cookies)
-        else:
-            print("No request object available")
-        tenant_id = await get_or_create_tenant_id(
-            user_create.email, user_create.referral_source
-        )
+            referral_source = request.cookies.get("referral_source", None)
+
+        tenant_id = await get_or_create_tenant_id(user_create.email, referral_source)
 
         async with get_async_session_with_tenant(tenant_id) as db_session:
             token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
@@ -291,39 +286,7 @@ async def oauth_callback(
         associate_by_email: bool = False,
         is_verified_by_default: bool = False,
     ) -> models.UOAP:
-        print("within callback")
-
-        if request:
-            cookies = request.cookies
-            print("Cookies from request:", cookies)
-        else:
-            print("No request object available")
-
-        referral_source = request.query_params.get("referral_source")
-        print("referral_source")
-
-        print(referral_source)
-
-        referral_source = None
-        if request and request.query_params.get("state"):
-            try:
-                state_json = base64.b64decode(request.query_params["state"]).decode(
-                    "utf-8"
-                )
-                state_data = json.loads(state_json)
-                referral_source = state_data.get("referralSource")
-            except Exception as e:
-                # Log the error for debugging
-                logger.error(f"Error parsing state parameter: {str(e)}")
-                # Continue with the flow, treating it as if there was no referral source
-
-        # Log the referral source for debugging
-        logger.info(f"zpzpzpReferral source: {referral_source}")
-
-        print("referral_source")
-        print(referral_source)
-        return
-        tenant_id = await get_or_create_tenant_id(account_email, referral_source)
+        tenant_id = await get_or_create_tenant_id(account_email)
 
         if not tenant_id:
             raise HTTPException(status_code=401, detail="User not found")

backend/ee/danswer/server/tenants/models.py

@@ -38,3 +38,4 @@ class ImpersonateRequest(BaseModel):
 class TenantCreationPayload(BaseModel):
     tenant_id: str
     email: str
+    referral_source: str | None = None

backend/ee/danswer/server/tenants/provisioning.py

@@ -70,9 +70,9 @@ async def create_tenant(email: str, referral_source: str | None = None) -> str:
     tenant_id = TENANT_ID_PREFIX + str(uuid.uuid4())
     try:
         # Provision tenant on data plane
-        await provision_tenant(tenant_id, email, referral_source)
+        await provision_tenant(tenant_id, email)
         # Notify control plane
-        await notify_control_plane(tenant_id, email)
+        await notify_control_plane(tenant_id, email, referral_source)
     except Exception as e:
         logger.error(f"Tenant provisioning failed: {e}")
         await rollback_tenant_provisioning(tenant_id)
@@ -119,14 +119,18 @@ async def provision_tenant(tenant_id: str, email: str) -> None:
             CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
 
-async def notify_control_plane(tenant_id: str, email: str) -> None:
+async def notify_control_plane(
+    tenant_id: str, email: str, referral_source: str | None = None
+) -> None:
     logger.info("Fetching billing information")
     token = generate_data_plane_token()
     headers = {
         "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
     }
-    payload = TenantCreationPayload(tenant_id=tenant_id, email=email)
+    payload = TenantCreationPayload(
+        tenant_id=tenant_id, email=email, referral_source=referral_source
+    )
 
     async with aiohttp.ClientSession() as session:
         async with session.post(

backend/tests/integration/multitenant_tests/tenants/test_tenant_creation.py

@@ -11,7 +11,7 @@
 
 # Test flow from creating tenant to registering as a user
 def test_tenant_creation(reset_multitenant: None) -> None:
-    TenantManager.create("tenant_dev", "[email protected]")
+    TenantManager.create("tenant_dev", "[email protected]", "Data Plane Registration")
     test_user: DATestUser = UserManager.create(name="test", email="[email protected]")
 
     assert UserManager.verify_role(test_user, UserRole.ADMIN)

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

@@ -14,9 +14,11 @@ import { Spinner } from "@/components/Spinner";
 export function EmailPasswordForm({
   isSignup = false,
   shouldVerify,
+  referralSource,
 }: {
   isSignup?: boolean;
   shouldVerify?: boolean;
+  referralSource?: string;
 }) {
   const router = useRouter();
   const { popup, setPopup } = usePopup();
@@ -39,7 +41,11 @@ export function EmailPasswordForm({
           if (isSignup) {
             // login is fast, no need to show a spinner
             setIsWorking(true);
-            const response = await basicSignup(values.email, values.password);
+            const response = await basicSignup(
+              values.email,
+              values.password,
+              referralSource
+            );
 
             if (!response.ok) {
               const errorDetail = (await response.json()).detail;

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

@@ -1,4 +1,3 @@
-import Cookies from "js-cookie";
 import { AuthType } from "@/lib/constants";
 import { FaGoogle } from "react-icons/fa";
 
@@ -41,7 +40,9 @@ export function SignInButton({
 
   const url = new URL(authorizeUrl);
 
+  // if (referralSource) {
   url.searchParams.set("referral_source", "docs");
+  // }
 
   const finalAuthorizeUrl = url.toString();
 
@@ -51,19 +52,12 @@ export function SignInButton({
     throw new Error(`Unhandled authType: ${authType}`);
   }
 
-  const handleClick = () => {
-    Cookies.set("referral_source", "docs", {
-      sameSite: "lax",
-      expires: new Date(Date.now() + 1000 * 60 * 60 * 24 * 30),
-    });
-    window.location.href = finalAuthorizeUrl;
-  };
   return (
-    <button
+    <a
       className="mx-auto mt-6 py-3 w-72 text-text-100 bg-accent flex rounded cursor-pointer hover:bg-indigo-800"
-      onClick={handleClick}
+      href={finalAuthorizeUrl}
     >
       {button}
-    </button>
+    </a>
   );
 }

web/src/app/auth/signup/ReferralSourceSelector.tsx

@@ -11,12 +11,10 @@ import {
 
 interface ReferralSourceSelectorProps {
   defaultValue?: string;
-  onChange: (value: string) => void;
 }
 
 const ReferralSourceSelector: React.FC<ReferralSourceSelectorProps> = ({
   defaultValue,
-  onChange,
 }) => {
   const [referralSource, setReferralSource] = useState(defaultValue);
 
@@ -36,7 +34,12 @@ const ReferralSourceSelector: React.FC<ReferralSourceSelectorProps> = ({
 
   const handleChange = (value: string) => {
     setReferralSource(value);
-    onChange(value);
+    const cookies = require("js-cookie");
+    cookies.set("referral_source", value, {
+      expires: 365,
+      path: "/",
+      sameSite: "strict",
+    });
   };
 
   return (

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

@@ -1,3 +1,4 @@
+import { HealthCheckBanner } from "@/components/health/healthcheck";
 import { User } from "@/lib/types";
 import {
   getCurrentUserSS,
@@ -6,13 +7,14 @@ import {
   getAuthUrlSS,
 } from "@/lib/userSS";
 import { redirect } from "next/navigation";
-import SignupPageContent from "./SignupPageContent";
-
-const Page = async (props: {
-  searchParams?: Promise<{ [key: string]: string | string[] | undefined }>;
-}) => {
-  const searchParams = await props.searchParams;
+import { EmailPasswordForm } from "../login/EmailPasswordForm";
+import Text from "@/components/ui/text";
+import Link from "next/link";
+import { SignInButton } from "../login/SignInButton";
+import AuthFlowContainer from "@/components/auth/AuthFlowContainer";
+import ReferralSourceSelector from "./ReferralSourceSelector";
 
+const Page = async () => {
   // catch cases where the backend is completely unreachable here
   // without try / catch, will just raise an exception and the page
   // will not render
@@ -51,18 +53,45 @@ const Page = async (props: {
     authUrl = await getAuthUrlSS(authTypeMetadata.authType, null);
   }
 
-  // Initialize referralSource with the value from URL parameters
-  const referralSource = searchParams
-    ? (searchParams.referral as string | undefined)
-    : undefined;
-
   return (
-    <SignupPageContent
-      authTypeMetadata={authTypeMetadata}
-      cloud={cloud}
-      authUrl={authUrl}
-      initialReferralSource={referralSource}
-    />
+    <AuthFlowContainer>
+      <HealthCheckBanner />
+
+      <>
+        <div className="absolute top-10x w-full"></div>
+        <div className="flex w-full flex-col justify-center">
+          <h2 className="text-center text-xl text-strong font-bold">
+            {cloud ? "Complete your sign up" : "Sign Up for Danswer"}
+          </h2>
+
+          {cloud && <ReferralSourceSelector />}
+          {cloud && authUrl && (
+            <div className="w-full justify-center">
+              <SignInButton authorizeUrl={authUrl} authType="cloud" />
+              <div className="flex items-center w-full my-4">
+                <div className="flex-grow border-t border-background-300"></div>
+                <span className="px-4 text-gray-500">or</span>
+                <div className="flex-grow border-t border-background-300"></div>
+              </div>
+            </div>
+          )}
+
+          <EmailPasswordForm
+            isSignup
+            shouldVerify={authTypeMetadata?.requiresVerification}
+          />
+
+          <div className="flex">
+            <Text className="mt-4 mx-auto">
+              Already have an account?{" "}
+              <Link href="/auth/login" className="text-link font-medium">
+                Log In
+              </Link>
+            </Text>
+          </div>
+        </div>
+      </>
+    </AuthFlowContainer>
   );
 };
 

web/src/lib/user.ts

@@ -43,7 +43,11 @@ export const basicLogin = async (
   return response;
 };
 
-export const basicSignup = async (email: string, password: string) => {
+export const basicSignup = async (
+  email: string,
+  password: string,
+  referralSource?: string
+) => {
   const response = await fetch("/api/auth/register", {
     method: "POST",
     credentials: "include",
@@ -54,6 +58,7 @@ export const basicSignup = async (email: string, password: string) => {
       email,
       username: email,
       password,
+      referral_source: referralSource,
     }),
   });
   return response;