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(backend): byor key alias cleanup on 404 (#12491)

Browse Source
This commit is contained in:
Hiep Le
2026-01-19 13:12:04 +07:00
committed by GitHub
parent 57929b73ee
commit e8ba741f28
4 changed files with 273 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
enterprise/server/routes/api_keys.py
View File

@@ -99,10 +99,19 @@ async def generate_byor_key(user_id: str) -> str | None:
async def delete_byor_key_from_litellm(user_id: str, byor_key: str) -> bool:
"""Delete the BYOR key from LiteLLM using the key directly."""
"""Delete the BYOR key from LiteLLM using the key directly.
Also attempts to delete by key alias if the key is not found,
to clean up orphaned aliases that could block key regeneration.
"""
try:
await LiteLlmManager.delete_key(byor_key)
# Get user to construct the key alias
user = await UserStore.get_user_by_id_async(user_id)
key_alias = None
if user and user.current_org_id:
key_alias = f'BYOR Key - user {user_id}, org {user.current_org_id}'
await LiteLlmManager.delete_key(byor_key, key_alias=key_alias)
logger.info(
'Successfully deleted BYOR key from LiteLLM',
extra={'user_id': user_id},

42
enterprise/storage/lite_llm_manager.py
View File

@@ -699,10 +699,45 @@ class LiteLlmManager:
'key_spend': key_info.get('spend'),
}
@staticmethod
async def _delete_key_by_alias(
client: httpx.AsyncClient,
key_alias: str,
):
"""Delete a key from LiteLLM by its alias.
This is a best-effort operation that logs but does not raise on failure.
"""
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
logger.warning('LiteLLM API configuration not found')
return
response = await client.post(
f'{LITE_LLM_API_URL}/key/delete',
json={
'key_aliases': [key_alias],
},
)
if response.is_success:
logger.info(
'LiteLlmManager:_delete_key_by_alias:key_deleted',
extra={'key_alias': key_alias},
)
elif response.status_code != 404:
# Log non-404 errors but don't fail
logger.warning(
'error_deleting_key_by_alias',
extra={
'key_alias': key_alias,
'status_code': response.status_code,
'text': response.text,
},
)
@staticmethod
async def _delete_key(
client: httpx.AsyncClient,
key_id: str,
key_alias: str | None = None,
):
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
logger.warning('LiteLLM API configuration not found')
@@ -713,10 +748,13 @@ class LiteLlmManager:
'keys': [key_id],
},
)
# Failed to key...
# Failed to delete key...
if not response.is_success:
if response.status_code == 404:
# key doesn't exist, just return
# Key doesn't exist by key_id. If we have a key_alias,
# try deleting by alias to clean up any orphaned alias.
if key_alias:
await LiteLlmManager._delete_key_by_alias(client, key_alias)
return
logger.error(
'error_deleting_key',

97
enterprise/tests/unit/server/routes/test_api_keys.py
View File

@@ -6,6 +6,7 @@ import httpx
import pytest
from fastapi import HTTPException
from server.routes.api_keys import (
delete_byor_key_from_litellm,
get_llm_api_key_for_byor,
)
from storage.lite_llm_manager import LiteLlmManager
@@ -328,3 +329,99 @@ class TestGetLlmApiKeyForByor:
assert exc_info.value.status_code == 500
assert 'Failed to retrieve BYOR LLM API key' in exc_info.value.detail
class TestDeleteByorKeyFromLitellm:
"""Test the delete_byor_key_from_litellm function with alias cleanup."""
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LiteLlmManager.delete_key')
@patch('storage.user_store.UserStore.get_user_by_id_async')
async def test_delete_constructs_alias_from_user(
self, mock_get_user, mock_delete_key
):
"""Test that delete_byor_key_from_litellm constructs key alias from user."""
# Arrange
user_id = 'user-123'
org_id = 'org-456'
byor_key = 'sk-byor-key-to-delete'
expected_alias = f'BYOR Key - user {user_id}, org {org_id}'
mock_user = MagicMock()
mock_user.current_org_id = org_id
mock_get_user.return_value = mock_user
mock_delete_key.return_value = None
# Act
result = await delete_byor_key_from_litellm(user_id, byor_key)
# Assert
assert result is True
mock_get_user.assert_called_once_with(user_id)
mock_delete_key.assert_called_once_with(byor_key, key_alias=expected_alias)
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LiteLlmManager.delete_key')
@patch('storage.user_store.UserStore.get_user_by_id_async')
async def test_delete_without_user_passes_no_alias(
self, mock_get_user, mock_delete_key
):
"""Test that when user is not found, no alias is passed."""
# Arrange
user_id = 'user-123'
byor_key = 'sk-byor-key-to-delete'
mock_get_user.return_value = None
mock_delete_key.return_value = None
# Act
result = await delete_byor_key_from_litellm(user_id, byor_key)
# Assert
assert result is True
mock_delete_key.assert_called_once_with(byor_key, key_alias=None)
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LiteLlmManager.delete_key')
@patch('storage.user_store.UserStore.get_user_by_id_async')
async def test_delete_without_org_id_passes_no_alias(
self, mock_get_user, mock_delete_key
):
"""Test that when user has no current_org_id, no alias is passed."""
# Arrange
user_id = 'user-123'
byor_key = 'sk-byor-key-to-delete'
mock_user = MagicMock()
mock_user.current_org_id = None
mock_get_user.return_value = mock_user
mock_delete_key.return_value = None
# Act
result = await delete_byor_key_from_litellm(user_id, byor_key)
# Assert
assert result is True
mock_delete_key.assert_called_once_with(byor_key, key_alias=None)
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LiteLlmManager.delete_key')
@patch('storage.user_store.UserStore.get_user_by_id_async')
async def test_delete_returns_false_on_exception(
self, mock_get_user, mock_delete_key
):
"""Test that exceptions during deletion return False."""
# Arrange
user_id = 'user-123'
byor_key = 'sk-byor-key-to-delete'
mock_user = MagicMock()
mock_user.current_org_id = 'org-456'
mock_get_user.return_value = mock_user
mock_delete_key.side_effect = Exception('LiteLLM API error')
# Act
result = await delete_byor_key_from_litellm(user_id, byor_key)
# Assert
assert result is False

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

@@ -557,6 +557,131 @@ class TestLiteLlmManager:
# Should not raise an exception for 404
await LiteLlmManager._delete_key(mock_http_client, 'test-key-id')
@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_delete_key_not_found_with_alias_triggers_alias_deletion(
self, mock_http_client
):
"""Test _delete_key falls back to alias deletion when key_id returns 404."""
# Arrange
not_found_response = MagicMock()
not_found_response.is_success = False
not_found_response.status_code = 404
not_found_response.text = 'Key not found'
alias_success_response = MagicMock()
alias_success_response.is_success = True
alias_success_response.status_code = 200
mock_http_client.post.side_effect = [not_found_response, alias_success_response]
# Act
await LiteLlmManager._delete_key(
mock_http_client, 'test-key-id', key_alias='BYOR Key - user 123, org 456'
)
# Assert
assert mock_http_client.post.call_count == 2
first_call = mock_http_client.post.call_args_list[0]
assert first_call[1]['json']['keys'] == ['test-key-id']
second_call = mock_http_client.post.call_args_list[1]
assert second_call[1]['json']['key_aliases'] == ['BYOR Key - user 123, org 456']
@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_delete_key_not_found_without_alias_no_fallback(
self, mock_http_client
):
"""Test _delete_key without alias does not attempt alias deletion on 404."""
# Arrange
not_found_response = MagicMock()
not_found_response.is_success = False
not_found_response.status_code = 404
not_found_response.text = 'Key not found'
mock_http_client.post.return_value = not_found_response
# Act
await LiteLlmManager._delete_key(mock_http_client, 'test-key-id')
# Assert
assert mock_http_client.post.call_count == 1
@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_delete_key_by_alias_success(self, mock_http_client, mock_response):
"""Test successful _delete_key_by_alias operation."""
# Arrange
mock_http_client.post.return_value = mock_response
# Act
await LiteLlmManager._delete_key_by_alias(
mock_http_client, 'BYOR Key - user 123, org 456'
)
# Assert
mock_http_client.post.assert_called_once()
call_args = mock_http_client.post.call_args
assert 'http://test.com/key/delete' in call_args[0]
assert call_args[1]['json']['key_aliases'] == ['BYOR Key - user 123, org 456']
@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_delete_key_by_alias_not_found(self, mock_http_client):
"""Test _delete_key_by_alias when alias is not found (404)."""
# Arrange
not_found_response = MagicMock()
not_found_response.is_success = False
not_found_response.status_code = 404
not_found_response.text = 'Key alias not found'
mock_http_client.post.return_value = not_found_response
# Act & Assert - should not raise exception for 404
await LiteLlmManager._delete_key_by_alias(
mock_http_client, 'BYOR Key - user 123, org 456'
)
@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_delete_key_by_alias_server_error_logs_warning(
self, mock_logger, mock_http_client
):
"""Test _delete_key_by_alias logs warning for non-404 errors."""
# Arrange
error_response = MagicMock()
error_response.is_success = False
error_response.status_code = 500
error_response.text = 'Internal server error'
mock_http_client.post.return_value = error_response
# Act
await LiteLlmManager._delete_key_by_alias(
mock_http_client, 'BYOR Key - user 123, org 456'
)
# Assert
mock_logger.warning.assert_called_once()
call_args = mock_logger.warning.call_args
assert call_args[0][0] == 'error_deleting_key_by_alias'
@pytest.mark.asyncio
@patch('storage.lite_llm_manager.LITE_LLM_API_URL', None)
@patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None)
async def test_delete_key_by_alias_missing_config(self, mock_http_client):
"""Test _delete_key_by_alias returns early when config is missing."""
# Act
await LiteLlmManager._delete_key_by_alias(
mock_http_client, 'BYOR Key - user 123, org 456'
)
# Assert
mock_http_client.post.assert_not_called()
@pytest.mark.asyncio
async def test_with_http_client_decorator(self):
"""Test the with_http_client decorator functionality."""