feat(proxy): persist allowlisted OIDC claims in CLI SSO poll#28463
Merged
mateo-berri merged 13 commits intoMay 22, 2026
Conversation
Map CLI_SSO_CLAIM_MAP sources into user metadata and return scalar attribution_metadata from /sso/cli/poll. Build SSOUserDefinedValues in cli_sso_callback so first-time CLI logins can upsert users. Add mock OIDC scripts and tests for claim extraction and poll exposure. Co-authored-by: Cursor <[email protected]>
Co-authored-by: Cursor <[email protected]>
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Contributor
Greptile SummaryThis PR extends the CLI SSO callback and poll flow to map allowlisted OIDC claims (via
Confidence Score: 4/5Safe to merge with minor attention to the ordering of the user_custom_sso hook relative to the restricted-group guard. The auth-adjacent changes (group check enforcement, SSOUserDefinedValues on first login) are correct and well-tested. The one ordering concern — that the operator-defined user_custom_sso hook fires before unauthorized users are rejected by the group check — is a contained, low-probability issue that only matters when both features are configured simultaneously. No data-loss or auth-bypass risk was found. litellm/proxy/management_endpoints/ui_sso.py — specifically the ordering of _build_cli_sso_user_defined_values vs verify_user_in_restricted_sso_group in cli_sso_callback.
|
| Filename | Overview |
|---|---|
| litellm/proxy/management_endpoints/ui_sso.py | Adds CLI SSO attribution metadata persistence and poll exposure; refactors callback to pass SSOUserDefinedValues and enforce restricted-group check. user_custom_sso hook is invoked before the restricted-group guard, and _persist_cli_sso_user_metadata issues raw Prisma calls (already flagged). |
| litellm/constants.py | Adds CLI_SSO_CLAIM_MAP env var (with LITELLM_ prefix fallback) and CLI_SSO_CLAIM_MAX_SCALAR_LENGTH constant; consistent with existing patterns. |
| litellm/proxy/client/README.md | One-line update to the CLI polling step description to document the new attribution_metadata poll response field. |
| tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | Adds TestCliSsoAttributionMetadata class with unit and async tests covering claim map parsing, metadata filtering, group rejection, and poll response. Tests use mocks only. One test omits user_custom_sso=None patch (flagged in previous review). |
Reviews (5): Last reviewed commit: "fix(sso): update CLI SSO test for new re..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
- Replace _update_dictionary with a local recursive merge so string OIDC claim values that happen to look numeric are not silently coerced to int/float when persisting CLI SSO attribution metadata. - Use a local dot-path resolver in _extract_sso_claim_value so that source claim paths beginning with 'metadata.' are not silently stripped by get_nested_value (which is designed for LiteLLM JWT metadata, not arbitrary OIDC claims). Co-authored-by: Yassin Kortam <[email protected]>
|
|
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high mode and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Double
metadata.prefix stripping corrupts nested dest paths- Removed the redundant metadata. prefix strip in _set_nested_metadata_value, leaving _parse_cli_sso_claim_map as the single place where the prefix is normalized.
Preview (cb103861af)
diff --git a/litellm/constants.py b/litellm/constants.py
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1443,6 +1443,12 @@
or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS")
or 24
)
+# Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g.
+# "employment_type->acme_employment_type,org_info.department->department"
+CLI_SSO_CLAIM_MAP = (
+ os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or ""
+)
+CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024
########################### UI SESSION DURATION ###########################
# Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d"
diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md
--- a/litellm/proxy/client/README.md
+++ b/litellm/proxy/client/README.md
@@ -350,7 +350,7 @@
4. **User Authentication**: User completes SSO authentication in browser
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
-7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
+7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready. When `CLI_SSO_CLAIM_MAP` is configured on the proxy, the poll response may include `attribution_metadata` (allowlisted scalar OIDC claims for client attribution).
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
### Benefits of This Approach
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -43,6 +43,8 @@
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
+ CLI_SSO_CLAIM_MAP,
+ CLI_SSO_CLAIM_MAX_SCALAR_LENGTH,
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
LITELLM_CLI_SOURCE_IDENTIFIER,
@@ -140,6 +142,21 @@
_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
+_CLI_SSO_SCALAR_TYPES = (str, int, float, bool)
+_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
+_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset(
+ {
+ "access_token",
+ "api_key",
+ "client_secret",
+ "id_token",
+ "password",
+ "private_key",
+ "refresh_token",
+ "secret",
+ "token",
+ }
+)
def _hash_cli_sso_secret(secret: str) -> str:
@@ -225,6 +242,232 @@
return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)
+def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]:
+ """
+ Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP.
+
+ Format: comma-separated ``source_claim->metadata_key`` pairs, e.g.
+ ``employment_type->acme_employment_type,org_info.department->department``.
+ Destination keys may use an optional ``metadata.`` prefix; values are stored
+ on the LiteLLM user's ``metadata`` JSON column.
+ """
+ claim_map_raw = CLI_SSO_CLAIM_MAP.strip()
+ if not claim_map_raw:
+ return []
+
+ parsed: List[Tuple[str, str]] = []
+ for entry in claim_map_raw.split(","):
+ entry = entry.strip()
+ if not entry or "->" not in entry:
+ continue
+ source_claim, dest_key = entry.split("->", 1)
+ source_claim = source_claim.strip()
+ dest_key = dest_key.strip()
+ if dest_key.startswith("metadata."):
+ dest_key = dest_key[len("metadata.") :]
+ if source_claim and dest_key:
+ parsed.append((source_claim, dest_key))
+ return parsed
+
+
+def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool:
+ if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key):
+ return False
+ lowered = dest_key.lower()
+ return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS)
+
+
+def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
+ if not isinstance(value, _CLI_SSO_SCALAR_TYPES):
+ return False
+ if isinstance(value, str):
+ if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH:
+ return False
+ if value.startswith("eyJ") and value.count(".") >= 2:
+ return False
+ return True
+
+
+def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]:
+ if isinstance(result, dict):
+ return result
+ if hasattr(result, "model_dump"):
+ dumped = result.model_dump()
+ if isinstance(dumped, dict):
+ return cast(Dict[str, Any], dumped)
+ return {}
+
+
+def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any:
+ """Resolve a dot-notation claim path against an SSO result dict.
+
+ Unlike ``get_nested_value``, this does not strip a leading ``metadata.``
+ prefix, since OIDC claims may legitimately use ``metadata`` as a top-level
+ key.
+ """
+ if not claim_path:
+ return None
+ if claim_path in data:
+ return data[claim_path]
+ placeholder = "\x00"
+ parts = claim_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = data
+ for part in parts:
+ if isinstance(current, dict) and part in current:
+ current = current[part]
+ else:
+ return None
+ return current
+
+
+def _extract_sso_claim_value(
+ result: Union[CustomOpenID, OpenID, dict], claim_path: str
+) -> Any:
+ extra_fields = getattr(result, "extra_fields", None)
+ if isinstance(extra_fields, dict) and claim_path in extra_fields:
+ return extra_fields[claim_path]
+
+ if isinstance(result, dict):
+ return _get_nested_claim_value(result, claim_path)
+
+ result_dict = _sso_result_to_dict(result)
+ return _get_nested_claim_value(result_dict, claim_path)
+
+
+def _set_nested_metadata_value(
+ metadata: Dict[str, Any], key_path: str, value: Any
+) -> None:
+ placeholder = "\x00"
+ parts = key_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = metadata
+ for part in parts[:-1]:
+ existing = current.get(part)
+ if not isinstance(existing, dict):
+ existing = {}
+ current[part] = existing
+ current = existing
+ current[parts[-1]] = value
+
+
+def _flatten_cli_sso_metadata_for_poll(
+ metadata: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ """Expose scalar attribution metadata as a flat dict for CLI poll responses."""
+
+ def _walk(prefix: str, value: Any) -> Dict[str, Union[str, int, float, bool]]:
+ flattened: Dict[str, Union[str, int, float, bool]] = {}
+ if isinstance(value, dict):
+ for key, nested in value.items():
+ nested_prefix = f"{prefix}.{key}" if prefix else key
+ flattened.update(_walk(nested_prefix, nested))
+ return flattened
+ if _is_safe_cli_sso_scalar_claim_value(value):
+ flattened[prefix] = value
+ return flattened
+
+ return _walk("", metadata)
+
+
+def build_cli_sso_attribution_metadata(
+ result: Union[CustomOpenID, OpenID, dict],
+) -> Dict[str, Any]:
+ """
+ Build allowlisted, non-secret scalar attribution metadata from an SSO result.
+
+ Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and
+ may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID.
+ """
+ claim_map = _parse_cli_sso_claim_map()
+ if not claim_map:
+ return {}
+
+ metadata: Dict[str, Any] = {}
+ for source_claim, dest_key in claim_map:
+ if not _is_safe_cli_sso_metadata_dest_key(dest_key):
+ verbose_proxy_logger.debug(
+ f"Skipping unsafe CLI SSO metadata destination key: {dest_key}"
+ )
+ continue
+
+ raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim)
+ if not _is_safe_cli_sso_scalar_claim_value(raw_value):
+ continue
+
+ _set_nested_metadata_value(
+ metadata=metadata, key_path=dest_key, value=raw_value
+ )
+
+ return metadata
+
+
+def _merge_cli_sso_attribution_metadata(
+ existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any]
+) -> Dict[str, Any]:
+ """Merge attribution metadata into existing user metadata in-place.
+
+ Preserves original value types (in particular, string claim values that
+ happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts
+ are merged recursively so attribution claims do not clobber unrelated keys
+ under the same parent.
+ """
+ for key, value in attribution_metadata.items():
+ if value is None:
+ continue
+ existing_value = existing_metadata.get(key)
+ if isinstance(value, dict) and isinstance(existing_value, dict):
+ _merge_cli_sso_attribution_metadata(existing_value, value)
+ else:
+ existing_metadata[key] = value
+ return existing_metadata
+
+
+async def _persist_cli_sso_user_metadata(
+ prisma_client: PrismaClient,
+ user_id: str,
+ attribution_metadata: Dict[str, Any],
+) -> None:
+ if not attribution_metadata:
+ return
+
+ try:
+ user_row = await prisma_client.db.litellm_usertable.find_unique(
+ where={"user_id": user_id}
+ )
+ existing_metadata: Dict[str, Any] = {}
+ if user_row is not None:
+ row_metadata = user_row.metadata
+ if isinstance(row_metadata, dict):
+ existing_metadata = deepcopy(row_metadata)
+
+ merged_metadata = _merge_cli_sso_attribution_metadata(
+ existing_metadata=existing_metadata,
+ attribution_metadata=attribution_metadata,
+ )
+ await prisma_client.db.litellm_usertable.update_many(
+ where={"user_id": user_id},
+ data={"metadata": merged_metadata},
+ )
+ verbose_proxy_logger.info(
+ f"Persisted CLI SSO attribution metadata for user {user_id}: "
+ f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}"
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}"
+ )
+
+
+def _cli_poll_attribution_metadata_from_session(
+ session_data: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ stored = session_data.get("attribution_metadata")
+ if isinstance(stored, dict):
+ return _flatten_cli_sso_metadata_for_poll(stored)
+ return {}
+
+
def _render_cli_sso_verification_page(
verify_url: str, browser_complete_token: str
) -> str:
@@ -1722,21 +1965,44 @@
# After None check, cast to non-None type for type checker
result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result)
+ generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result(
- result=result_non_none
+ result=result_non_none,
+ generic_client_id=generic_client_id,
)
verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
+ user_id = parsed_openid_result.get("user_id")
+ user_email = parsed_openid_result.get("user_email")
+ user_role = parsed_openid_result.get("user_role")
+
+ user_defined_values: Optional[SSOUserDefinedValues] = None
+ from litellm.proxy.proxy_server import user_custom_sso
+
+ if user_custom_sso is not None:
+ if inspect.iscoroutinefunction(user_custom_sso):
+ user_defined_values = await user_custom_sso(result_non_none) # type: ignore
+ else:
+ raise ValueError("user_custom_sso must be a coroutine function")
+ elif user_id is not None:
+ user_defined_values = SSOUserDefinedValues(
+ models=[],
+ user_id=user_id,
+ user_email=user_email,
+ max_budget=litellm.max_internal_user_budget,
+ user_role=user_role,
+ budget_duration=litellm.internal_user_budget_duration,
+ )
+
try:
- # Get full user info from DB
user_info = await get_user_info_from_db(
result=result_non_none,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
- user_email=parsed_openid_result.get("user_email"),
- user_defined_values=None,
- alternate_user_id=parsed_openid_result.get("user_id"),
+ user_email=user_email,
+ user_defined_values=user_defined_values,
+ alternate_user_id=user_id,
)
if user_info is None:
@@ -1774,6 +2040,16 @@
f"Error fetching team details for CLI SSO session: {e}"
)
+ attribution_metadata = build_cli_sso_attribution_metadata(
+ result=result_non_none
+ )
+ if attribution_metadata:
+ await _persist_cli_sso_user_metadata(
+ prisma_client=prisma_client,
+ user_id=user_info.user_id,
+ attribution_metadata=attribution_metadata,
+ )
+
session_data = {
"user_id": user_info.user_id,
"user_role": user_info.user_role,
@@ -1782,6 +2058,7 @@
"teams": teams,
# Optional rich metadata for clients that want nicer display
"team_details": team_details,
+ "attribution_metadata": attribution_metadata,
}
flow["session_data"] = session_data
@@ -1873,13 +2150,19 @@
team_details_response = [
{"team_id": t, "team_alias": None} for t in user_teams
]
- return {
+ poll_response: Dict[str, Any] = {
"status": "ready",
"user_id": user_id,
"teams": user_teams,
"team_details": team_details_response,
"requires_team_selection": True,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
# Validate team_id if provided
if team_id is not None:
@@ -1912,7 +2195,7 @@
verbose_proxy_logger.info(
f"CLI JWT generated for user: {user_id}, team: {team_id}"
)
- return {
+ poll_response = {
"status": "ready",
"key": jwt_token,
"user_id": user_id,
@@ -1922,6 +2205,12 @@
# present nicer information if needed.
"team_details": user_team_details,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
else:
return {"status": "pending"}
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -5552,6 +5552,235 @@
assert result.extra_fields["another_missing"] is None
+class TestCliSsoAttributionMetadata:
+ """CLI SSO allowlisted OIDC claim persistence and poll exposure."""
+
+ def test_parse_cli_sso_claim_map(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->metadata.acme_employment_type, org_info.department -> department",
+ )
+ assert ui_sso._parse_cli_sso_claim_map() == [
+ ("employment_type", "acme_employment_type"),
+ ("org_info.department", "department"),
+ ]
+
+ def test_build_cli_sso_attribution_metadata_filters_non_scalars(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->acme_employment_type,access_token->should_drop,group->groups",
+ )
+
+ result = CustomOpenID(
+ id="user-1",
+ email="[email protected]",
+ display_name="User",
+ provider="generic",
+ team_ids=[],
+ extra_fields={
+ "employment_type": "full_time",
+ "access_token": "eyJhbGciOiJIUzI1NiJ9.payload.signature",
+ "group": ["team-a", "team-b"],
+ },
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(result=result)
+ assert metadata == {"acme_employment_type": "full_time"}
+
+ def test_build_cli_sso_attribution_metadata_from_oidc_dict(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "org_info.department->department",
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(
+ result={
+ "sub": "user-1",
+ "email": "[email protected]",
+ "org_info": {"department": "Engineering"},
+ }
+ )
+ assert metadata == {"department": "Engineering"}
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self):
+ """First CLI SSO login must supply SSOUserDefinedValues so upsert can create the user."""
+ from litellm.proxy._types import LiteLLM_UserTable
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ session_key = "cli-session-new-user"
+ mock_user_info = LiteLLM_UserTable(
+ user_id="cli-test-user",
+ user_role="internal_user",
+ teams=[],
+ models=[],
+ )
+ mock_sso_result = CustomOpenID(
+ id="cli-test-user",
+ email="[email protected]",
+ display_name="cli-test-user",
+ provider="generic",
+ team_ids=[],
+ )
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ get_user_info_mock = AsyncMock(return_value=mock_user_info)
+
+ with (
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ get_user_info_mock,
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ ):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key=session_key,
+ result=mock_sso_result,
+ )
+
+ get_user_info_mock.assert_awaited_once()
+ assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None
+ assert (
+ get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"]
+ == "cli-test-user"
+ )
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_persists_attribution_metadata(self, monkeypatch):
+ from litellm.proxy._types import LiteLLM_UserTable
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->acme_employment_type",
+ )
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ session_key = "cli-session-4567890"
+ mock_user_info = LiteLLM_UserTable(
+ user_id="test-user-123",
+ user_role="internal_user",
+ teams=["team1"],
+ models=["gpt-4"],
+ )
+ mock_sso_result = {
+ "user_email": "[email protected]",
+ "user_id": "test-user-123",
+ "employment_type": "contractor",
+ }
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ mock_prisma = MagicMock()
+ mock_prisma.db.litellm_usertable.find_unique = AsyncMock(
+ return_value=MagicMock(metadata={"auth_provider": "generic"})
+ )
+ mock_prisma.db.litellm_usertable.update_many = AsyncMock()
+
+ with (
+ patch.dict(
+ os.environ,
+ {
+ "PROXY_BASE_URL": "https://test.example.com",
+ "SERVER_ROOT_PATH": "",
+ },
+ ),
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ return_value=mock_user_info,
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch(
+ "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page",
+ return_value="<html>Success</html>",
+ ),
+ ):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key=session_key,
+ result=mock_sso_result,
+ )
+
+ flow_data = mock_cache.set_cache.call_args.kwargs["value"]
+ assert flow_data["session_data"]["attribution_metadata"] == {
+ "acme_employment_type": "contractor"
+ }
+ mock_prisma.db.litellm_usertable.update_many.assert_awaited_once()
+ update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs[
+ "data"
+ ]
+ assert update_data["metadata"]["acme_employment_type"] == "contractor"
+ assert update_data["metadata"]["auth_provider"] == "generic"
+
+ @pytest.mark.asyncio
+ async def test_cli_poll_key_returns_attribution_metadata(self, monkeypatch):
+ from litellm.proxy.management_endpoints.ui_sso import (
+ _hash_cli_sso_secret,
+ cli_poll_key,
+ )
+
+ session_key = "cli-session-789123"
+ session_data = {
+ "user_id": "test-user-456",
+ "user_role": "internal_user",
+ "teams": ["team-a", "team-b"],
+ "models": ["gpt-4"],
+ "attribution_metadata": {
+ "acme_employment_type": "full_time",
+ "org": {"cost_center": "CC-42"},
+ },
+ }
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": _hash_cli_sso_secret("poll-secret"),
+ "sso_complete": True,
+ "user_code_verified": True,
+ "session_data": session_data,
+ }
+
+ with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
+ result = await cli_poll_key(
+ key_id=session_key,
+ team_id=None,
+ x_litellm_cli_poll_secret="poll-secret",
+ )
+
+ assert result["attribution_metadata"] == {
+ "acme_employment_type": "full_time",
+ "org.cost_center": "CC-42",
+ }
+
+
class TestValidateReturnTo:
"""Tests for SSOAuthenticationHandler._validate_return_to"""You can send follow-ups to the cloud agent here.
The _parse_cli_sso_claim_map already strips the metadata. prefix from dest keys before reaching the setter. The duplicate strip in _set_nested_metadata_value was a no-op in normal flow but could mis-place values for dest keys like metadata.metadata.foo. Co-authored-by: Yassin Kortam <[email protected]>
Collaborator
Author
|
@greptile re review |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high mode and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Custom SSO handler errors bypass error handling block
- Moved the parsed_openid_result and _build_cli_sso_user_defined_values calls inside the existing try/except block so custom SSO handler errors are logged and wrapped in the friendly HTTP 500 response.
Preview (7cf692e6d7)
diff --git a/litellm/constants.py b/litellm/constants.py
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1443,6 +1443,12 @@
or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS")
or 24
)
+# Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g.
+# "employment_type->acme_employment_type,org_info.department->department"
+CLI_SSO_CLAIM_MAP = (
+ os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or ""
+)
+CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024
########################### UI SESSION DURATION ###########################
# Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d"
diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md
--- a/litellm/proxy/client/README.md
+++ b/litellm/proxy/client/README.md
@@ -350,7 +350,7 @@
4. **User Authentication**: User completes SSO authentication in browser
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
-7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
+7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready. When `CLI_SSO_CLAIM_MAP` is configured on the proxy, the poll response may include `attribution_metadata` (allowlisted scalar OIDC claims for client attribution).
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
### Benefits of This Approach
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -43,6 +43,8 @@
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
+ CLI_SSO_CLAIM_MAP,
+ CLI_SSO_CLAIM_MAX_SCALAR_LENGTH,
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
LITELLM_CLI_SOURCE_IDENTIFIER,
@@ -140,6 +142,21 @@
_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
+_CLI_SSO_SCALAR_TYPES = (str, int, float, bool)
+_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
+_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset(
+ {
+ "access_token",
+ "api_key",
+ "client_secret",
+ "id_token",
+ "password",
+ "private_key",
+ "refresh_token",
+ "secret",
+ "token",
+ }
+)
def _hash_cli_sso_secret(secret: str) -> str:
@@ -225,6 +242,232 @@
return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)
+def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]:
+ """
+ Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP.
+
+ Format: comma-separated ``source_claim->metadata_key`` pairs, e.g.
+ ``employment_type->acme_employment_type,org_info.department->department``.
+ Destination keys may use an optional ``metadata.`` prefix; values are stored
+ on the LiteLLM user's ``metadata`` JSON column.
+ """
+ claim_map_raw = CLI_SSO_CLAIM_MAP.strip()
+ if not claim_map_raw:
+ return []
+
+ parsed: List[Tuple[str, str]] = []
+ for entry in claim_map_raw.split(","):
+ entry = entry.strip()
+ if not entry or "->" not in entry:
+ continue
+ source_claim, dest_key = entry.split("->", 1)
+ source_claim = source_claim.strip()
+ dest_key = dest_key.strip()
+ if dest_key.startswith("metadata."):
+ dest_key = dest_key[len("metadata.") :]
+ if source_claim and dest_key:
+ parsed.append((source_claim, dest_key))
+ return parsed
+
+
+def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool:
+ if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key):
+ return False
+ lowered = dest_key.lower()
+ return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS)
+
+
+def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
+ if not isinstance(value, _CLI_SSO_SCALAR_TYPES):
+ return False
+ if isinstance(value, str):
+ if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH:
+ return False
+ if value.startswith("eyJ") and value.count(".") >= 2:
+ return False
+ return True
+
+
+def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]:
+ if isinstance(result, dict):
+ return result
+ if hasattr(result, "model_dump"):
+ dumped = result.model_dump()
+ if isinstance(dumped, dict):
+ return cast(Dict[str, Any], dumped)
+ return {}
+
+
+def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any:
+ """Resolve a dot-notation claim path against an SSO result dict.
+
+ Unlike ``get_nested_value``, this does not strip a leading ``metadata.``
+ prefix, since OIDC claims may legitimately use ``metadata`` as a top-level
+ key.
+ """
+ if not claim_path:
+ return None
+ if claim_path in data:
+ return data[claim_path]
+ placeholder = "\x00"
+ parts = claim_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = data
+ for part in parts:
+ if isinstance(current, dict) and part in current:
+ current = current[part]
+ else:
+ return None
+ return current
+
+
+def _extract_sso_claim_value(
+ result: Union[CustomOpenID, OpenID, dict], claim_path: str
+) -> Any:
+ extra_fields = getattr(result, "extra_fields", None)
+ if isinstance(extra_fields, dict) and claim_path in extra_fields:
+ return extra_fields[claim_path]
+
+ if isinstance(result, dict):
+ return _get_nested_claim_value(result, claim_path)
+
+ result_dict = _sso_result_to_dict(result)
+ return _get_nested_claim_value(result_dict, claim_path)
+
+
+def _set_nested_metadata_value(
+ metadata: Dict[str, Any], key_path: str, value: Any
+) -> None:
+ placeholder = "\x00"
+ parts = key_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = metadata
+ for part in parts[:-1]:
+ existing = current.get(part)
+ if not isinstance(existing, dict):
+ existing = {}
+ current[part] = existing
+ current = existing
+ current[parts[-1]] = value
+
+
+def _flatten_cli_sso_metadata_for_poll(
+ metadata: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ """Expose scalar attribution metadata as a flat dict for CLI poll responses."""
+
+ def _walk(prefix: str, value: Any) -> Dict[str, Union[str, int, float, bool]]:
+ flattened: Dict[str, Union[str, int, float, bool]] = {}
+ if isinstance(value, dict):
+ for key, nested in value.items():
+ nested_prefix = f"{prefix}.{key}" if prefix else key
+ flattened.update(_walk(nested_prefix, nested))
+ return flattened
+ if _is_safe_cli_sso_scalar_claim_value(value):
+ flattened[prefix] = value
+ return flattened
+
+ return _walk("", metadata)
+
+
+def build_cli_sso_attribution_metadata(
+ result: Union[CustomOpenID, OpenID, dict],
+) -> Dict[str, Any]:
+ """
+ Build allowlisted, non-secret scalar attribution metadata from an SSO result.
+
+ Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and
+ may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID.
+ """
+ claim_map = _parse_cli_sso_claim_map()
+ if not claim_map:
+ return {}
+
+ metadata: Dict[str, Any] = {}
+ for source_claim, dest_key in claim_map:
+ if not _is_safe_cli_sso_metadata_dest_key(dest_key):
+ verbose_proxy_logger.debug(
+ f"Skipping unsafe CLI SSO metadata destination key: {dest_key}"
+ )
+ continue
+
+ raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim)
+ if not _is_safe_cli_sso_scalar_claim_value(raw_value):
+ continue
+
+ _set_nested_metadata_value(
+ metadata=metadata, key_path=dest_key, value=raw_value
+ )
+
+ return metadata
+
+
+def _merge_cli_sso_attribution_metadata(
+ existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any]
+) -> Dict[str, Any]:
+ """Merge attribution metadata into existing user metadata in-place.
+
+ Preserves original value types (in particular, string claim values that
+ happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts
+ are merged recursively so attribution claims do not clobber unrelated keys
+ under the same parent.
+ """
+ for key, value in attribution_metadata.items():
+ if value is None:
+ continue
+ existing_value = existing_metadata.get(key)
+ if isinstance(value, dict) and isinstance(existing_value, dict):
+ _merge_cli_sso_attribution_metadata(existing_value, value)
+ else:
+ existing_metadata[key] = value
+ return existing_metadata
+
+
+async def _persist_cli_sso_user_metadata(
+ prisma_client: PrismaClient,
+ user_id: str,
+ attribution_metadata: Dict[str, Any],
+) -> None:
+ if not attribution_metadata:
+ return
+
+ try:
+ user_row = await prisma_client.db.litellm_usertable.find_unique(
+ where={"user_id": user_id}
+ )
+ existing_metadata: Dict[str, Any] = {}
+ if user_row is not None:
+ row_metadata = user_row.metadata
+ if isinstance(row_metadata, dict):
+ existing_metadata = deepcopy(row_metadata)
+
+ merged_metadata = _merge_cli_sso_attribution_metadata(
+ existing_metadata=existing_metadata,
+ attribution_metadata=attribution_metadata,
+ )
+ await prisma_client.db.litellm_usertable.update_many(
+ where={"user_id": user_id},
+ data={"metadata": merged_metadata},
+ )
+ verbose_proxy_logger.info(
+ f"Persisted CLI SSO attribution metadata for user {user_id}: "
+ f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}"
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}"
+ )
+
+
+def _cli_poll_attribution_metadata_from_session(
+ session_data: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ stored = session_data.get("attribution_metadata")
+ if isinstance(stored, dict):
+ return _flatten_cli_sso_metadata_for_poll(stored)
+ return {}
+
+
def _render_cli_sso_verification_page(
verify_url: str, browser_complete_token: str
) -> str:
@@ -1692,6 +1935,129 @@
)
+async def _build_cli_sso_user_defined_values(
+ result: Union[OpenID, dict],
+ parsed_openid_result: ParsedOpenIDResult,
+) -> Optional[SSOUserDefinedValues]:
+ from litellm.proxy.proxy_server import user_custom_sso
+
+ user_id = parsed_openid_result.get("user_id")
+ if user_custom_sso is not None:
+ if inspect.iscoroutinefunction(user_custom_sso):
+ return await user_custom_sso(result) # type: ignore
+ raise ValueError("user_custom_sso must be a coroutine function")
+ if user_id is None:
+ return None
+ return SSOUserDefinedValues(
+ models=[],
+ user_id=user_id,
+ user_email=parsed_openid_result.get("user_email"),
+ max_budget=litellm.max_internal_user_budget,
+ user_role=parsed_openid_result.get("user_role"),
+ budget_duration=litellm.internal_user_budget_duration,
+ )
+
+
+async def _fetch_cli_sso_team_details(
+ prisma_client: PrismaClient,
+ teams: List[str],
+) -> List[Dict[str, Any]]:
+ team_details: List[Dict[str, Any]] = []
+ try:
+ if teams:
+ prisma_teams = await prisma_client.db.litellm_teamtable.find_many(
+ where={"team_id": {"in": teams}}
+ )
+ for team_row in prisma_teams:
+ team_dict = team_row.model_dump()
+ team_details.append(
+ {
+ "team_id": team_dict.get("team_id"),
+ "team_alias": team_dict.get("team_alias"),
+ }
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Error fetching team details for CLI SSO session: {e}"
+ )
+ return team_details
+
+
+async def _complete_cli_sso_callback_session(
+ *,
+ request: Request,
+ key: str,
+ flow: dict,
+ result: Union[OpenID, dict],
+ parsed_openid_result: ParsedOpenIDResult,
+ user_defined_values: Optional[SSOUserDefinedValues],
+ prisma_client: PrismaClient,
+ user_api_key_cache: DualCache,
+ proxy_logging_obj: ProxyLogging,
+):
+ from fastapi.responses import HTMLResponse
+
+ user_id = parsed_openid_result.get("user_id")
+ user_email = parsed_openid_result.get("user_email")
+ user_info = await get_user_info_from_db(
+ result=result,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ user_email=user_email,
+ user_defined_values=user_defined_values,
+ alternate_user_id=user_id,
+ )
+ if user_info is None:
+ raise HTTPException(
+ status_code=500, detail="Failed to retrieve user information from SSO"
+ )
+
+ teams: List[str] = []
+ if hasattr(user_info, "teams") and user_info.teams:
+ teams = user_info.teams if isinstance(user_info.teams, list) else []
+
+ team_details = await _fetch_cli_sso_team_details(
+ prisma_client=prisma_client, teams=teams
+ )
+ attribution_metadata = build_cli_sso_attribution_metadata(result=result)
+ if attribution_metadata:
+ await _persist_cli_sso_user_metadata(
+ prisma_client=prisma_client,
+ user_id=user_info.user_id,
+ attribution_metadata=attribution_metadata,
+ )
+
+ flow["session_data"] = {
+ "user_id": user_info.user_id,
+ "user_role": user_info.user_role,
+ "models": user_info.models if hasattr(user_info, "models") else [],
+ "user_email": user_email,
+ "teams": teams,
+ "team_details": team_details,
+ "attribution_metadata": attribution_metadata,
+ }
+ flow["sso_complete"] = True
+ browser_complete_token = secrets.token_urlsafe(32)
+ flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token)
+ _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow)
+
+ verbose_proxy_logger.info(
+ f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
+ )
+ verify_url = get_custom_url(
+ request_base_url=str(request.base_url),
+ route=f"sso/cli/complete/{key}",
+ )
+ return HTMLResponse(
+ content=_render_cli_sso_verification_page(
+ verify_url=verify_url,
+ browser_complete_token=browser_complete_token,
+ ),
+ status_code=200,
+ )
+
+
async def cli_sso_callback(
request: Request,
key: Optional[str] = None,
@@ -1722,92 +2088,32 @@
# After None check, cast to non-None type for type checker
result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result)
- parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result(
- result=result_non_none
- )
- verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
-
try:
- # Get full user info from DB
- user_info = await get_user_info_from_db(
+ parsed_openid_result = (
+ SSOAuthenticationHandler._get_user_email_and_id_from_result(
+ result=result_non_none,
+ generic_client_id=os.getenv("GENERIC_CLIENT_ID", None),
+ )
+ )
+ verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
+ user_defined_values = await _build_cli_sso_user_defined_values(
result=result_non_none,
+ parsed_openid_result=parsed_openid_result,
+ )
+
+ return await _complete_cli_sso_callback_session(
+ request=request,
+ key=cast(str, key),
+ flow=flow,
+ result=result_non_none,
+ parsed_openid_result=parsed_openid_result,
+ user_defined_values=user_defined_values,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
- user_email=parsed_openid_result.get("user_email"),
- user_defined_values=None,
- alternate_user_id=parsed_openid_result.get("user_id"),
)
-
- if user_info is None:
- raise HTTPException(
- status_code=500, detail="Failed to retrieve user information from SSO"
- )
-
- # Get all teams from user_info - CLI will let user select which one
- teams: List[str] = []
- if hasattr(user_info, "teams") and user_info.teams:
- teams = user_info.teams if isinstance(user_info.teams, list) else []
-
- # Also fetch team aliases for a better CLI UX. We keep the original
- # "teams" list of IDs for backwards compatibility and add an
- # optional "team_details" field containing objects with both
- # team_id and team_alias.
- team_details: List[Dict[str, Any]] = []
- try:
- if teams:
- prisma_teams = await prisma_client.db.litellm_teamtable.find_many(
- where={"team_id": {"in": teams}}
- )
- for team_row in prisma_teams:
- team_dict = team_row.model_dump()
- team_details.append(
- {
- "team_id": team_dict.get("team_id"),
- "team_alias": team_dict.get("team_alias"),
- }
- )
- except Exception as e:
- # If anything goes wrong here, fall back gracefully without
- # impacting the SSO flow.
- verbose_proxy_logger.error(
- f"Error fetching team details for CLI SSO session: {e}"
- )
-
- session_data = {
- "user_id": user_info.user_id,
- "user_role": user_info.user_role,
- "models": user_info.models if hasattr(user_info, "models") else [],
- "user_email": parsed_openid_result.get("user_email"),
- "teams": teams,
- # Optional rich metadata for clients that want nicer display
- "team_details": team_details,
- }
-
- flow["session_data"] = session_data
- flow["sso_complete"] = True
- browser_complete_token = secrets.token_urlsafe(32)
- flow["browser_complete_token_hash"] = _hash_cli_sso_secret(
- browser_complete_token
- )
- _set_cli_sso_flow(login_id=cast(str, key), cache=user_api_key_cache, flow=flow)
-
- verbose_proxy_logger.info(
- f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
- )
-
- from fastapi.responses import HTMLResponse
-
- verify_url = get_custom_url(
- request_base_url=str(request.base_url),
- route=f"sso/cli/complete/{key}",
- )
- html_content = _render_cli_sso_verification_page(
- verify_url=verify_url,
- browser_complete_token=browser_complete_token,
- )
- return HTMLResponse(content=html_content, status_code=200)
-
+ except HTTPException:
+ raise
except Exception as e:
verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}")
raise HTTPException(
@@ -1873,13 +2179,19 @@
team_details_response = [
{"team_id": t, "team_alias": None} for t in user_teams
]
- return {
+ poll_response: Dict[str, Any] = {
"status": "ready",
"user_id": user_id,
"teams": user_teams,
"team_details": team_details_response,
"requires_team_selection": True,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
# Validate team_id if provided
if team_id is not None:
@@ -1912,7 +2224,7 @@
verbose_proxy_logger.info(
f"CLI JWT generated for user: {user_id}, team: {team_id}"
)
- return {
+ poll_response = {
"status": "ready",
"key": jwt_token,
"user_id": user_id,
@@ -1922,6 +2234,12 @@
# present nicer information if needed.
"team_details": user_team_details,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
else:
return {"status": "pending"}
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -5552,6 +5552,236 @@
assert result.extra_fields["another_missing"] is None
+class TestCliSsoAttributionMetadata:
+ """CLI SSO allowlisted OIDC claim persistence and poll exposure."""
+
+ def test_parse_cli_sso_claim_map(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->metadata.acme_employment_type, org_info.department -> department",
+ )
+ assert ui_sso._parse_cli_sso_claim_map() == [
+ ("employment_type", "acme_employment_type"),
+ ("org_info.department", "department"),
+ ]
+
+ def test_build_cli_sso_attribution_metadata_filters_non_scalars(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->acme_employment_type,access_token->should_drop,group->groups",
+ )
+
+ result = CustomOpenID(
+ id="user-1",
+ email="[email protected]",
+ display_name="User",
+ provider="generic",
+ team_ids=[],
+ extra_fields={
+ "employment_type": "full_time",
+ "access_token": "eyJhbGciOiJIUzI1NiJ9.payload.signature",
+ "group": ["team-a", "team-b"],
+ },
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(result=result)
+ assert metadata == {"acme_employment_type": "full_time"}
+
+ def test_build_cli_sso_attribution_metadata_from_oidc_dict(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "org_info.department->department",
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(
+ result={
+ "sub": "user-1",
+ "email": "[email protected]",
+ "org_info": {"department": "Engineering"},
+ }
+ )
+ assert metadata == {"department": "Engineering"}
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self):
+ """First CLI SSO login must supply SSOUserDefinedValues so upsert can create the user."""
+ from litellm.proxy._types import LiteLLM_UserTable
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ session_key = "cli-session-new-user"
+ mock_user_info = LiteLLM_UserTable(
+ user_id="cli-test-user",
+ user_role="internal_user",
+ teams=[],
+ models=[],
+ )
+ mock_sso_result = CustomOpenID(
+ id="cli-test-user",
+ email="[email protected]",
+ display_name="cli-test-user",
+ provider="generic",
+ team_ids=[],
+ )
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ get_user_info_mock = AsyncMock(return_value=mock_user_info)
+
+ with (
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ get_user_info_mock,
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ ):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key=session_key,
+ result=mock_sso_result,
+ )
+
+ get_user_info_mock.assert_awaited_once()
+ assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None
+ assert (
+ get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"]
+ == "cli-test-user"
+ )
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_persists_attribution_metadata(self, monkeypatch):
+ from litellm.proxy._types import LiteLLM_UserTable
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->acme_employment_type",
+ )
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ session_key = "cli-session-4567890"
+ mock_user_info = LiteLLM_UserTable(
+ user_id="test-user-123",
+ user_role="internal_user",
+ teams=["team1"],
+ models=["gpt-4"],
+ )
+ mock_sso_result = {
+ "user_email": "[email protected]",
+ "user_id": "test-user-123",
+ "employment_type": "contractor",
+ }
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ mock_prisma = MagicMock()
+ mock_prisma.db.litellm_usertable.find_unique = AsyncMock(
+ return_value=MagicMock(metadata={"auth_provider": "generic"})
+ )
+ mock_prisma.db.litellm_usertable.update_many = AsyncMock()
+
+ with (
+ patch.dict(
+ os.environ,
+ {
+ "PROXY_BASE_URL": "https://test.example.com",
+ "SERVER_ROOT_PATH": "",
+ },
+ ),
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ return_value=mock_user_info,
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ patch(
+ "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page",
+ return_value="<html>Success</html>",
+ ),
+ ):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key=session_key,
+ result=mock_sso_result,
+ )
+
+ flow_data = mock_cache.set_cache.call_args.kwargs["value"]
+ assert flow_data["session_data"]["attribution_metadata"] == {
+ "acme_employment_type": "contractor"
+ }
+ mock_prisma.db.litellm_usertable.update_many.assert_awaited_once()
+ update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs[
+ "data"
+ ]
+ assert update_data["metadata"]["acme_employment_type"] == "contractor"
+ assert update_data["metadata"]["auth_provider"] == "generic"
+
+ @pytest.mark.asyncio
+ async def test_cli_poll_key_returns_attribution_metadata(self, monkeypatch):
+ from litellm.proxy.management_endpoints.ui_sso import (
+ _hash_cli_sso_secret,
+ cli_poll_key,
+ )
+
+ session_key = "cli-session-789123"
+ session_data = {
+ "user_id": "test-user-456",
+ "user_role": "internal_user",
+ "teams": ["team-a", "team-b"],
+ "models": ["gpt-4"],
+ "attribution_metadata": {
+ "acme_employment_type": "full_time",
+ "org": {"cost_center": "CC-42"},
+ },
+ }
... diff truncated: showing 800 of 823 linesYou can send follow-ups to the cloud agent here.
…nt error handling Co-authored-by: Yassin Kortam <[email protected]>
Collaborator
Contributor
PR overviewCLI SSO attribution metadata persistenceThis PR adds an operator-configured OIDC claim allowlist for CLI SSO, persists scalar attribution metadata on the user record, and returns it only through the existing poll-secret-protected CLI polling flow. I checked the callback, restricted-group enforcement, metadata filtering, persistence, and poll response paths and did not find a new attacker-controlled security issue. Security review
Risk: 2/10 |
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
Contributor
|
@Sameerlite — heads up: the Code Quality Checks, LiteLLM Linting, Documentation Validation, and Integrations checks are failing on this PR. They appear related to your changes, since they pass on other PRs against the same base branch (the broader CircleCI failures look pre-existing on the base branch and unrelated). Could you take a look and fix the failing checks before this gets reviewed? Thanks! |
Apply verify_user_in_restricted_sso_group before CLI session completion and user upsert, matching the UI SSO path. Re-raise ProxyException so restricted-group denials return 403 instead of 500. Co-authored-by: Cursor <[email protected]>
… merge Use stack-based flatten/merge to satisfy recursive_detector CI. Fix mypy types for UserApiKeyCache and user_id on CLI SSO session completion. Co-authored-by: Cursor <[email protected]>
Collaborator
Author
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high mode and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Nested claim extraction fails for CustomOpenID extra_fields
- Added a dotted-path fallback using _get_nested_claim_value on extra_fields when the exact flat key lookup misses, so parent-object captures like {'org_info': {'department': '...'}} now resolve org_info.department.
Preview (a0ecb87fda)
diff --git a/litellm/constants.py b/litellm/constants.py
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1443,6 +1443,12 @@
or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS")
or 24
)
+# Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g.
+# "employment_type->acme_employment_type,org_info.department->department"
+CLI_SSO_CLAIM_MAP = (
+ os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or ""
+)
+CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024
########################### UI SESSION DURATION ###########################
# Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d"
diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md
--- a/litellm/proxy/client/README.md
+++ b/litellm/proxy/client/README.md
@@ -350,7 +350,7 @@
4. **User Authentication**: User completes SSO authentication in browser
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
-7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
+7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready. When `CLI_SSO_CLAIM_MAP` is configured on the proxy, the poll response may include `attribution_metadata` (allowlisted scalar OIDC claims for client attribution).
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
### Benefits of This Approach
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -43,6 +43,8 @@
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
+ CLI_SSO_CLAIM_MAP,
+ CLI_SSO_CLAIM_MAX_SCALAR_LENGTH,
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
LITELLM_CLI_SOURCE_IDENTIFIER,
@@ -140,6 +142,21 @@
_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
+_CLI_SSO_SCALAR_TYPES = (str, int, float, bool)
+_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
+_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset(
+ {
+ "access_token",
+ "api_key",
+ "client_secret",
+ "id_token",
+ "password",
+ "private_key",
+ "refresh_token",
+ "secret",
+ "token",
+ }
+)
def _hash_cli_sso_secret(secret: str) -> str:
@@ -225,6 +242,239 @@
return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)
+def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]:
+ """
+ Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP.
+
+ Format: comma-separated ``source_claim->metadata_key`` pairs, e.g.
+ ``employment_type->acme_employment_type,org_info.department->department``.
+ Destination keys may use an optional ``metadata.`` prefix; values are stored
+ on the LiteLLM user's ``metadata`` JSON column.
+ """
+ claim_map_raw = CLI_SSO_CLAIM_MAP.strip()
+ if not claim_map_raw:
+ return []
+
+ parsed: List[Tuple[str, str]] = []
+ for entry in claim_map_raw.split(","):
+ entry = entry.strip()
+ if not entry or "->" not in entry:
+ continue
+ source_claim, dest_key = entry.split("->", 1)
+ source_claim = source_claim.strip()
+ dest_key = dest_key.strip()
+ if dest_key.startswith("metadata."):
+ dest_key = dest_key[len("metadata.") :]
+ if source_claim and dest_key:
+ parsed.append((source_claim, dest_key))
+ return parsed
+
+
+def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool:
+ if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key):
+ return False
+ lowered = dest_key.lower()
+ return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS)
+
+
+def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
+ if not isinstance(value, _CLI_SSO_SCALAR_TYPES):
+ return False
+ if isinstance(value, str):
+ if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH:
+ return False
+ if value.startswith("eyJ") and value.count(".") >= 2:
+ return False
+ return True
+
+
+def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]:
+ if isinstance(result, dict):
+ return result
+ if hasattr(result, "model_dump"):
+ dumped = result.model_dump()
+ if isinstance(dumped, dict):
+ return cast(Dict[str, Any], dumped)
+ return {}
+
+
+def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any:
+ """Resolve a dot-notation claim path against an SSO result dict.
+
+ Unlike ``get_nested_value``, this does not strip a leading ``metadata.``
+ prefix, since OIDC claims may legitimately use ``metadata`` as a top-level
+ key.
+ """
+ if not claim_path:
+ return None
+ if claim_path in data:
+ return data[claim_path]
+ placeholder = "\x00"
+ parts = claim_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = data
+ for part in parts:
+ if isinstance(current, dict) and part in current:
+ current = current[part]
+ else:
+ return None
+ return current
+
+
+def _extract_sso_claim_value(
+ result: Union[CustomOpenID, OpenID, dict], claim_path: str
+) -> Any:
+ extra_fields = getattr(result, "extra_fields", None)
+ if isinstance(extra_fields, dict):
+ if claim_path in extra_fields:
+ return extra_fields[claim_path]
+ nested = _get_nested_claim_value(extra_fields, claim_path)
+ if nested is not None:
+ return nested
+
+ if isinstance(result, dict):
+ return _get_nested_claim_value(result, claim_path)
+
+ result_dict = _sso_result_to_dict(result)
+ return _get_nested_claim_value(result_dict, claim_path)
+
+
+def _set_nested_metadata_value(
+ metadata: Dict[str, Any], key_path: str, value: Any
+) -> None:
+ placeholder = "\x00"
+ parts = key_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = metadata
+ for part in parts[:-1]:
+ existing = current.get(part)
+ if not isinstance(existing, dict):
+ existing = {}
+ current[part] = existing
+ current = existing
+ current[parts[-1]] = value
+
+
+def _flatten_cli_sso_metadata_for_poll(
+ metadata: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ """Expose scalar attribution metadata as a flat dict for CLI poll responses."""
+ flattened: Dict[str, Union[str, int, float, bool]] = {}
+ stack: List[Tuple[str, Any]] = [("", metadata)]
+ while stack:
+ prefix, value = stack.pop()
+ if isinstance(value, dict):
+ for key, nested in value.items():
+ nested_prefix = f"{prefix}.{key}" if prefix else key
+ stack.append((nested_prefix, nested))
+ elif _is_safe_cli_sso_scalar_claim_value(value):
+ flattened[prefix] = value
+ return flattened
+
+
+def build_cli_sso_attribution_metadata(
+ result: Union[CustomOpenID, OpenID, dict],
+) -> Dict[str, Any]:
+ """
+ Build allowlisted, non-secret scalar attribution metadata from an SSO result.
+
+ Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and
+ may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID.
+ """
+ claim_map = _parse_cli_sso_claim_map()
+ if not claim_map:
+ return {}
+
+ metadata: Dict[str, Any] = {}
+ for source_claim, dest_key in claim_map:
+ if not _is_safe_cli_sso_metadata_dest_key(dest_key):
+ verbose_proxy_logger.debug(
+ f"Skipping unsafe CLI SSO metadata destination key: {dest_key}"
+ )
+ continue
+
+ raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim)
+ if not _is_safe_cli_sso_scalar_claim_value(raw_value):
+ continue
+
+ _set_nested_metadata_value(
+ metadata=metadata, key_path=dest_key, value=raw_value
+ )
+
+ return metadata
+
+
+def _merge_cli_sso_attribution_metadata(
+ existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any]
+) -> Dict[str, Any]:
+ """Merge attribution metadata into existing user metadata in-place.
+
+ Preserves original value types (in particular, string claim values that
+ happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts
+ are merged iteratively so attribution claims do not clobber unrelated keys
+ under the same parent.
+ """
+ pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [
+ (existing_metadata, attribution_metadata)
+ ]
+ while pending:
+ target, source = pending.pop()
+ for key, value in source.items():
+ if value is None:
+ continue
+ existing_value = target.get(key)
+ if isinstance(value, dict) and isinstance(existing_value, dict):
+ pending.append((existing_value, value))
+ else:
+ target[key] = value
+ return existing_metadata
+
+
+async def _persist_cli_sso_user_metadata(
+ prisma_client: PrismaClient,
+ user_id: str,
+ attribution_metadata: Dict[str, Any],
+) -> None:
+ if not attribution_metadata:
+ return
+
+ try:
+ user_row = await prisma_client.db.litellm_usertable.find_unique(
+ where={"user_id": user_id}
+ )
+ existing_metadata: Dict[str, Any] = {}
+ if user_row is not None:
+ row_metadata = user_row.metadata
+ if isinstance(row_metadata, dict):
+ existing_metadata = deepcopy(row_metadata)
+
+ merged_metadata = _merge_cli_sso_attribution_metadata(
+ existing_metadata=existing_metadata,
+ attribution_metadata=attribution_metadata,
+ )
+ await prisma_client.db.litellm_usertable.update_many(
+ where={"user_id": user_id},
+ data={"metadata": merged_metadata},
+ )
+ verbose_proxy_logger.info(
+ f"Persisted CLI SSO attribution metadata for user {user_id}: "
+ f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}"
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}"
+ )
+
+
+def _cli_poll_attribution_metadata_from_session(
+ session_data: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ stored = session_data.get("attribution_metadata")
+ if isinstance(stored, dict):
+ return _flatten_cli_sso_metadata_for_poll(stored)
+ return {}
+
+
def _render_cli_sso_verification_page(
verify_url: str, browser_complete_token: str
) -> str:
@@ -1674,7 +1924,12 @@
key_id = state_parts[1] if len(state_parts) > 1 else None
verbose_proxy_logger.info("CLI SSO callback detected")
- return await cli_sso_callback(request=request, key=key_id, result=result)
+ return await cli_sso_callback(
+ request=request,
+ key=key_id,
+ result=result,
+ received_response=received_response,
+ )
# Control-plane cross-origin: read return_to from cookie.
# Starlette's cookie_parser already handles RFC 2109 unquoting.
@@ -1692,15 +1947,144 @@
)
+async def _build_cli_sso_user_defined_values(
+ result: Union[OpenID, dict],
+ parsed_openid_result: ParsedOpenIDResult,
+) -> Optional[SSOUserDefinedValues]:
+ from litellm.proxy.proxy_server import user_custom_sso
+
+ user_id = parsed_openid_result.get("user_id")
+ if user_custom_sso is not None:
+ if inspect.iscoroutinefunction(user_custom_sso):
+ return await user_custom_sso(result) # type: ignore
+ raise ValueError("user_custom_sso must be a coroutine function")
+ if user_id is None:
+ return None
+ return SSOUserDefinedValues(
+ models=[],
+ user_id=user_id,
+ user_email=parsed_openid_result.get("user_email"),
+ max_budget=litellm.max_internal_user_budget,
+ user_role=parsed_openid_result.get("user_role"),
+ budget_duration=litellm.internal_user_budget_duration,
+ )
+
+
+async def _fetch_cli_sso_team_details(
+ prisma_client: PrismaClient,
+ teams: List[str],
+) -> List[Dict[str, Any]]:
+ team_details: List[Dict[str, Any]] = []
+ try:
+ if teams:
+ prisma_teams = await prisma_client.db.litellm_teamtable.find_many(
+ where={"team_id": {"in": teams}}
+ )
+ for team_row in prisma_teams:
+ team_dict = team_row.model_dump()
+ team_details.append(
+ {
+ "team_id": team_dict.get("team_id"),
+ "team_alias": team_dict.get("team_alias"),
+ }
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Error fetching team details for CLI SSO session: {e}"
+ )
+ return team_details
+
+
+async def _complete_cli_sso_callback_session(
+ *,
+ request: Request,
+ key: str,
+ flow: dict,
+ result: Union[OpenID, dict],
+ parsed_openid_result: ParsedOpenIDResult,
+ user_defined_values: Optional[SSOUserDefinedValues],
+ prisma_client: PrismaClient,
+ user_api_key_cache: UserApiKeyCache,
+ proxy_logging_obj: ProxyLogging,
+):
+ from fastapi.responses import HTMLResponse
+
+ user_id = parsed_openid_result.get("user_id")
+ user_email = parsed_openid_result.get("user_email")
+ user_info = await get_user_info_from_db(
+ result=result,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ user_email=user_email,
+ user_defined_values=user_defined_values,
+ alternate_user_id=user_id,
+ )
+ if user_info is None:
+ raise HTTPException(
+ status_code=500, detail="Failed to retrieve user information from SSO"
+ )
+ if not user_info.user_id:
+ raise HTTPException(
+ status_code=500, detail="Failed to retrieve user information from SSO"
+ )
+
+ teams: List[str] = []
+ if hasattr(user_info, "teams") and user_info.teams:
+ teams = user_info.teams if isinstance(user_info.teams, list) else []
+
+ team_details = await _fetch_cli_sso_team_details(
+ prisma_client=prisma_client, teams=teams
+ )
+ attribution_metadata = build_cli_sso_attribution_metadata(result=result)
+ if attribution_metadata:
+ await _persist_cli_sso_user_metadata(
+ prisma_client=prisma_client,
+ user_id=cast(str, user_info.user_id),
+ attribution_metadata=attribution_metadata,
+ )
+
+ flow["session_data"] = {
+ "user_id": cast(str, user_info.user_id),
+ "user_role": user_info.user_role,
+ "models": user_info.models if hasattr(user_info, "models") else [],
+ "user_email": user_email,
+ "teams": teams,
+ "team_details": team_details,
+ "attribution_metadata": attribution_metadata,
+ }
+ flow["sso_complete"] = True
+ browser_complete_token = secrets.token_urlsafe(32)
+ flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token)
+ _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow)
+
+ verbose_proxy_logger.info(
+ f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
+ )
+ verify_url = get_custom_url(
+ request_base_url=str(request.base_url),
+ route=f"sso/cli/complete/{key}",
+ )
+ return HTMLResponse(
+ content=_render_cli_sso_verification_page(
+ verify_url=verify_url,
+ browser_complete_token=browser_complete_token,
+ ),
+ status_code=200,
+ )
+
+
async def cli_sso_callback(
request: Request,
key: Optional[str] = None,
result: Optional[Union[OpenID, dict]] = None,
+ received_response: Optional[dict] = None,
):
"""CLI SSO callback - stores session info for JWT generation on polling"""
verbose_proxy_logger.info("CLI SSO callback")
from litellm.proxy.proxy_server import (
+ general_settings,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
@@ -1722,92 +2106,40 @@
# After None check, cast to non-None type for type checker
result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result)
- parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result(
- result=result_non_none
- )
- verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
-
try:
- # Get full user info from DB
- user_info = await get_user_info_from_db(
+ parsed_openid_result = (
+ SSOAuthenticationHandler._get_user_email_and_id_from_result(
+ result=result_non_none,
+ generic_client_id=os.getenv("GENERIC_CLIENT_ID", None),
+ )
+ )
+ verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
+ user_defined_values = await _build_cli_sso_user_defined_values(
result=result_non_none,
- prisma_client=prisma_client,
- user_api_key_cache=user_api_key_cache,
- proxy_logging_obj=proxy_logging_obj,
- user_email=parsed_openid_result.get("user_email"),
- user_defined_values=None,
- alternate_user_id=parsed_openid_result.get("user_id"),
+ parsed_openid_result=parsed_openid_result,
)
- if user_info is None:
- raise HTTPException(
- status_code=500, detail="Failed to retrieve user information from SSO"
- )
-
- # Get all teams from user_info - CLI will let user select which one
- teams: List[str] = []
- if hasattr(user_info, "teams") and user_info.teams:
- teams = user_info.teams if isinstance(user_info.teams, list) else []
-
- # Also fetch team aliases for a better CLI UX. We keep the original
- # "teams" list of IDs for backwards compatibility and add an
- # optional "team_details" field containing objects with both
- # team_id and team_alias.
- team_details: List[Dict[str, Any]] = []
- try:
- if teams:
- prisma_teams = await prisma_client.db.litellm_teamtable.find_many(
- where={"team_id": {"in": teams}}
- )
- for team_row in prisma_teams:
- team_dict = team_row.model_dump()
- team_details.append(
- {
- "team_id": team_dict.get("team_id"),
- "team_alias": team_dict.get("team_alias"),
- }
- )
- except Exception as e:
- # If anything goes wrong here, fall back gracefully without
- # impacting the SSO flow.
- verbose_proxy_logger.error(
- f"Error fetching team details for CLI SSO session: {e}"
- )
-
- session_data = {
- "user_id": user_info.user_id,
- "user_role": user_info.user_role,
- "models": user_info.models if hasattr(user_info, "models") else [],
- "user_email": parsed_openid_result.get("user_email"),
- "teams": teams,
- # Optional rich metadata for clients that want nicer display
- "team_details": team_details,
- }
-
- flow["session_data"] = session_data
- flow["sso_complete"] = True
- browser_complete_token = secrets.token_urlsafe(32)
- flow["browser_complete_token_hash"] = _hash_cli_sso_secret(
- browser_complete_token
+ SSOAuthenticationHandler.verify_user_in_restricted_sso_group(
+ general_settings=general_settings,
+ result=result_non_none,
+ received_response=received_response,
)
- _set_cli_sso_flow(login_id=cast(str, key), cache=user_api_key_cache, flow=flow)
- verbose_proxy_logger.info(
- f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
+ return await _complete_cli_sso_callback_session(
+ request=request,
+ key=cast(str, key),
+ flow=flow,
+ result=result_non_none,
+ parsed_openid_result=parsed_openid_result,
+ user_defined_values=user_defined_values,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
)
-
- from fastapi.responses import HTMLResponse
-
- verify_url = get_custom_url(
- request_base_url=str(request.base_url),
- route=f"sso/cli/complete/{key}",
- )
- html_content = _render_cli_sso_verification_page(
- verify_url=verify_url,
- browser_complete_token=browser_complete_token,
- )
- return HTMLResponse(content=html_content, status_code=200)
-
+ except ProxyException:
+ raise
+ except HTTPException:
+ raise
except Exception as e:
verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}")
raise HTTPException(
@@ -1873,13 +2205,19 @@
team_details_response = [
{"team_id": t, "team_alias": None} for t in user_teams
]
- return {
+ poll_response: Dict[str, Any] = {
"status": "ready",
"user_id": user_id,
"teams": user_teams,
"team_details": team_details_response,
"requires_team_selection": True,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
# Validate team_id if provided
if team_id is not None:
@@ -1912,7 +2250,7 @@
verbose_proxy_logger.info(
f"CLI JWT generated for user: {user_id}, team: {team_id}"
)
- return {
+ poll_response = {
"status": "ready",
"key": jwt_token,
"user_id": user_id,
@@ -1922,6 +2260,12 @@
# present nicer information if needed.
"team_details": user_team_details,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
else:
return {"status": "pending"}
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -5552,6 +5552,289 @@
assert result.extra_fields["another_missing"] is None
+class TestCliSsoAttributionMetadata:
+ """CLI SSO allowlisted OIDC claim persistence and poll exposure."""
+
+ def test_parse_cli_sso_claim_map(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->metadata.acme_employment_type, org_info.department -> department",
+ )
+ assert ui_sso._parse_cli_sso_claim_map() == [
+ ("employment_type", "acme_employment_type"),
+ ("org_info.department", "department"),
+ ]
+
+ def test_build_cli_sso_attribution_metadata_filters_non_scalars(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->acme_employment_type,access_token->should_drop,group->groups",
+ )
+
+ result = CustomOpenID(
+ id="user-1",
+ email="[email protected]",
+ display_name="User",
+ provider="generic",
+ team_ids=[],
+ extra_fields={
+ "employment_type": "full_time",
+ "access_token": "eyJhbGciOiJIUzI1NiJ9.payload.signature",
+ "group": ["team-a", "team-b"],
+ },
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(result=result)
+ assert metadata == {"acme_employment_type": "full_time"}
+
+ def test_build_cli_sso_attribution_metadata_from_oidc_dict(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "org_info.department->department",
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(
+ result={
+ "sub": "user-1",
+ "email": "[email protected]",
+ "org_info": {"department": "Engineering"},
+ }
+ )
+ assert metadata == {"department": "Engineering"}
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self):
+ """First CLI SSO login must supply SSOUserDefinedValues so upsert can create the user."""
+ from litellm.proxy._types import LiteLLM_UserTable
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ session_key = "cli-session-new-user"
+ mock_user_info = LiteLLM_UserTable(
+ user_id="cli-test-user",
+ user_role="internal_user",
+ teams=[],
+ models=[],
+ )
+ mock_sso_result = CustomOpenID(
+ id="cli-test-user",
+ email="[email protected]",
+ display_name="cli-test-user",
+ provider="generic",
+ team_ids=[],
+ )
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ get_user_info_mock = AsyncMock(return_value=mock_user_info)
+
+ with (
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ get_user_info_mock,
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ ):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key=session_key,
+ result=mock_sso_result,
+ )
+
+ get_user_info_mock.assert_awaited_once()
+ assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None
+ assert (
+ get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"]
+ == "cli-test-user"
+ )
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_rejects_restricted_sso_group(self):
+ """CLI SSO must enforce restricted_sso_group before upserting the user."""
+ from litellm.proxy._types import ProxyException
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ mock_sso_result = CustomOpenID(
+ id="cli-test-user",
+ email="[email protected]",
+ display_name="cli-test-user",
+ provider="generic",
+ team_ids=["other-group"],
+ )
+
+ with (
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ new=AsyncMock(),
+ ) as get_user_info_mock,
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ patch(
+ "litellm.proxy.proxy_server.general_settings",
+ {
+ "ui_access_mode": {
+ "type": "restricted_sso_group",
+ "restricted_sso_group": "required-group",
+ }
+ },
+ ),
+ ):
+ with pytest.raises(ProxyException):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key="cli-session-restricted",
+ result=mock_sso_result,
+ received_response={"groups": ["other-group"]},
+ )
+
... diff truncated: showing 800 of 919 linesYou can send follow-ups to the cloud agent here.
…tion
When GENERIC_USER_EXTRA_ATTRIBUTES captures a parent object (e.g. org_info),
extra_fields stores it as {"org_info": {"department": "..."}}. A CLI claim
map entry using a dotted path like org_info.department would silently fail
because the lookup only checked the exact flat key. Fall back to dotted-path
resolution on extra_fields before model_dump().
Co-authored-by: Yassin Kortam <[email protected]>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high mode and found 2 potential issues.
Bugbot Autofix prepared fixes for both issues found in the latest run.
- ✅ Fixed: Auth callback breaks existing test with new parameter
- Updated test_auth_callback_routes_to_cli to include received_response=None in assert_called_once_with so it matches the new cli_sso_callback signature.
- ✅ Fixed: Secret fragment "token" makes specific entries redundant
- Removed the overly broad 'token' substring from _CLI_SSO_SECRET_KEY_FRAGMENTS so legitimate keys like 'token_type' are no longer silently blocked, while still rejecting access_token/id_token/refresh_token.
Preview (26101a0b77)
diff --git a/litellm/constants.py b/litellm/constants.py
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1443,6 +1443,12 @@
or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS")
or 24
)
+# Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g.
+# "employment_type->acme_employment_type,org_info.department->department"
+CLI_SSO_CLAIM_MAP = (
+ os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or ""
+)
+CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024
########################### UI SESSION DURATION ###########################
# Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d"
diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md
--- a/litellm/proxy/client/README.md
+++ b/litellm/proxy/client/README.md
@@ -350,7 +350,7 @@
4. **User Authentication**: User completes SSO authentication in browser
5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
6. **User Code Verification**: Browser confirms the verification code shown in the CLI
-7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
+7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready. When `CLI_SSO_CLAIM_MAP` is configured on the proxy, the poll response may include `attribution_metadata` (allowlisted scalar OIDC claims for client attribution).
8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
### Benefits of This Approach
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -43,6 +43,8 @@
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
+ CLI_SSO_CLAIM_MAP,
+ CLI_SSO_CLAIM_MAX_SCALAR_LENGTH,
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
LITELLM_CLI_SOURCE_IDENTIFIER,
@@ -140,6 +142,20 @@
_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
+_CLI_SSO_SCALAR_TYPES = (str, int, float, bool)
+_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
+_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset(
+ {
+ "access_token",
+ "api_key",
+ "client_secret",
+ "id_token",
+ "password",
+ "private_key",
+ "refresh_token",
+ "secret",
+ }
+)
def _hash_cli_sso_secret(secret: str) -> str:
@@ -225,6 +241,239 @@
return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)
+def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]:
+ """
+ Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP.
+
+ Format: comma-separated ``source_claim->metadata_key`` pairs, e.g.
+ ``employment_type->acme_employment_type,org_info.department->department``.
+ Destination keys may use an optional ``metadata.`` prefix; values are stored
+ on the LiteLLM user's ``metadata`` JSON column.
+ """
+ claim_map_raw = CLI_SSO_CLAIM_MAP.strip()
+ if not claim_map_raw:
+ return []
+
+ parsed: List[Tuple[str, str]] = []
+ for entry in claim_map_raw.split(","):
+ entry = entry.strip()
+ if not entry or "->" not in entry:
+ continue
+ source_claim, dest_key = entry.split("->", 1)
+ source_claim = source_claim.strip()
+ dest_key = dest_key.strip()
+ if dest_key.startswith("metadata."):
+ dest_key = dest_key[len("metadata.") :]
+ if source_claim and dest_key:
+ parsed.append((source_claim, dest_key))
+ return parsed
+
+
+def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool:
+ if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key):
+ return False
+ lowered = dest_key.lower()
+ return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS)
+
+
+def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
+ if not isinstance(value, _CLI_SSO_SCALAR_TYPES):
+ return False
+ if isinstance(value, str):
+ if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH:
+ return False
+ if value.startswith("eyJ") and value.count(".") >= 2:
+ return False
+ return True
+
+
+def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]:
+ if isinstance(result, dict):
+ return result
+ if hasattr(result, "model_dump"):
+ dumped = result.model_dump()
+ if isinstance(dumped, dict):
+ return cast(Dict[str, Any], dumped)
+ return {}
+
+
+def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any:
+ """Resolve a dot-notation claim path against an SSO result dict.
+
+ Unlike ``get_nested_value``, this does not strip a leading ``metadata.``
+ prefix, since OIDC claims may legitimately use ``metadata`` as a top-level
+ key.
+ """
+ if not claim_path:
+ return None
+ if claim_path in data:
+ return data[claim_path]
+ placeholder = "\x00"
+ parts = claim_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = data
+ for part in parts:
+ if isinstance(current, dict) and part in current:
+ current = current[part]
+ else:
+ return None
+ return current
+
+
+def _extract_sso_claim_value(
+ result: Union[CustomOpenID, OpenID, dict], claim_path: str
+) -> Any:
+ extra_fields = getattr(result, "extra_fields", None)
+ if isinstance(extra_fields, dict):
+ if claim_path in extra_fields:
+ return extra_fields[claim_path]
+ nested = _get_nested_claim_value(extra_fields, claim_path)
+ if nested is not None:
+ return nested
+
+ if isinstance(result, dict):
+ return _get_nested_claim_value(result, claim_path)
+
+ result_dict = _sso_result_to_dict(result)
+ return _get_nested_claim_value(result_dict, claim_path)
+
+
+def _set_nested_metadata_value(
+ metadata: Dict[str, Any], key_path: str, value: Any
+) -> None:
+ placeholder = "\x00"
+ parts = key_path.replace("\\.", placeholder).split(".")
+ parts = [p.replace(placeholder, ".") for p in parts]
+ current: Any = metadata
+ for part in parts[:-1]:
+ existing = current.get(part)
+ if not isinstance(existing, dict):
+ existing = {}
+ current[part] = existing
+ current = existing
+ current[parts[-1]] = value
+
+
+def _flatten_cli_sso_metadata_for_poll(
+ metadata: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ """Expose scalar attribution metadata as a flat dict for CLI poll responses."""
+ flattened: Dict[str, Union[str, int, float, bool]] = {}
+ stack: List[Tuple[str, Any]] = [("", metadata)]
+ while stack:
+ prefix, value = stack.pop()
+ if isinstance(value, dict):
+ for key, nested in value.items():
+ nested_prefix = f"{prefix}.{key}" if prefix else key
+ stack.append((nested_prefix, nested))
+ elif _is_safe_cli_sso_scalar_claim_value(value):
+ flattened[prefix] = value
+ return flattened
+
+
+def build_cli_sso_attribution_metadata(
+ result: Union[CustomOpenID, OpenID, dict],
+) -> Dict[str, Any]:
+ """
+ Build allowlisted, non-secret scalar attribution metadata from an SSO result.
+
+ Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and
+ may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID.
+ """
+ claim_map = _parse_cli_sso_claim_map()
+ if not claim_map:
+ return {}
+
+ metadata: Dict[str, Any] = {}
+ for source_claim, dest_key in claim_map:
+ if not _is_safe_cli_sso_metadata_dest_key(dest_key):
+ verbose_proxy_logger.debug(
+ f"Skipping unsafe CLI SSO metadata destination key: {dest_key}"
+ )
+ continue
+
+ raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim)
+ if not _is_safe_cli_sso_scalar_claim_value(raw_value):
+ continue
+
+ _set_nested_metadata_value(
+ metadata=metadata, key_path=dest_key, value=raw_value
+ )
+
+ return metadata
+
+
+def _merge_cli_sso_attribution_metadata(
+ existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any]
+) -> Dict[str, Any]:
+ """Merge attribution metadata into existing user metadata in-place.
+
+ Preserves original value types (in particular, string claim values that
+ happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts
+ are merged iteratively so attribution claims do not clobber unrelated keys
+ under the same parent.
+ """
+ pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [
+ (existing_metadata, attribution_metadata)
+ ]
+ while pending:
+ target, source = pending.pop()
+ for key, value in source.items():
+ if value is None:
+ continue
+ existing_value = target.get(key)
+ if isinstance(value, dict) and isinstance(existing_value, dict):
+ pending.append((existing_value, value))
+ else:
+ target[key] = value
+ return existing_metadata
+
+
+async def _persist_cli_sso_user_metadata(
+ prisma_client: PrismaClient,
+ user_id: str,
+ attribution_metadata: Dict[str, Any],
+) -> None:
+ if not attribution_metadata:
+ return
+
+ try:
+ user_row = await prisma_client.db.litellm_usertable.find_unique(
+ where={"user_id": user_id}
+ )
+ existing_metadata: Dict[str, Any] = {}
+ if user_row is not None:
+ row_metadata = user_row.metadata
+ if isinstance(row_metadata, dict):
+ existing_metadata = deepcopy(row_metadata)
+
+ merged_metadata = _merge_cli_sso_attribution_metadata(
+ existing_metadata=existing_metadata,
+ attribution_metadata=attribution_metadata,
+ )
+ await prisma_client.db.litellm_usertable.update_many(
+ where={"user_id": user_id},
+ data={"metadata": merged_metadata},
+ )
+ verbose_proxy_logger.info(
+ f"Persisted CLI SSO attribution metadata for user {user_id}: "
+ f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}"
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}"
+ )
+
+
+def _cli_poll_attribution_metadata_from_session(
+ session_data: Dict[str, Any],
+) -> Dict[str, Union[str, int, float, bool]]:
+ stored = session_data.get("attribution_metadata")
+ if isinstance(stored, dict):
+ return _flatten_cli_sso_metadata_for_poll(stored)
+ return {}
+
+
def _render_cli_sso_verification_page(
verify_url: str, browser_complete_token: str
) -> str:
@@ -1674,7 +1923,12 @@
key_id = state_parts[1] if len(state_parts) > 1 else None
verbose_proxy_logger.info("CLI SSO callback detected")
- return await cli_sso_callback(request=request, key=key_id, result=result)
+ return await cli_sso_callback(
+ request=request,
+ key=key_id,
+ result=result,
+ received_response=received_response,
+ )
# Control-plane cross-origin: read return_to from cookie.
# Starlette's cookie_parser already handles RFC 2109 unquoting.
@@ -1692,15 +1946,144 @@
)
+async def _build_cli_sso_user_defined_values(
+ result: Union[OpenID, dict],
+ parsed_openid_result: ParsedOpenIDResult,
+) -> Optional[SSOUserDefinedValues]:
+ from litellm.proxy.proxy_server import user_custom_sso
+
+ user_id = parsed_openid_result.get("user_id")
+ if user_custom_sso is not None:
+ if inspect.iscoroutinefunction(user_custom_sso):
+ return await user_custom_sso(result) # type: ignore
+ raise ValueError("user_custom_sso must be a coroutine function")
+ if user_id is None:
+ return None
+ return SSOUserDefinedValues(
+ models=[],
+ user_id=user_id,
+ user_email=parsed_openid_result.get("user_email"),
+ max_budget=litellm.max_internal_user_budget,
+ user_role=parsed_openid_result.get("user_role"),
+ budget_duration=litellm.internal_user_budget_duration,
+ )
+
+
+async def _fetch_cli_sso_team_details(
+ prisma_client: PrismaClient,
+ teams: List[str],
+) -> List[Dict[str, Any]]:
+ team_details: List[Dict[str, Any]] = []
+ try:
+ if teams:
+ prisma_teams = await prisma_client.db.litellm_teamtable.find_many(
+ where={"team_id": {"in": teams}}
+ )
+ for team_row in prisma_teams:
+ team_dict = team_row.model_dump()
+ team_details.append(
+ {
+ "team_id": team_dict.get("team_id"),
+ "team_alias": team_dict.get("team_alias"),
+ }
+ )
+ except Exception as e:
+ verbose_proxy_logger.error(
+ f"Error fetching team details for CLI SSO session: {e}"
+ )
+ return team_details
+
+
+async def _complete_cli_sso_callback_session(
+ *,
+ request: Request,
+ key: str,
+ flow: dict,
+ result: Union[OpenID, dict],
+ parsed_openid_result: ParsedOpenIDResult,
+ user_defined_values: Optional[SSOUserDefinedValues],
+ prisma_client: PrismaClient,
+ user_api_key_cache: UserApiKeyCache,
+ proxy_logging_obj: ProxyLogging,
+):
+ from fastapi.responses import HTMLResponse
+
+ user_id = parsed_openid_result.get("user_id")
+ user_email = parsed_openid_result.get("user_email")
+ user_info = await get_user_info_from_db(
+ result=result,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ user_email=user_email,
+ user_defined_values=user_defined_values,
+ alternate_user_id=user_id,
+ )
+ if user_info is None:
+ raise HTTPException(
+ status_code=500, detail="Failed to retrieve user information from SSO"
+ )
+ if not user_info.user_id:
+ raise HTTPException(
+ status_code=500, detail="Failed to retrieve user information from SSO"
+ )
+
+ teams: List[str] = []
+ if hasattr(user_info, "teams") and user_info.teams:
+ teams = user_info.teams if isinstance(user_info.teams, list) else []
+
+ team_details = await _fetch_cli_sso_team_details(
+ prisma_client=prisma_client, teams=teams
+ )
+ attribution_metadata = build_cli_sso_attribution_metadata(result=result)
+ if attribution_metadata:
+ await _persist_cli_sso_user_metadata(
+ prisma_client=prisma_client,
+ user_id=cast(str, user_info.user_id),
+ attribution_metadata=attribution_metadata,
+ )
+
+ flow["session_data"] = {
+ "user_id": cast(str, user_info.user_id),
+ "user_role": user_info.user_role,
+ "models": user_info.models if hasattr(user_info, "models") else [],
+ "user_email": user_email,
+ "teams": teams,
+ "team_details": team_details,
+ "attribution_metadata": attribution_metadata,
+ }
+ flow["sso_complete"] = True
+ browser_complete_token = secrets.token_urlsafe(32)
+ flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token)
+ _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow)
+
+ verbose_proxy_logger.info(
+ f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
+ )
+ verify_url = get_custom_url(
+ request_base_url=str(request.base_url),
+ route=f"sso/cli/complete/{key}",
+ )
+ return HTMLResponse(
+ content=_render_cli_sso_verification_page(
+ verify_url=verify_url,
+ browser_complete_token=browser_complete_token,
+ ),
+ status_code=200,
+ )
+
+
async def cli_sso_callback(
request: Request,
key: Optional[str] = None,
result: Optional[Union[OpenID, dict]] = None,
+ received_response: Optional[dict] = None,
):
"""CLI SSO callback - stores session info for JWT generation on polling"""
verbose_proxy_logger.info("CLI SSO callback")
from litellm.proxy.proxy_server import (
+ general_settings,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
@@ -1722,92 +2105,40 @@
# After None check, cast to non-None type for type checker
result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result)
- parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result(
- result=result_non_none
- )
- verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
-
try:
- # Get full user info from DB
- user_info = await get_user_info_from_db(
+ parsed_openid_result = (
+ SSOAuthenticationHandler._get_user_email_and_id_from_result(
+ result=result_non_none,
+ generic_client_id=os.getenv("GENERIC_CLIENT_ID", None),
+ )
+ )
+ verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
+ user_defined_values = await _build_cli_sso_user_defined_values(
result=result_non_none,
- prisma_client=prisma_client,
- user_api_key_cache=user_api_key_cache,
- proxy_logging_obj=proxy_logging_obj,
- user_email=parsed_openid_result.get("user_email"),
- user_defined_values=None,
- alternate_user_id=parsed_openid_result.get("user_id"),
+ parsed_openid_result=parsed_openid_result,
)
- if user_info is None:
- raise HTTPException(
- status_code=500, detail="Failed to retrieve user information from SSO"
- )
-
- # Get all teams from user_info - CLI will let user select which one
- teams: List[str] = []
- if hasattr(user_info, "teams") and user_info.teams:
- teams = user_info.teams if isinstance(user_info.teams, list) else []
-
- # Also fetch team aliases for a better CLI UX. We keep the original
- # "teams" list of IDs for backwards compatibility and add an
- # optional "team_details" field containing objects with both
- # team_id and team_alias.
- team_details: List[Dict[str, Any]] = []
- try:
- if teams:
- prisma_teams = await prisma_client.db.litellm_teamtable.find_many(
- where={"team_id": {"in": teams}}
- )
- for team_row in prisma_teams:
- team_dict = team_row.model_dump()
- team_details.append(
- {
- "team_id": team_dict.get("team_id"),
- "team_alias": team_dict.get("team_alias"),
- }
- )
- except Exception as e:
- # If anything goes wrong here, fall back gracefully without
- # impacting the SSO flow.
- verbose_proxy_logger.error(
- f"Error fetching team details for CLI SSO session: {e}"
- )
-
- session_data = {
- "user_id": user_info.user_id,
- "user_role": user_info.user_role,
- "models": user_info.models if hasattr(user_info, "models") else [],
- "user_email": parsed_openid_result.get("user_email"),
- "teams": teams,
- # Optional rich metadata for clients that want nicer display
- "team_details": team_details,
- }
-
- flow["session_data"] = session_data
- flow["sso_complete"] = True
- browser_complete_token = secrets.token_urlsafe(32)
- flow["browser_complete_token_hash"] = _hash_cli_sso_secret(
- browser_complete_token
+ SSOAuthenticationHandler.verify_user_in_restricted_sso_group(
+ general_settings=general_settings,
+ result=result_non_none,
+ received_response=received_response,
)
- _set_cli_sso_flow(login_id=cast(str, key), cache=user_api_key_cache, flow=flow)
- verbose_proxy_logger.info(
- f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
+ return await _complete_cli_sso_callback_session(
+ request=request,
+ key=cast(str, key),
+ flow=flow,
+ result=result_non_none,
+ parsed_openid_result=parsed_openid_result,
+ user_defined_values=user_defined_values,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
)
-
- from fastapi.responses import HTMLResponse
-
- verify_url = get_custom_url(
- request_base_url=str(request.base_url),
- route=f"sso/cli/complete/{key}",
- )
- html_content = _render_cli_sso_verification_page(
- verify_url=verify_url,
- browser_complete_token=browser_complete_token,
- )
- return HTMLResponse(content=html_content, status_code=200)
-
+ except ProxyException:
+ raise
+ except HTTPException:
+ raise
except Exception as e:
verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}")
raise HTTPException(
@@ -1873,13 +2204,19 @@
team_details_response = [
{"team_id": t, "team_alias": None} for t in user_teams
]
- return {
+ poll_response: Dict[str, Any] = {
"status": "ready",
"user_id": user_id,
"teams": user_teams,
"team_details": team_details_response,
"requires_team_selection": True,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
# Validate team_id if provided
if team_id is not None:
@@ -1912,7 +2249,7 @@
verbose_proxy_logger.info(
f"CLI JWT generated for user: {user_id}, team: {team_id}"
)
- return {
+ poll_response = {
"status": "ready",
"key": jwt_token,
"user_id": user_id,
@@ -1922,6 +2259,12 @@
# present nicer information if needed.
"team_details": user_team_details,
}
+ attribution_metadata = _cli_poll_attribution_metadata_from_session(
+ session_data
+ )
+ if attribution_metadata:
+ poll_response["attribution_metadata"] = attribution_metadata
+ return poll_response
else:
return {"status": "pending"}
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -2438,6 +2438,7 @@
request=mock_request,
key="cli-new-session-key-456",
result=mock_result,
+ received_response=None,
)
def test_get_redirect_url_does_not_include_existing_key_in_url(self):
@@ -5552,6 +5553,289 @@
assert result.extra_fields["another_missing"] is None
+class TestCliSsoAttributionMetadata:
+ """CLI SSO allowlisted OIDC claim persistence and poll exposure."""
+
+ def test_parse_cli_sso_claim_map(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->metadata.acme_employment_type, org_info.department -> department",
+ )
+ assert ui_sso._parse_cli_sso_claim_map() == [
+ ("employment_type", "acme_employment_type"),
+ ("org_info.department", "department"),
+ ]
+
+ def test_build_cli_sso_attribution_metadata_filters_non_scalars(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "employment_type->acme_employment_type,access_token->should_drop,group->groups",
+ )
+
+ result = CustomOpenID(
+ id="user-1",
+ email="[email protected]",
+ display_name="User",
+ provider="generic",
+ team_ids=[],
+ extra_fields={
+ "employment_type": "full_time",
+ "access_token": "eyJhbGciOiJIUzI1NiJ9.payload.signature",
+ "group": ["team-a", "team-b"],
+ },
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(result=result)
+ assert metadata == {"acme_employment_type": "full_time"}
+
+ def test_build_cli_sso_attribution_metadata_from_oidc_dict(self, monkeypatch):
+ from litellm.proxy.management_endpoints import ui_sso
+
+ monkeypatch.setattr(
+ ui_sso,
+ "CLI_SSO_CLAIM_MAP",
+ "org_info.department->department",
+ )
+
+ metadata = ui_sso.build_cli_sso_attribution_metadata(
+ result={
+ "sub": "user-1",
+ "email": "[email protected]",
+ "org_info": {"department": "Engineering"},
+ }
+ )
+ assert metadata == {"department": "Engineering"}
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self):
+ """First CLI SSO login must supply SSOUserDefinedValues so upsert can create the user."""
+ from litellm.proxy._types import LiteLLM_UserTable
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ session_key = "cli-session-new-user"
+ mock_user_info = LiteLLM_UserTable(
+ user_id="cli-test-user",
+ user_role="internal_user",
+ teams=[],
+ models=[],
+ )
+ mock_sso_result = CustomOpenID(
+ id="cli-test-user",
+ email="[email protected]",
+ display_name="cli-test-user",
+ provider="generic",
+ team_ids=[],
+ )
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ get_user_info_mock = AsyncMock(return_value=mock_user_info)
+
+ with (
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ get_user_info_mock,
+ ),
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ ):
+ await ui_sso.cli_sso_callback(
+ request=mock_request,
+ key=session_key,
+ result=mock_sso_result,
+ )
+
+ get_user_info_mock.assert_awaited_once()
+ assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None
+ assert (
+ get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"]
+ == "cli-test-user"
+ )
+
+ @pytest.mark.asyncio
+ async def test_cli_sso_callback_rejects_restricted_sso_group(self):
+ """CLI SSO must enforce restricted_sso_group before upserting the user."""
+ from litellm.proxy._types import ProxyException
+ from litellm.proxy.management_endpoints import ui_sso
+ from litellm.proxy.management_endpoints.types import CustomOpenID
+
+ mock_request = MagicMock(spec=Request)
+ mock_request.base_url = "http://internal-proxy.local/"
+ mock_cache = MagicMock()
+ mock_cache.get_cache.return_value = {
+ "poll_secret_hash": "poll-secret-hash",
+ "user_code_hash": "user-code-hash",
+ "sso_complete": False,
+ "user_code_verified": False,
+ "session_data": None,
+ }
+ mock_sso_result = CustomOpenID(
+ id="cli-test-user",
+ email="[email protected]",
+ display_name="cli-test-user",
+ provider="generic",
+ team_ids=["other-group"],
+ )
+
+ with (
+ patch(
+ "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db",
+ new=AsyncMock(),
+ ) as get_user_info_mock,
+ patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+ patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
+ patch("litellm.proxy.proxy_server.user_custom_sso", None),
+ patch(
+ "litellm.proxy.proxy_server.general_settings",
+ {
+ "ui_access_mode": {
+ "type": "restricted_sso_group",
+ "restricted_sso_group": "required-group",
+ }
+ },
+ ),
+ ):
+ with pytest.raises(ProxyException):
... diff truncated: showing 800 of 926 linesYou can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit a0ecb87. Configure here.
…ove redundant 'token' secret fragment Co-authored-by: Yassin Kortam <[email protected]>
Collaborator
Author
Collaborator
Author
|
Docs is here: BerriAI/litellm-docs#190 |
mateo-berri
merged commit 50a3f10
into
litellm_internal_staging
116 of 119 checks passed
fzowl
pushed a commit
to fzowl/litellm
that referenced
this pull request
Jun 24, 2026
…#28463) * feat(proxy): persist allowlisted OIDC claims in CLI SSO poll Map CLI_SSO_CLAIM_MAP sources into user metadata and return scalar attribution_metadata from /sso/cli/poll. Build SSOUserDefinedValues in cli_sso_callback so first-time CLI logins can upsert users. Add mock OIDC scripts and tests for claim extraction and poll exposure. Co-authored-by: Cursor <[email protected]> * docs(proxy): document CLI SSO attribution_metadata in client README Co-authored-by: Cursor <[email protected]> * Delete scripts/mock_oidc_server_for_cli_sso.py * Delete scripts/test_cli_sso_claims_e2e.py * fix(ui_sso): preserve claim types and avoid metadata. prefix stripping - Replace _update_dictionary with a local recursive merge so string OIDC claim values that happen to look numeric are not silently coerced to int/float when persisting CLI SSO attribution metadata. - Use a local dot-path resolver in _extract_sso_claim_value so that source claim paths beginning with 'metadata.' are not silently stripped by get_nested_value (which is designed for LiteLLM JWT metadata, not arbitrary OIDC claims). Co-authored-by: Yassin Kortam <[email protected]> * Remove redundant metadata. prefix strip in _set_nested_metadata_value The _parse_cli_sso_claim_map already strips the metadata. prefix from dest keys before reaching the setter. The duplicate strip in _set_nested_metadata_value was a no-op in normal flow but could mis-place values for dest keys like metadata.metadata.foo. Co-authored-by: Yassin Kortam <[email protected]> * Fix greptile * Fix ruff * Move CLI SSO user defined values build inside try/except for consistent error handling Co-authored-by: Yassin Kortam <[email protected]> * fix(proxy): enforce restricted SSO group on CLI SSO callback Apply verify_user_in_restricted_sso_group before CLI session completion and user upsert, matching the UI SSO path. Re-raise ProxyException so restricted-group denials return 403 instead of 500. Co-authored-by: Cursor <[email protected]> * fix(proxy): replace recursive CLI SSO metadata helpers with iterative merge Use stack-based flatten/merge to satisfy recursive_detector CI. Fix mypy types for UserApiKeyCache and user_id on CLI SSO session completion. Co-authored-by: Cursor <[email protected]> * fix: resolve nested CustomOpenID extra_fields in CLI SSO claim extraction When GENERIC_USER_EXTRA_ATTRIBUTES captures a parent object (e.g. org_info), extra_fields stores it as {"org_info": {"department": "..."}}. A CLI claim map entry using a dotted path like org_info.department would silently fail because the lookup only checked the exact flat key. Fall back to dotted-path resolution on extra_fields before model_dump(). Co-authored-by: Yassin Kortam <[email protected]> * fix(sso): update CLI SSO test for new received_response kwarg and remove redundant 'token' secret fragment Co-authored-by: Yassin Kortam <[email protected]> --------- Co-authored-by: Cursor <[email protected]> Co-authored-by: Yassin Kortam <[email protected]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
CLI_SSO_CLAIM_MAP/LITELLM_CLI_SSO_CLAIM_MAPto map OIDC claims (includingGENERIC_USER_EXTRA_ATTRIBUTES) into usermetadataand expose scalarattribution_metadataon/sso/cli/poll.SSOUserDefinedValueson first login so new users are upserted (fixes 500 when mock/real IdP returns a user not yet in DB).scripts/for testing without a real SSO provider.Test plan
pytest tests/test_litellm/proxy/management_endpoints/test_ui_sso.py::TestCliSsoAttributionMetadata -vpytest tests/test_litellm/proxy/management_endpoints/test_ui_sso.py::TestCLIKeyRegenerationFlow::test_cli_sso_callback_stores_session -vpython scripts/mock_oidc_server_for_cli_sso.py+ proxy with Generic SSO env +python scripts/test_cli_sso_claims_e2e.pyConfig example
Fixes LIT-2631
Note
Medium Risk
Touches the CLI SSO authentication callback/poll flow and now persists extra SSO-derived metadata into the user table, which could impact login behavior or data written to the DB if misconfigured. Safeguards/allowlisting reduce exposure, but it remains auth-adjacent logic.
Overview
CLI SSO polling can now return attribution data and persist it to users. The proxy adds
CLI_SSO_CLAIM_MAP(env-configured) to allowlist specific OIDC claims (including nested paths and Generic SSOextra_fields), store them into the user’smetadata, and expose flattened scalarattribution_metadatain/sso/cli/pollresponses.CLI SSO callback flow is refactored and hardened. The callback now forwards
received_response, enforces restricted SSO group checks before user upsert, and passesSSOUserDefinedValuesso first-time CLI SSO logins can create missing users; tests and docs are updated accordingly.Reviewed by Cursor Bugbot for commit 26101a0. Bugbot is set up for automated code reviews on this repo. Configure here.