From f694878155823a3694fc5ec9b6c8a9ee5b137656 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 26 Apr 2025 16:17:38 -0700
Subject: [PATCH] fix SSO (#10344)

---
 litellm/proxy/management_endpoints/ui_sso.py | 68 ++++++++++----------
 1 file changed, 35 insertions(+), 33 deletions(-)

diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index 7414488bc8a..163678ab853 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -356,24 +356,28 @@ async def get_user_info_from_db(
     user_defined_values: Optional[SSOUserDefinedValues],
 ) -> Optional[Union[LiteLLM_UserTable, NewUserResponse]]:
     try:
-        user_info: Optional[
-            Union[LiteLLM_UserTable, NewUserResponse]
-        ] = await get_existing_user_info_from_db(
-            user_id=cast(
-                Optional[str],
-                getattr(result, "id", None)
-                if not isinstance(result, dict)
-                else result.get("id", None),
-            ),
-            user_email=cast(
-                Optional[str],
-                getattr(result, "email", None)
-                if not isinstance(result, dict)
-                else result.get("email", None),
-            ),
-            prisma_client=prisma_client,
-            user_api_key_cache=user_api_key_cache,
-            proxy_logging_obj=proxy_logging_obj,
+        user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]] = (
+            await get_existing_user_info_from_db(
+                user_id=cast(
+                    Optional[str],
+                    (
+                        getattr(result, "id", None)
+                        if not isinstance(result, dict)
+                        else result.get("id", None)
+                    ),
+                ),
+                user_email=cast(
+                    Optional[str],
+                    (
+                        getattr(result, "email", None)
+                        if not isinstance(result, dict)
+                        else result.get("email", None)
+                    ),
+                ),
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
         )
         verbose_proxy_logger.debug(
             f"user_info: {user_info}; litellm.default_internal_user_params: {litellm.default_internal_user_params}"
@@ -666,9 +670,7 @@ async def auth_callback(request: Request):  # noqa: PLR0915
         litellm_dashboard_ui += "?login=success"
     verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")
     redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
-    redirect_response.set_cookie(
-        key="token", value=jwt_token, secure=True, httponly=True
-    )
+    redirect_response.set_cookie(key="token", value=jwt_token, secure=True)
     return redirect_response
 
 
@@ -705,9 +707,9 @@ async def insert_sso_user(
         if user_defined_values.get("max_budget") is None:
             user_defined_values["max_budget"] = litellm.max_internal_user_budget
         if user_defined_values.get("budget_duration") is None:
-            user_defined_values[
-                "budget_duration"
-            ] = litellm.internal_user_budget_duration
+            user_defined_values["budget_duration"] = (
+                litellm.internal_user_budget_duration
+            )
 
     if user_defined_values["user_role"] is None:
         user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
@@ -902,9 +904,9 @@ class SSOAuthenticationHandler:
                 if state:
                     redirect_params["state"] = state
                 elif "okta" in generic_authorization_endpoint:
-                    redirect_params[
-                        "state"
-                    ] = uuid.uuid4().hex  # set state param for okta - required
+                    redirect_params["state"] = (
+                        uuid.uuid4().hex
+                    )  # set state param for okta - required
                 return await generic_sso.get_login_redirect(**redirect_params)  # type: ignore
         raise ValueError(
             "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso"
@@ -1149,9 +1151,9 @@ class MicrosoftSSOHandler:
 
         # if user is trying to get the raw sso response for debugging, return the raw sso response
         if return_raw_sso_response:
-            original_msft_result[
-                MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY
-            ] = user_team_ids
+            original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = (
+                user_team_ids
+            )
             return original_msft_result or {}
 
         result = MicrosoftSSOHandler.openid_from_response(
@@ -1219,9 +1221,9 @@ class MicrosoftSSOHandler:
 
             # Fetch user membership from Microsoft Graph API
             all_group_ids = []
-            next_link: Optional[
-                str
-            ] = MicrosoftSSOHandler.graph_api_user_groups_endpoint
+            next_link: Optional[str] = (
+                MicrosoftSSOHandler.graph_api_user_groups_endpoint
+            )
             auth_headers = {"Authorization": f"Bearer {access_token}"}
             page_count = 0