(proxy ui sso flow) - fix invite user sso flow (#6093)

* return if sso setup on ui_settings

* use helper to get invite link

litellm/proxy/management_endpoints/ui_sso.py

@@ -23,6 +23,7 @@ from litellm.proxy._types import (
     SSOUserDefinedValues,
     UserAPIKeyAuth,
 )
+from litellm.proxy.auth.auth_utils import _has_user_setup_sso
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.admin_ui_utils import (
     admin_ui_disabled,
@@ -640,6 +641,7 @@ async def get_ui_settings(request: Request):
 
     _proxy_base_url = os.getenv("PROXY_BASE_URL", None)
     _logout_url = os.getenv("PROXY_LOGOUT_URL", None)
+    _is_sso_enabled = _has_user_setup_sso()
 
     default_team_disabled = general_settings.get("default_team_disabled", False)
     if "PROXY_DEFAULT_TEAM_DISABLED" in os.environ:
@@ -650,4 +652,5 @@ async def get_ui_settings(request: Request):
         "PROXY_BASE_URL": _proxy_base_url,
         "PROXY_LOGOUT_URL": _logout_url,
         "DEFAULT_TEAM_DISABLED": default_team_disabled,
+        "SSO_ENABLED": _is_sso_enabled,
     }

ui/litellm-dashboard/src/components/create_user_button.tsx

@@ -17,6 +17,7 @@ import {
   userCreateCall,
   modelAvailableCall,
   invitationCreateCall,
+  getProxyBaseUrlAndLogoutUrl,
 } from "./networking";
 const { Option } = Select;
 
@@ -27,12 +28,21 @@ interface CreateuserProps {
   possibleUIRoles: null | Record<string, Record<string, string>>;
 }
 
+// Define an interface for the UI settings
+interface UISettings {
+  PROXY_BASE_URL: string | null;
+  PROXY_LOGOUT_URL: string | null;
+  DEFAULT_TEAM_DISABLED: boolean;
+  SSO_ENABLED: boolean;
+}
+
 const Createuser: React.FC<CreateuserProps> = ({
   userID,
   accessToken,
   teams,
   possibleUIRoles,
 }) => {
+  const [uiSettings, setUISettings] = useState<UISettings | null>(null);
   const [form] = Form.useForm();
   const [isModalVisible, setIsModalVisible] = useState(false);
   const [apiuser, setApiuser] = useState<string | null>(null);
@@ -70,11 +80,19 @@ const Createuser: React.FC<CreateuserProps> = ({
 
         // Assuming modelDataResponse.data contains an array of model names
         setUserModels(availableModels);
+
+        // get ui settings 
+        const uiSettingsResponse = await getProxyBaseUrlAndLogoutUrl(accessToken);
+        console.log("uiSettingsResponse:", uiSettingsResponse);
+        
+        setUISettings(uiSettingsResponse);
       } catch (error) {
         console.error("Error fetching model data:", error);
       }
     };
 
+    
+
     fetchData(); // Call the function to fetch model data when the component mounts
   }, []); // Empty dependency array to run only once
 
@@ -105,10 +123,33 @@ const Createuser: React.FC<CreateuserProps> = ({
       console.log("user create Response:", response);
       setApiuser(response["key"]);
       const user_id = response.data?.user_id || response.user_id;
-      invitationCreateCall(accessToken, user_id).then((data) => {
-        setInvitationLinkData(data);
+      
+      // only do invite link flow if sso is not enabled
+      if (!uiSettings?.SSO_ENABLED) {
+        invitationCreateCall(accessToken, user_id).then((data) => {
+          data.has_user_setup_sso = false;
+          setInvitationLinkData(data);
+          setIsInvitationLinkModalVisible(true);
+        });
+      } else {
+        // create an InvitationLink Object for this user for the SSO flow 
+        // for SSO the invite link is the proxy base url since the User just needs to login
+        const invitationLink: InvitationLink = {
+          id: crypto.randomUUID(), // Generate a unique ID
+          user_id: user_id,
+          is_accepted: false,
+          accepted_at: null,
+          expires_at: new Date(Date.now() + 7 * 24 * 60 * 60 * 1000), // Set expiry to 7 days from now
+          created_at: new Date(),
+          created_by: userID, // Assuming userID is the current user creating the invitation
+          updated_at: new Date(),
+          updated_by: userID,
+          has_user_setup_sso: true,
+        };
+        setInvitationLinkData(invitationLink);
         setIsInvitationLinkModalVisible(true);
-      });
+      }
+
       message.success("API user Created");
       form.resetFields();
       localStorage.removeItem("userData" + userID);

ui/litellm-dashboard/src/components/onboarding_link.tsx

@@ -21,6 +21,7 @@ export interface InvitationLink {
   created_by: string;
   updated_at: Date;
   updated_by: string;
+  has_user_setup_sso: boolean;
 }
 
 interface OnboardingProps {
@@ -47,6 +48,13 @@ const OnboardingModal: React.FC<OnboardingProps> = ({
     setIsInvitationLinkModalVisible(false);
   };
 
+  const getInvitationUrl = () => {
+    if (invitationLinkData?.has_user_setup_sso) {
+      return `${baseUrl}/ui`;
+    }
+    return `${baseUrl}/ui?invitation_id=${invitationLinkData?.id}`;
+  };
+
   return (
     <Modal
       title="Invitation Link"
@@ -67,13 +75,13 @@ const OnboardingModal: React.FC<OnboardingProps> = ({
       <div className="flex justify-between pt-5 pb-2">
         <Text>Invitation Link</Text>
         <Text>
-          {baseUrl}/ui?invitation_id={invitationLinkData?.id}
+        <Text>{getInvitationUrl()}</Text>
         </Text>
       </div>
       <div className="flex justify-end mt-5">
         <div></div>
         <CopyToClipboard
-          text={`${baseUrl}/ui?invitation_id=${invitationLinkData?.id}`}
+          text={getInvitationUrl()}
           onCopy={() => message.success("Copied!")}
         >
           <Button variant="primary">Copy invitation link</Button>

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -32,6 +32,7 @@ export interface ProxySettings {
   PROXY_BASE_URL: string | null;
   PROXY_LOGOUT_URL: string | null;
   DEFAULT_TEAM_DISABLED: boolean;
+  SSO_ENABLED: boolean;
 }
 
 function getCookie(name: string) {