From 51af0d5d949086c85c4bfa9a6c00a96a08da2058 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 7 Oct 2024 12:32:08 +0530 Subject: [PATCH] (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 | 3 ++ .../src/components/create_user_button.tsx | 47 +++++++++++++++++-- .../src/components/onboarding_link.tsx | 12 ++++- .../src/components/user_dashboard.tsx | 1 + 4 files changed, 58 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index c0ad13c47..d2b0e551a 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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, } diff --git a/ui/litellm-dashboard/src/components/create_user_button.tsx b/ui/litellm-dashboard/src/components/create_user_button.tsx index dc8298d38..21b99c666 100644 --- a/ui/litellm-dashboard/src/components/create_user_button.tsx +++ b/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>; } +// 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 = ({ userID, accessToken, teams, possibleUIRoles, }) => { + const [uiSettings, setUISettings] = useState(null); const [form] = Form.useForm(); const [isModalVisible, setIsModalVisible] = useState(false); const [apiuser, setApiuser] = useState(null); @@ -70,11 +80,19 @@ const Createuser: React.FC = ({ // 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 = ({ 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); diff --git a/ui/litellm-dashboard/src/components/onboarding_link.tsx b/ui/litellm-dashboard/src/components/onboarding_link.tsx index 3889e1ef9..78f0dff8a 100644 --- a/ui/litellm-dashboard/src/components/onboarding_link.tsx +++ b/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 = ({ setIsInvitationLinkModalVisible(false); }; + const getInvitationUrl = () => { + if (invitationLinkData?.has_user_setup_sso) { + return `${baseUrl}/ui`; + } + return `${baseUrl}/ui?invitation_id=${invitationLinkData?.id}`; + }; + return ( = ({
Invitation Link - {baseUrl}/ui?invitation_id={invitationLinkData?.id} + {getInvitationUrl()}
message.success("Copied!")} > diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index a1866b6b8..7db57112b 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/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) {