From 0abfb07ab8b45a9af91d7d477f2805e9f0d196f3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 21 Nov 2025 16:11:58 -0800 Subject: [PATCH] Remove UI Session Token from user/info return (#16851) --- .../internal_user_endpoints.py | 6 + .../test_internal_user_endpoints.py | 122 ++++++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 66085b69b3d..2ca1c4cc483 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -705,6 +705,7 @@ def _process_keys_for_user_info( keys: Optional[List[LiteLLM_VerificationToken]], all_teams: Optional[Union[List[LiteLLM_TeamTable], List[TeamListResponseObject]]], ): + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash returned_keys = [] @@ -724,6 +725,11 @@ def _process_keys_for_user_info( except Exception: # if using pydantic v1 _key = key.dict() + + # Filter out UI session tokens (team_id="litellm-dashboard") + if _key.get("team_id") == UI_SESSION_TOKEN_TEAM_ID: + continue + if ( "team_id" in _key and _key["team_id"] is not None diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 266056bcdd2..a112ae046ab 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -690,3 +690,125 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): await _check_duplicate_user_email( None, mock_prisma_client ) # Should not raise exception + + +def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): + """ + Test that _process_keys_for_user_info filters out keys with team_id='litellm-dashboard' + + UI session tokens (team_id='litellm-dashboard') should be excluded from user info responses + to prevent confusion, as these are automatically created during dashboard login. + """ + from unittest.mock import MagicMock + + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _process_keys_for_user_info, + ) + + # Create mock keys with different team_ids + mock_key_dashboard = MagicMock() + mock_key_dashboard.model_dump.return_value = { + "token": "sk-dashboard-token", + "team_id": UI_SESSION_TOKEN_TEAM_ID, + "user_id": "test-user", + "key_alias": "dashboard-session-key", + } + + mock_key_regular = MagicMock() + mock_key_regular.model_dump.return_value = { + "token": "sk-regular-token", + "team_id": "regular-team", + "user_id": "test-user", + "key_alias": "regular-key", + } + + mock_key_no_team = MagicMock() + mock_key_no_team.model_dump.return_value = { + "token": "sk-no-team-token", + "team_id": None, + "user_id": "test-user", + "key_alias": "no-team-key", + } + + keys = [mock_key_dashboard, mock_key_regular, mock_key_no_team] + + # Mock general_settings and litellm_master_key_hash (they're imported from proxy_server) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {}, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.litellm_master_key_hash", + "different-hash", + ) + + # Call the function + result = _process_keys_for_user_info(keys=keys, all_teams=None) + + # Verify that dashboard key is filtered out + assert len(result) == 2, "Should return 2 keys (dashboard key filtered out)" + + # Verify dashboard key is not in results + result_team_ids = [key.get("team_id") for key in result] + assert UI_SESSION_TOKEN_TEAM_ID not in result_team_ids, "Dashboard key should be filtered out" + + # Verify regular keys are included + assert "regular-team" in result_team_ids, "Regular team key should be included" + assert None in result_team_ids, "No-team key should be included" + + # Verify the correct keys are returned + result_tokens = [key.get("token") for key in result] + assert "sk-regular-token" in result_tokens, "Regular key should be included" + assert "sk-no-team-token" in result_tokens, "No-team key should be included" + assert "sk-dashboard-token" not in result_tokens, "Dashboard key should not be included" + + +def test_process_keys_for_user_info_handles_none_keys(monkeypatch): + """ + Test that _process_keys_for_user_info handles None keys gracefully + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _process_keys_for_user_info, + ) + + # Mock general_settings and litellm_master_key_hash (they're imported from proxy_server) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {}, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.litellm_master_key_hash", + "different-hash", + ) + + # Call with None keys + result = _process_keys_for_user_info(keys=None, all_teams=None) + + # Should return empty list + assert result == [], "Should return empty list when keys is None" + + +def test_process_keys_for_user_info_handles_empty_keys(monkeypatch): + """ + Test that _process_keys_for_user_info handles empty keys list + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _process_keys_for_user_info, + ) + + # Mock general_settings and litellm_master_key_hash (they're imported from proxy_server) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {}, + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.litellm_master_key_hash", + "different-hash", + ) + + # Call with empty list + result = _process_keys_for_user_info(keys=[], all_teams=None) + + # Should return empty list + assert result == [], "Should return empty list when keys is empty"