From 51285c2fa981e0e3e633a8a053ed46f4d2d7c791 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 9 Apr 2025 16:04:38 -0700
Subject: [PATCH] add_user_to_teams_from_sso_response

---
 litellm/proxy/management_endpoints/ui_sso.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index bdbcf79b82b..fb5550e5d7d 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -495,8 +495,10 @@ async def auth_callback(request: Request):  # noqa: PLR0915
                 user_role = (
                     user_info.user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
                 )
-            sso_teams = getattr(result, "team_ids", [])
-            await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)
+            await SSOAuthenticationHandler.add_user_to_teams_from_sso_response(
+                result=result,
+                user_info=user_info,
+            )
 
     except Exception as e:
         verbose_proxy_logger.debug(
@@ -840,6 +842,20 @@ class SSOAuthenticationHandler:
             redirect_url += "/" + sso_callback_route
         return redirect_url
 
+    @staticmethod
+    async def add_user_to_teams_from_sso_response(
+        result: Optional[Union[CustomOpenID, OpenID, dict]],
+        user_info: Union[NewUserResponse, LiteLLM_UserTable],
+    ):
+        """
+        Adds the user as a team member to the teams specified in the SSO responses `team_ids` field
+
+
+        The `team_ids` field is populated by litellm after processing the SSO response
+        """
+        sso_teams = getattr(result, "team_ids", [])
+        await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)
+
 
 class MicrosoftSSOHandler:
     """