github/OpenHands
1
0
Fork 0
mirror of https://github.com/OpenHands/OpenHands.git synced 2026-03-22 05:37:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: ensure LiteLLM user exists before generating API keys (#12667)

Browse Source 
Co-authored-by: openhands <openhands@all-hands.dev>
This commit is contained in:
Saurya Velagapudi
2026-03-18 17:16:43 -07:00
committed by GitHub
parent dcb2e21b87
commit a96760eea7
2 changed files with 331 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
enterprise/storage/lite_llm_manager.py
View File

@@ -164,9 +164,33 @@ class LiteLlmManager:
)
if create_user:
await LiteLlmManager._create_user(
user_created = await LiteLlmManager._create_user(
client, keycloak_user_info.get('email'), keycloak_user_id
)
if not user_created:
logger.error(
'create_entries_failed_user_creation',
extra={
'org_id': org_id,
'user_id': keycloak_user_id,
},
)
return None
# Verify user exists before proceeding with key generation
user_exists = await LiteLlmManager._user_exists(
client, keycloak_user_id
)
if not user_exists:
logger.error(
'create_entries_user_not_found_before_key_generation',
extra={
'org_id': org_id,
'user_id': keycloak_user_id,
'create_user_flag': create_user,
},
)
return None
await LiteLlmManager._add_user_to_team(
client, keycloak_user_id, org_id, team_budget
@@ -655,15 +679,48 @@ class LiteLlmManager:
)
response.raise_for_status()
@staticmethod
async def _user_exists(
client: httpx.AsyncClient,
user_id: str,
) -> bool:
"""Check if a user exists in LiteLLM.
Returns True if the user exists, False otherwise.
"""
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
return False
try:
response = await client.get(
f'{LITE_LLM_API_URL}/user/info?user_id={user_id}',
)
if response.is_success:
user_data = response.json()
# Check that user_info exists and has the user_id
user_info = user_data.get('user_info', {})
return user_info.get('user_id') == user_id
return False
except Exception as e:
logger.warning(
'litellm_user_exists_check_failed',
extra={'user_id': user_id, 'error': str(e)},
)
return False
@staticmethod
async def _create_user(
client: httpx.AsyncClient,
email: str | None,
keycloak_user_id: str,
):
) -> bool:
"""Create a user in LiteLLM.
Returns True if the user was created or already exists and is verified,
False if creation failed and user does not exist.
"""
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
logger.warning('LiteLLM API configuration not found')
return
return False
response = await client.post(
f'{LITE_LLM_API_URL}/user/new',
json={
@@ -716,17 +773,33 @@ class LiteLlmManager:
'user_id': keycloak_user_id,
},
)
return
# Verify the user actually exists before returning success
user_exists = await LiteLlmManager._user_exists(
client, keycloak_user_id
)
if not user_exists:
logger.error(
'litellm_user_claimed_exists_but_not_found',
extra={
'user_id': keycloak_user_id,
'status_code': response.status_code,
'text': response.text,
},
)
return False
return True
logger.error(
'error_creating_litellm_user',
extra={
'status_code': response.status_code,
'text': response.text,
'user_id': [keycloak_user_id],
'user_id': keycloak_user_id,
'email': None,
},
)
return False
response.raise_for_status()
return True
@staticmethod
async def _get_user(client: httpx.AsyncClient, user_id: str) -> dict | None:
@@ -1450,6 +1523,7 @@ class LiteLlmManager:
create_team = staticmethod(with_http_client(_create_team))
get_team = staticmethod(with_http_client(_get_team))
update_team = staticmethod(with_http_client(_update_team))
user_exists = staticmethod(with_http_client(_user_exists))
create_user = staticmethod(with_http_client(_create_user))
get_user = staticmethod(with_http_client(_get_user))
update_user = staticmethod(with_http_client(_update_user))

284
enterprise/tests/unit/test_lite_llm_manager.py
View File

