Improve SSO login flow (#4206)

* Improve SSO login flow

* Expose initiate login url in SSO drawer

* ✨

front/components/workspace/connection.tsx

@@ -28,8 +28,9 @@ interface EnterpriseConnectionDetailsProps {
 }
 
 export interface EnterpriseConnectionStrategyDetails {
-  strategy: SupportedEnterpriseConnectionStrategies;
   callbackUrl: string;
+  initiateLoginUrl: string;
+  strategy: SupportedEnterpriseConnectionStrategies;
 }
 
 export function EnterpriseConnectionDetails({
@@ -170,7 +171,7 @@ function CreateOktaEnterpriseConnectionModal({
       domain?: string;
     }>({});
 
-  const { callbackUrl } = strategyDetails;
+  const { callbackUrl, initiateLoginUrl } = strategyDetails;
 
   const sendNotification = useContext(SendNotificationsContext);
 
@@ -225,15 +226,19 @@ function CreateOktaEnterpriseConnectionModal({
             Callback URL:
             <Input
               name="Callback URL"
-              placeholder="callback url"
+              placeholder="Callback url"
               value={callbackUrl}
               disabled={true}
-              onChange={(value) =>
-                setEnterpriseConnectionDetails({
-                  ...enterpriseConnectionDetails,
-                  domain: value,
-                })
-              }
+              className="max-w-sm"
+            />
+          </Page.P>
+          <Page.P>
+            Initiate login URI:
+            <Input
+              name="Initiate login URI"
+              placeholder="Initiate login URI"
+              value={initiateLoginUrl}
+              disabled={true}
               className="max-w-sm"
             />
           </Page.P>

front/lib/api/config.ts

@@ -1,6 +1,9 @@
 import { EnvironmentConfig } from "@dust-tt/types";
 
 const config = {
+  getAppUrl: (): string => {
+    return EnvironmentConfig.getEnvVariable("URL");
+  },
   getAuth0TenantUrl: (): string => {
     return EnvironmentConfig.getEnvVariable("AUTH0_TENANT_DOMAIN_URL");
   },

front/lib/api/enterprise_connection.ts

@@ -5,16 +5,30 @@ import { ManagementClient } from "auth0";
 import config from "@app/lib/api/config";
 import type { Authenticator } from "@app/lib/auth";
 
-const management = new ManagementClient({
-  domain: config.getAuth0TenantUrl(),
-  clientId: config.getAuth0M2MClientId(),
-  clientSecret: config.getAuth0M2MClientSecret(),
-});
+let auth0ManagemementClient: ManagementClient | null = null;
+
+function getAuth0ManagemementClient(): ManagementClient {
+  if (!auth0ManagemementClient) {
+    auth0ManagemementClient = new ManagementClient({
+      domain: config.getAuth0TenantUrl(),
+      clientId: config.getAuth0M2MClientId(),
+      clientSecret: config.getAuth0M2MClientSecret(),
+    });
+  }
+
+  return auth0ManagemementClient;
+}
 
 function makeEnterpriseConnectionName(workspaceId: string) {
   return `workspace-${workspaceId}`;
 }
 
+export function makeEnterpriseConnectionInitiateLoginUrl(workspaceId: string) {
+  return `${config.getAppUrl()}/api/auth/login?connection=${makeEnterpriseConnectionName(
+    workspaceId
+  )}`;
+}
+
 export async function getEnterpriseConnectionForWorkspace(
   auth: Authenticator,
   strategy: SupportedEnterpriseConnectionStrategies = "okta"
@@ -26,7 +40,7 @@ export async function getEnterpriseConnectionForWorkspace(
 
   // This endpoint supports fetching up to 1000 connections in one page.
   // In the future, consider implementing pagination to handle larger datasets.
-  const connections = await management.connections.getAll({
+  const connections = await getAuth0ManagemementClient().connections.getAll({
     strategy: [strategy],
   });
 
@@ -54,7 +68,7 @@ export async function createEnterpriseConnection(
   }
 
   const { sId } = owner;
-  const connection = await management.connections.create({
+  const connection = await getAuth0ManagemementClient().connections.create({
     name: makeEnterpriseConnectionName(sId),
     display_name: makeEnterpriseConnectionName(sId),
     strategy: connectionDetails.strategy,
@@ -93,7 +107,7 @@ export async function deleteEnterpriseConnection(
     throw new Error("Enterprise connection not found.");
   }
 
-  return management.connections.delete({
+  return getAuth0ManagemementClient().connections.delete({
     id: existingConnection.id,
   });
 }

front/pages/api/auth/[auth0].ts

@@ -1,9 +1,21 @@
 import { handleAuth, handleLogin } from "@auth0/nextjs-auth0";
+import type { AuthorizationParameters } from "@auth0/nextjs-auth0/dist/auth0-session";
 
 export default handleAuth({
-  login: handleLogin({
-    authorizationParams: {
+  login: handleLogin((req) => {
+    const connection = "query" in req ? req.query.connection : undefined;
+
+    const defaultAuthorizationParams: Partial<AuthorizationParameters> = {
       scope: "openid profile email",
-    },
+    };
+
+    // Set the Auth0 connection based on the provided connection param, redirecting the user to the correct screen.
+    if (connection) {
+      defaultAuthorizationParams.connection = connection;
+    }
+
+    return {
+      authorizationParams: defaultAuthorizationParams,
+    };
   }),
 });

front/pages/w/[wId]/members/index.tsx

@@ -35,6 +35,7 @@ import { SendNotificationsContext } from "@app/components/sparkle/Notification";
 import type { EnterpriseConnectionStrategyDetails } from "@app/components/workspace/connection";
 import { EnterpriseConnectionDetails } from "@app/components/workspace/connection";
 import config from "@app/lib/api/config";
+import { makeEnterpriseConnectionInitiateLoginUrl } from "@app/lib/api/enterprise_connection";
 import {
   checkWorkspaceSeatAvailabilityUsingAuth,
   getWorkspaceVerifiedDomain,
@@ -80,9 +81,11 @@ export const getServerSideProps = withDefaultGetServerSidePropsRequirements<{
 
   const enterpriseConnectionStrategyDetails: EnterpriseConnectionStrategyDetails =
     {
-      strategy: "okta",
       callbackUrl: config.getAuth0TenantUrl(),
+      initiateLoginUrl: makeEnterpriseConnectionInitiateLoginUrl(owner.sId),
+      strategy: "okta",
     };
+
   return {
     props: {
       user,