From 75b98a79091fcd46650872e6426a29fa762f3e9d Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 19 Sep 2025 17:02:49 -0700
Subject: [PATCH] fix(ui_sso.py): initial commit, adding checking token
 endpoint response for allowed team ids

Allows ui access control to work for the given team ids
---
 litellm/proxy/auth/handle_jwt.py             | 22 ++++++++++++++++----
 litellm/proxy/management_endpoints/ui_sso.py | 20 +++++++++++++++++-
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py
index 5f78efbdf4..3e18db2d02 100644
--- a/litellm/proxy/auth/handle_jwt.py
+++ b/litellm/proxy/auth/handle_jwt.py
@@ -163,6 +163,7 @@ class JWTHandler:
         return False
 
     def get_team_ids_from_jwt(self, token: dict) -> List[str]:
+
         if self.litellm_jwtauth.team_ids_jwt_field is not None:
             team_ids: Optional[List[str]] = get_nested_value(
                 data=token,
@@ -483,7 +484,18 @@ class JWTHandler:
         # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
         # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
         #   the key in different ways (e.g. HS* and RS*)."
-        algorithms = ["RS256", "RS384", "RS512", "PS256", "PS384", "PS512", "ES256", "ES384", "ES512", "EdDSA"]
+        algorithms = [
+            "RS256",
+            "RS384",
+            "RS512",
+            "PS256",
+            "PS384",
+            "PS512",
+            "ES256",
+            "ES384",
+            "ES512",
+            "EdDSA",
+        ]
 
         audience = os.getenv("JWT_AUDIENCE")
         decode_options = None
@@ -540,7 +552,9 @@ class JWTHandler:
                 raise Exception(f"Validation fails: {str(e)}")
         elif public_key is not None and isinstance(public_key, str):
             try:
-                cert = x509.load_pem_x509_certificate(public_key.encode(), default_backend())
+                cert = x509.load_pem_x509_certificate(
+                    public_key.encode(), default_backend()
+                )
 
                 # Extract public key
                 key = cert.public_key().public_bytes(
@@ -565,7 +579,7 @@ class JWTHandler:
                 raise Exception(f"Validation fails: {str(e)}")
 
         raise Exception("Invalid JWT Submitted")
-    
+
     async def close(self):
         await self.http_handler.close()
 
@@ -1214,4 +1228,4 @@ class JWTAuthManager:
             end_user_object=end_user_object,
             token=api_key,
             team_membership=team_membership_object,
-        )
\ No newline at end of file
+        )
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index 98165a5a0c..1900cb1575 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -193,7 +193,7 @@ def generic_response_convertor(
     response,
     jwt_handler: JWTHandler,
     sso_jwt_handler: Optional[JWTHandler] = None,
-):
+) -> CustomOpenID:
     generic_user_id_attribute_name = os.getenv(
         "GENERIC_USER_ID_ATTRIBUTE", "preferred_username"
     )
@@ -226,6 +226,7 @@ def generic_response_convertor(
 
     team_ids = jwt_handler.get_team_ids_from_jwt(cast(dict, response))
     all_teams.extend(team_ids)
+
     return CustomOpenID(
         id=response.get(generic_user_id_attribute_name),
         display_name=response.get(generic_user_display_name_attribute_name),
@@ -340,6 +341,22 @@ async def get_generic_sso_response(
             params={"include_client_id": generic_include_client_id},
             headers=additional_generic_sso_headers_dict,
         )
+
+        access_token_str: Optional[str] = generic_sso.access_token
+        if access_token_str and sso_jwt_handler:
+            import jwt
+
+            access_token_payload = jwt.decode(
+                access_token_str, options={"verify_signature": False}
+            )
+
+            result_team_ids: Optional[List[str]] = (
+                getattr(result, "team_ids", []) if result else []
+            )
+            if not result_team_ids:
+                team_ids = sso_jwt_handler.get_team_ids_from_jwt(access_token_payload)
+                setattr(result, "team_ids", team_ids)
+
     except Exception as e:
         verbose_proxy_logger.exception(
             f"Error verifying and processing generic SSO: {e}. Passed in headers: {additional_generic_sso_headers_dict}"
@@ -612,6 +629,7 @@ async def auth_callback(request: Request, state: Optional[str] = None):  # noqa:
             microsoft_client_id=microsoft_client_id,
             redirect_url=redirect_url,
         )
+
     elif generic_client_id is not None:
         result, received_response = await get_generic_sso_response(
             request=request,