@@ -239,6 +239,16 @@ class TestLiteLlmManager:
mock_404_response = MagicMock()
mock_404_response.status_code = 404
mock_404_response.is_success = False
mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError(
message='Not Found', request=MagicMock(), response=mock_404_response
)
# Mock user exists check response
mock_user_exists_response = MagicMock()
mock_user_exists_response.is_success = True
mock_user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_token_manager = MagicMock()
mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock(
@@ -246,12 +256,8 @@ class TestLiteLlmManager:
)
mock_client = AsyncMock()
mock_client.get.return_value = mock_404_response
mock_client.get.return_value.raise_for_status.side_effect = (
httpx.HTTPStatusError(
message='Not Found', request=MagicMock(), response=mock_404_response
)
)
# First GET is for _get_team (404), second GET is for _user_exists (success)
mock_client.get.side_effect = [mock_404_response, mock_user_exists_response]
mock_client.post.return_value = mock_response
mock_client_class = MagicMock()
@@ -274,8 +280,8 @@ class TestLiteLlmManager:
assert result.llm_api_key.get_secret_value() == 'test-api-key'
assert result.llm_base_url == 'http://test.com'
# Verify API calls were made (get_team + 4 posts)
assert mock_client.get.call_count == 1 # get_team
# Verify API calls were made (get_team + user_exists + 4 posts)
assert mock_client.get.call_count == 2 # get_team + user_exists
assert (
mock_client.post.call_count == 4
) # create_team, add_user_to_team, delete_key_by_alias, generate_key
@@ -294,13 +300,21 @@ class TestLiteLlmManager:
}
mock_team_response.raise_for_status = MagicMock()
# Mock user exists check response
mock_user_exists_response = MagicMock()
mock_user_exists_response.is_success = True
mock_user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_token_manager = MagicMock()
mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock(
return_value={'email': 'test@example.com'}
)
mock_client = AsyncMock()
mock_client.get.return_value = mock_team_response
# First GET is for _get_team (success), second GET is for _user_exists (success)
mock_client.get.side_effect = [mock_team_response, mock_user_exists_response]
mock_client.post.return_value = mock_response
mock_client_class = MagicMock()
@@ -320,8 +334,8 @@ class TestLiteLlmManager:
assert result is not None
# Verify _get_team was called first
mock_client.get.assert_called_once()
get_call_url = mock_client.get.call_args[0][0]
assert mock_client.get.call_count == 2 # get_team + user_exists
get_call_url = mock_client.get.call_args_list[0][0][0]
assert 'team/info' in get_call_url
assert 'test-org-id' in get_call_url
@@ -343,19 +357,25 @@ class TestLiteLlmManager:
mock_404_response = MagicMock()
mock_404_response.status_code = 404
mock_404_response.is_success = False
mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError(
message='Not Found', request=MagicMock(), response=mock_404_response
)
mock_token_manager = MagicMock()
mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock(
return_value={'email': 'test@example.com'}
)
# Mock user exists check response
mock_user_exists_response = MagicMock()
mock_user_exists_response.is_success = True
mock_user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_client = AsyncMock()
mock_client.get.return_value = mock_404_response
mock_client.get.return_value.raise_for_status.side_effect = (
httpx.HTTPStatusError(
message='Not Found', request=MagicMock(), response=mock_404_response
)
)
# First GET is for _get_team (404), second GET is for _user_exists (success)
mock_client.get.side_effect = [mock_404_response, mock_user_exists_response]
mock_client.post.return_value = mock_response
mock_client_class = MagicMock()
@@ -393,6 +413,16 @@ class TestLiteLlmManager:
mock_404_response = MagicMock()
mock_404_response.status_code = 404
mock_404_response.is_success = False
mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError(
message='Not Found', request=MagicMock(), response=mock_404_response
)
# Mock user exists check response
mock_user_exists_response = MagicMock()
mock_user_exists_response.is_success = True
mock_user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_token_manager = MagicMock()
mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock(
@@ -400,12 +430,8 @@ class TestLiteLlmManager:
)
mock_client = AsyncMock()
mock_client.get.return_value = mock_404_response
mock_client.get.return_value.raise_for_status.side_effect = (
httpx.HTTPStatusError(
message='Not Found', request=MagicMock(), response=mock_404_response
)
)
# First GET is for _get_team (404), second GET is for _user_exists (success)
mock_client.get.side_effect = [mock_404_response, mock_user_exists_response]
mock_client.post.return_value = mock_response
mock_client_class = MagicMock()
@@ -833,15 +859,16 @@ class TestLiteLlmManager:
@pytest.mark.asyncio
async def test_create_user_success(self, mock_http_client, mock_response):
"""Test successful _create_user operation."""
"""Test successful _create_user operation returns True."""
mock_http_client.post.return_value = mock_response
with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'):
with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'):
await LiteLlmManager._create_user(
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
assert result is True
mock_http_client.post.assert_called_once()
call_args = mock_http_client.post.call_args
assert 'http://test.com/user/new' in call_args[0]
@@ -850,7 +877,7 @@ class TestLiteLlmManager:
@pytest.mark.asyncio
async def test_create_user_duplicate_email(self, mock_http_client, mock_response):
"""Test _create_user with duplicate email handling."""
"""Test _create_user with duplicate email handling returns True."""
# First call fails with duplicate email
error_response = MagicMock()
error_response.is_success = False
@@ -862,23 +889,81 @@ class TestLiteLlmManager:
with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'):
with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'):
await LiteLlmManager._create_user(
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
assert result is True
assert mock_http_client.post.call_count == 2
# Second call should have None email
second_call_args = mock_http_client.post.call_args_list[1]
assert second_call_args[1]['json']['user_email'] is None
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_user_exists_returns_true(self, mock_http_client):
"""Test _user_exists returns True when user exists in LiteLLM."""
# Arrange
user_response = MagicMock()
user_response.is_success = True
user_response.json.return_value = {
'user_info': {'user_id': 'test-user-id', 'email': 'test@example.com'}
}
mock_http_client.get.return_value = user_response
# Act
result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id')
# Assert
assert result is True
mock_http_client.get.assert_called_once()
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_user_exists_returns_false_when_not_found(self, mock_http_client):
"""Test _user_exists returns False when user not found."""
# Arrange
user_response = MagicMock()
user_response.is_success = False
mock_http_client.get.return_value = user_response
# Act
result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id')
# Assert
assert result is False
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_user_exists_returns_false_on_mismatched_user_id(
self, mock_http_client
):
"""Test _user_exists returns False when returned user_id doesn't match."""
# Arrange
user_response = MagicMock()
user_response.is_success = True
user_response.json.return_value = {
'user_info': {'user_id': 'different-user-id'}
}
mock_http_client.get.return_value = user_response
# Act
result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id')
# Assert
assert result is False
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.logger')
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_create_user_already_exists_with_409_status_code(
async def test_create_user_already_exists_and_verified(
self, mock_logger, mock_http_client
):
"""Test _create_user handles 409 Conflict when user already exists."""
"""Test _create_user returns True when user already exists and is verified."""
# Arrange
first_response = MagicMock()
first_response.is_success = False
@@ -890,14 +975,141 @@ class TestLiteLlmManager:
second_response.status_code = 409
second_response.text = 'User with id test-user-id already exists'
user_exists_response = MagicMock()
user_exists_response.is_success = True
user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_http_client.post.side_effect = [first_response, second_response]
mock_http_client.get.return_value = user_exists_response
# Act
await LiteLlmManager._create_user(
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
# Assert
assert result is True
mock_logger.warning.assert_any_call(
'litellm_user_already_exists',
extra={'user_id': 'test-user-id'},
)
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.logger')
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_create_user_already_exists_but_not_found_returns_false(
self, mock_logger, mock_http_client
):
"""Test _create_user returns False when LiteLLM claims user exists but verification fails."""
# Arrange
first_response = MagicMock()
first_response.is_success = False
first_response.status_code = 400
first_response.text = 'duplicate email'
second_response = MagicMock()
second_response.is_success = False
second_response.status_code = 409
second_response.text = 'User with id test-user-id already exists'
user_not_exists_response = MagicMock()
user_not_exists_response.is_success = False
mock_http_client.post.side_effect = [first_response, second_response]
mock_http_client.get.return_value = user_not_exists_response
# Act
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
# Assert
assert result is False
mock_logger.error.assert_any_call(
'litellm_user_claimed_exists_but_not_found',
extra={
'user_id': 'test-user-id',
'status_code': 409,
'text': 'User with id test-user-id already exists',
},
)
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.logger')
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_create_user_failure_returns_false(
self, mock_logger, mock_http_client
):
"""Test _create_user returns False when creation fails with non-'already exists' error."""
# Arrange
first_response = MagicMock()
first_response.is_success = False
first_response.status_code = 400
first_response.text = 'duplicate email'
second_response = MagicMock()
second_response.is_success = False
second_response.status_code = 500
second_response.text = 'Internal server error'
mock_http_client.post.side_effect = [first_response, second_response]
# Act
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
# Assert
assert result is False
mock_logger.error.assert_any_call(
'error_creating_litellm_user',
extra={
'status_code': 500,
'text': 'Internal server error',
'user_id': 'test-user-id',
'email': None,
},
)
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.logger')
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com')
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key')
async def test_create_user_already_exists_with_409_status_code(
self, mock_logger, mock_http_client
):
"""Test _create_user handles 409 Conflict when user already exists and verifies."""
# Arrange
first_response = MagicMock()
first_response.is_success = False
first_response.status_code = 400
first_response.text = 'duplicate email'
second_response = MagicMock()
second_response.is_success = False
second_response.status_code = 409
second_response.text = 'User with id test-user-id already exists'
user_exists_response = MagicMock()
user_exists_response.is_success = True
user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_http_client.post.side_effect = [first_response, second_response]
mock_http_client.get.return_value = user_exists_response
# Act
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
# Assert
assert result is True
mock_logger.warning.assert_any_call(
'litellm_user_already_exists',
extra={'user_id': 'test-user-id'},
@@ -910,7 +1122,7 @@ class TestLiteLlmManager:
async def test_create_user_already_exists_with_400_status_code(
self, mock_logger, mock_http_client
):
"""Test _create_user handles 400 Bad Request when user already exists."""
"""Test _create_user handles 400 Bad Request when user already exists and verifies."""
# Arrange
first_response = MagicMock()
first_response.is_success = False
@@ -922,14 +1134,22 @@ class TestLiteLlmManager:
second_response.status_code = 400
second_response.text = 'User already exists'
user_exists_response = MagicMock()
user_exists_response.is_success = True
user_exists_response.json.return_value = {
'user_info': {'user_id': 'test-user-id'}
}
mock_http_client.post.side_effect = [first_response, second_response]
mock_http_client.get.return_value = user_exists_response
# Act
await LiteLlmManager._create_user(
result = await LiteLlmManager._create_user(
mock_http_client, 'test@example.com', 'test-user-id'
)
# Assert
assert result is True
mock_logger.warning.assert_any_call(
'litellm_user_already_exists',
extra={'user_id': 'test-user-id'},