From d43ff825340d1f7fc0c244f2d77f36fa9f8b8730 Mon Sep 17 00:00:00 2001 From: Tim O'Farrell Date: Fri, 6 Feb 2026 09:30:12 -0700 Subject: [PATCH] feat: Add BYOR export flag to org for LLM key access control (#12753) Co-authored-by: openhands Co-authored-by: hieptl --- .../091_add_byor_export_enabled_flag.py | 46 ++++++++++++++ enterprise/server/routes/api_keys.py | 60 ++++++++++++++++++- enterprise/storage/org.py | 1 + enterprise/storage/user_store.py | 14 +++++ .../tests/unit/server/routes/test_api_keys.py | 43 +++++++++++-- .../features/settings/api-keys-manager.tsx | 46 +++++++++++++- frontend/src/hooks/query/use-llm-api-key.ts | 28 ++++++++- frontend/src/i18n/declaration.ts | 2 + frontend/src/i18n/translation.json | 32 ++++++++++ 9 files changed, 262 insertions(+), 10 deletions(-) create mode 100644 enterprise/migrations/versions/091_add_byor_export_enabled_flag.py diff --git a/enterprise/migrations/versions/091_add_byor_export_enabled_flag.py b/enterprise/migrations/versions/091_add_byor_export_enabled_flag.py new file mode 100644 index 0000000000..98eb10c6be --- /dev/null +++ b/enterprise/migrations/versions/091_add_byor_export_enabled_flag.py @@ -0,0 +1,46 @@ +"""Add byor_export_enabled flag to org table. + +Revision ID: 091 +Revises: 090 +Create Date: 2025-01-15 00:00:00.000000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = '091' +down_revision: Union[str, None] = '090' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Add byor_export_enabled column to org table with default false + op.add_column( + 'org', + sa.Column( + 'byor_export_enabled', + sa.Boolean, + nullable=False, + server_default=sa.text('false'), + ), + ) + + # Set byor_export_enabled to true for orgs that have completed billing sessions + op.execute( + sa.text(""" + UPDATE org SET byor_export_enabled = TRUE + WHERE id IN ( + SELECT DISTINCT org_id FROM billing_sessions + WHERE status = 'completed' AND org_id IS NOT NULL + ) + """) + ) + + +def downgrade() -> None: + op.drop_column('org', 'byor_export_enabled') diff --git a/enterprise/server/routes/api_keys.py b/enterprise/server/routes/api_keys.py index 049df8dd4f..711937cf2e 100644 --- a/enterprise/server/routes/api_keys.py +++ b/enterprise/server/routes/api_keys.py @@ -6,12 +6,30 @@ from storage.api_key_store import ApiKeyStore from storage.lite_llm_manager import LiteLlmManager from storage.org_member import OrgMember from storage.org_member_store import OrgMemberStore +from storage.org_store import OrgStore from storage.user_store import UserStore from openhands.core.logger import openhands_logger as logger from openhands.server.user_auth import get_user_id +async def check_byor_export_enabled(user_id: str) -> bool: + """Check if BYOR export is enabled for the user's current org. + + Returns True if the user's current org has byor_export_enabled set to True. + Returns False if the user is not found, has no current org, or the flag is False. + """ + user = await UserStore.get_user_by_id_async(user_id) + if not user or not user.current_org_id: + return False + + org = OrgStore.get_org_by_id(user.current_org_id) + if not org: + return False + + return org.byor_export_enabled + + # Helper functions for BYOR API key management async def get_byor_key_from_db(user_id: str) -> str | None: """Get the BYOR key from the database for a user.""" @@ -52,7 +70,6 @@ async def store_byor_key_in_db(user_id: str, key: str) -> None: async def generate_byor_key(user_id: str) -> str | None: """Generate a new BYOR key for a user.""" - try: user = await UserStore.get_user_by_id_async(user_id) if not user: @@ -148,6 +165,26 @@ class LlmApiKeyResponse(BaseModel): key: str | None +class ByorPermittedResponse(BaseModel): + permitted: bool + + +@api_router.get('/llm/byor/permitted', response_model=ByorPermittedResponse) +async def check_byor_permitted(user_id: str = Depends(get_user_id)): + """Check if BYOR key export is permitted for the user's current org.""" + try: + permitted = await check_byor_export_enabled(user_id) + return {'permitted': permitted} + except Exception as e: + logger.exception( + 'Error checking BYOR export permission', extra={'error': str(e)} + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail='Failed to check BYOR export permission', + ) + + @api_router.post('', response_model=ApiKeyCreateResponse) async def create_api_key(key_data: ApiKeyCreate, user_id: str = Depends(get_user_id)): """Create a new API key for the authenticated user.""" @@ -253,8 +290,17 @@ async def get_llm_api_key_for_byor(user_id: str = Depends(get_user_id)): This endpoint validates that the key exists in LiteLLM before returning it. If validation fails, it automatically generates a new key to ensure users always receive a working key. + + Returns 402 Payment Required if BYOR export is not enabled for the user's org. """ try: + # Check if BYOR export is enabled for the user's org + if not await check_byor_export_enabled(user_id): + raise HTTPException( + status_code=status.HTTP_402_PAYMENT_REQUIRED, + detail='BYOR key export is not enabled. Purchase credits to enable this feature.', + ) + # Check if the BYOR key exists in the database byor_key = await get_byor_key_from_db(user_id) if byor_key: @@ -310,10 +356,20 @@ async def get_llm_api_key_for_byor(user_id: str = Depends(get_user_id)): @api_router.post('/llm/byor/refresh', response_model=LlmApiKeyResponse) async def refresh_llm_api_key_for_byor(user_id: str = Depends(get_user_id)): - """Refresh the LLM API key for BYOR (Bring Your Own Runtime) for the authenticated user.""" + """Refresh the LLM API key for BYOR (Bring Your Own Runtime) for the authenticated user. + + Returns 402 Payment Required if BYOR export is not enabled for the user's org. + """ logger.info('Starting BYOR LLM API key refresh', extra={'user_id': user_id}) try: + # Check if BYOR export is enabled for the user's org + if not await check_byor_export_enabled(user_id): + raise HTTPException( + status_code=status.HTTP_402_PAYMENT_REQUIRED, + detail='BYOR key export is not enabled. Purchase credits to enable this feature.', + ) + # Get the existing BYOR key from the database existing_byor_key = await get_byor_key_from_db(user_id) diff --git a/enterprise/storage/org.py b/enterprise/storage/org.py index 1e05007bcd..45d64a0b83 100644 --- a/enterprise/storage/org.py +++ b/enterprise/storage/org.py @@ -46,6 +46,7 @@ class Org(Base): # type: ignore v1_enabled = Column(Boolean, nullable=True) conversation_expiration = Column(Integer, nullable=True) condenser_max_size = Column(Integer, nullable=True) + byor_export_enabled = Column(Boolean, nullable=False, default=False) # Relationships org_members = relationship('OrgMember', back_populates='org') diff --git a/enterprise/storage/user_store.py b/enterprise/storage/user_store.py index e1f95658e6..fd909b11cc 100644 --- a/enterprise/storage/user_store.py +++ b/enterprise/storage/user_store.py @@ -172,6 +172,19 @@ class UserStore: ) decrypted_user_settings = UserSettings(**kwargs) with session_maker() as session: + # Check if user has completed billing sessions to enable BYOR export + from storage.billing_session import BillingSession + + has_completed_billing = ( + session.query(BillingSession) + .filter( + BillingSession.user_id == user_id, + BillingSession.status == 'completed', + ) + .first() + is not None + ) + # create personal org org = Org( id=uuid.UUID(user_id), @@ -180,6 +193,7 @@ class UserStore: contact_name=resolve_display_name(user_info) or user_info.get('username', ''), contact_email=user_info['email'], + byor_export_enabled=has_completed_billing, ) session.add(org) diff --git a/enterprise/tests/unit/server/routes/test_api_keys.py b/enterprise/tests/unit/server/routes/test_api_keys.py index 7f08425758..703a0cf467 100644 --- a/enterprise/tests/unit/server/routes/test_api_keys.py +++ b/enterprise/tests/unit/server/routes/test_api_keys.py @@ -182,16 +182,18 @@ class TestGetLlmApiKeyForByor: """Test the get_llm_api_key_for_byor endpoint.""" @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') @patch('server.routes.api_keys.store_byor_key_in_db') @patch('server.routes.api_keys.generate_byor_key') @patch('server.routes.api_keys.get_byor_key_from_db') async def test_no_key_in_database_generates_new( - self, mock_get_key, mock_generate_key, mock_store_key + self, mock_get_key, mock_generate_key, mock_store_key, mock_check_enabled ): """Test that when no key exists in database, a new one is generated.""" # Arrange user_id = 'user-123' new_key = 'sk-new-generated-key' + mock_check_enabled.return_value = True mock_get_key.return_value = None mock_generate_key.return_value = new_key mock_store_key.return_value = None @@ -201,20 +203,23 @@ class TestGetLlmApiKeyForByor: # Assert assert result == {'key': new_key} + mock_check_enabled.assert_called_once_with(user_id) mock_get_key.assert_called_once_with(user_id) mock_generate_key.assert_called_once_with(user_id) mock_store_key.assert_called_once_with(user_id, new_key) @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') @patch('storage.lite_llm_manager.LiteLlmManager.verify_key') @patch('server.routes.api_keys.get_byor_key_from_db') async def test_valid_key_in_database_returns_key( - self, mock_get_key, mock_verify_key + self, mock_get_key, mock_verify_key, mock_check_enabled ): """Test that when a valid key exists in database, it is returned.""" # Arrange user_id = 'user-123' existing_key = 'sk-existing-valid-key' + mock_check_enabled.return_value = True mock_get_key.return_value = existing_key mock_verify_key.return_value = True @@ -223,10 +228,12 @@ class TestGetLlmApiKeyForByor: # Assert assert result == {'key': existing_key} + mock_check_enabled.assert_called_once_with(user_id) mock_get_key.assert_called_once_with(user_id) mock_verify_key.assert_called_once_with(existing_key, user_id) @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') @patch('server.routes.api_keys.store_byor_key_in_db') @patch('server.routes.api_keys.generate_byor_key') @patch('server.routes.api_keys.delete_byor_key_from_litellm') @@ -239,12 +246,14 @@ class TestGetLlmApiKeyForByor: mock_delete_key, mock_generate_key, mock_store_key, + mock_check_enabled, ): """Test that when an invalid key exists in database, it is regenerated.""" # Arrange user_id = 'user-123' invalid_key = 'sk-invalid-key' new_key = 'sk-new-generated-key' + mock_check_enabled.return_value = True mock_get_key.return_value = invalid_key mock_verify_key.return_value = False mock_delete_key.return_value = True @@ -256,6 +265,7 @@ class TestGetLlmApiKeyForByor: # Assert assert result == {'key': new_key} + mock_check_enabled.assert_called_once_with(user_id) mock_get_key.assert_called_once_with(user_id) mock_verify_key.assert_called_once_with(invalid_key, user_id) mock_delete_key.assert_called_once_with(user_id, invalid_key) @@ -263,6 +273,7 @@ class TestGetLlmApiKeyForByor: mock_store_key.assert_called_once_with(user_id, new_key) @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') @patch('server.routes.api_keys.store_byor_key_in_db') @patch('server.routes.api_keys.generate_byor_key') @patch('server.routes.api_keys.delete_byor_key_from_litellm') @@ -275,12 +286,14 @@ class TestGetLlmApiKeyForByor: mock_delete_key, mock_generate_key, mock_store_key, + mock_check_enabled, ): """Test that even if deletion fails, regeneration still proceeds.""" # Arrange user_id = 'user-123' invalid_key = 'sk-invalid-key' new_key = 'sk-new-generated-key' + mock_check_enabled.return_value = True mock_get_key.return_value = invalid_key mock_verify_key.return_value = False mock_delete_key.return_value = False # Deletion fails @@ -292,19 +305,22 @@ class TestGetLlmApiKeyForByor: # Assert assert result == {'key': new_key} + mock_check_enabled.assert_called_once_with(user_id) mock_delete_key.assert_called_once_with(user_id, invalid_key) mock_generate_key.assert_called_once_with(user_id) mock_store_key.assert_called_once_with(user_id, new_key) @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') @patch('server.routes.api_keys.generate_byor_key') @patch('server.routes.api_keys.get_byor_key_from_db') async def test_key_generation_failure_raises_exception( - self, mock_get_key, mock_generate_key + self, mock_get_key, mock_generate_key, mock_check_enabled ): """Test that when key generation fails, an HTTPException is raised.""" # Arrange user_id = 'user-123' + mock_check_enabled.return_value = True mock_get_key.return_value = None mock_generate_key.return_value = None @@ -316,11 +332,15 @@ class TestGetLlmApiKeyForByor: assert 'Failed to generate new BYOR LLM API key' in exc_info.value.detail @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') @patch('server.routes.api_keys.get_byor_key_from_db') - async def test_database_error_raises_exception(self, mock_get_key): + async def test_database_error_raises_exception( + self, mock_get_key, mock_check_enabled + ): """Test that database errors are properly handled.""" # Arrange user_id = 'user-123' + mock_check_enabled.return_value = True mock_get_key.side_effect = Exception('Database connection error') # Act & Assert @@ -330,6 +350,21 @@ class TestGetLlmApiKeyForByor: assert exc_info.value.status_code == 500 assert 'Failed to retrieve BYOR LLM API key' in exc_info.value.detail + @pytest.mark.asyncio + @patch('server.routes.api_keys.check_byor_export_enabled') + async def test_byor_export_disabled_returns_402(self, mock_check_enabled): + """Test that when BYOR export is disabled, 402 is returned.""" + # Arrange + user_id = 'user-123' + mock_check_enabled.return_value = False + + # Act & Assert + with pytest.raises(HTTPException) as exc_info: + await get_llm_api_key_for_byor(user_id=user_id) + + assert exc_info.value.status_code == 402 + assert 'BYOR key export is not enabled' in exc_info.value.detail + class TestDeleteByorKeyFromLitellm: """Test the delete_byor_key_from_litellm function with alias cleanup.""" diff --git a/frontend/src/components/features/settings/api-keys-manager.tsx b/frontend/src/components/features/settings/api-keys-manager.tsx index 20a8807aa0..143aed0671 100644 --- a/frontend/src/components/features/settings/api-keys-manager.tsx +++ b/frontend/src/components/features/settings/api-keys-manager.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; import { useTranslation, Trans } from "react-i18next"; +import { useNavigate } from "react-router"; import { FaTrash, FaEye, FaEyeSlash, FaCopy } from "react-icons/fa6"; import { I18nKey } from "#/i18n/declaration"; import { BrandButton } from "#/components/features/settings/brand-button"; @@ -19,12 +20,41 @@ import { useRefreshLlmApiKey } from "#/hooks/mutation/use-refresh-llm-api-key"; interface LlmApiKeyManagerProps { llmApiKey: { key: string | null } | undefined; isLoadingLlmKey: boolean; + isPaymentRequired: boolean; refreshLlmApiKey: ReturnType; } +function LlmApiKeyPaywall() { + const { t } = useTranslation(); + const navigate = useNavigate(); + + return ( +
+

+ {t(I18nKey.SETTINGS$LLM_API_KEY)} +

+
+

+ {t(I18nKey.SETTINGS$LLM_API_KEY_PAYWALL_MESSAGE)} +

+
+ navigate("/settings/billing")} + > + {t(I18nKey.SETTINGS$LLM_API_KEY_BUY_NOW)} + +
+
+
+ ); +} + function LlmApiKeyManager({ llmApiKey, isLoadingLlmKey, + isPaymentRequired, refreshLlmApiKey, }: LlmApiKeyManagerProps) { const { t } = useTranslation(); @@ -45,6 +75,11 @@ function LlmApiKeyManager({ }); }; + // Show paywall if payment is required + if (isPaymentRequired) { + return ; + } + if (isLoadingLlmKey || !llmApiKey) { return null; } @@ -206,7 +241,11 @@ function ApiKeysTable({ apiKeys, isLoading, onDeleteKey }: ApiKeysTableProps) { export function ApiKeysManager() { const { t } = useTranslation(); const { data: apiKeys = [], isLoading, error } = useApiKeys(); - const { data: llmApiKey, isLoading: isLoadingLlmKey } = useLlmApiKey(); + const { + data: llmApiKey, + isLoading: isLoadingLlmKey, + isPaymentRequired, + } = useLlmApiKey(); const refreshLlmApiKey = useRefreshLlmApiKey(); const [createModalOpen, setCreateModalOpen] = useState(false); const [deleteModalOpen, setDeleteModalOpen] = useState(false); @@ -215,8 +254,8 @@ export function ApiKeysManager() { useState(null); const [showNewKeyModal, setShowNewKeyModal] = useState(false); - // Display error toast if the query fails - if (error) { + // Display error toast if the query fails (but not for payment required) + if (error && !isPaymentRequired) { displayErrorToast(t(I18nKey.ERROR$GENERIC)); } @@ -251,6 +290,7 @@ export function ApiKeysManager() { diff --git a/frontend/src/hooks/query/use-llm-api-key.ts b/frontend/src/hooks/query/use-llm-api-key.ts index 10b91b0e43..0970ee98e8 100644 --- a/frontend/src/hooks/query/use-llm-api-key.ts +++ b/frontend/src/hooks/query/use-llm-api-key.ts @@ -1,4 +1,5 @@ import { useQuery } from "@tanstack/react-query"; +import { AxiosError } from "axios"; import { openHands } from "#/api/open-hands-axios"; import { useConfig } from "./use-config"; @@ -8,10 +9,15 @@ export interface LlmApiKeyResponse { key: string | null; } +export interface LlmApiKeyError { + isPaymentRequired: boolean; + message?: string; +} + export function useLlmApiKey() { const { data: config } = useConfig(); - return useQuery({ + const query = useQuery({ queryKey: [LLM_API_KEY_QUERY_KEY], enabled: config?.app_mode === "saas", queryFn: async () => { @@ -21,5 +27,25 @@ export function useLlmApiKey() { }, staleTime: 1000 * 60 * 5, // 5 minutes gcTime: 1000 * 60 * 15, // 15 minutes + retry: (failureCount, error) => { + // Don't retry on 402 Payment Required + if (error instanceof AxiosError && error.response?.status === 402) { + return false; + } + return failureCount < 3; + }, + // Disable global error toast - we handle 402 errors in the UI + meta: { disableToast: true }, }); + + // Check if the error is a 402 Payment Required + const isPaymentRequired = + query.error instanceof AxiosError && query.error.response?.status === 402; + + return { + data: query.data, + error: query.error, + isLoading: query.isLoading, + isPaymentRequired, + }; } diff --git a/frontend/src/i18n/declaration.ts b/frontend/src/i18n/declaration.ts index 126894e6c9..b570bdbbaa 100644 --- a/frontend/src/i18n/declaration.ts +++ b/frontend/src/i18n/declaration.ts @@ -402,6 +402,8 @@ export enum I18nKey { SETTINGS$LLM_API_KEY = "SETTINGS$LLM_API_KEY", SETTINGS$LLM_API_KEY_DESCRIPTION = "SETTINGS$LLM_API_KEY_DESCRIPTION", SETTINGS$REFRESH_LLM_API_KEY = "SETTINGS$REFRESH_LLM_API_KEY", + SETTINGS$LLM_API_KEY_PAYWALL_MESSAGE = "SETTINGS$LLM_API_KEY_PAYWALL_MESSAGE", + SETTINGS$LLM_API_KEY_BUY_NOW = "SETTINGS$LLM_API_KEY_BUY_NOW", SETTINGS$CONFIRMATION_MODE = "SETTINGS$CONFIRMATION_MODE", SETTINGS$CONFIRMATION_MODE_TOOLTIP = "SETTINGS$CONFIRMATION_MODE_TOOLTIP", SETTINGS$CONFIRMATION_MODE_LOCK_TOOLTIP = "SETTINGS$CONFIRMATION_MODE_LOCK_TOOLTIP", diff --git a/frontend/src/i18n/translation.json b/frontend/src/i18n/translation.json index 3cabb8cbc5..787f287c6d 100644 --- a/frontend/src/i18n/translation.json +++ b/frontend/src/i18n/translation.json @@ -6431,6 +6431,38 @@ "ja": "APIキーを更新", "uk": "Оновити ключ API" }, + "SETTINGS$LLM_API_KEY_PAYWALL_MESSAGE": { + "en": "Purchase at least $10 in credits to get access to OpenHands LLM key for use with OpenHands CLI and SDK.", + "zh-CN": "购买至少10美元的积分，即可获得OpenHands LLM密钥，用于OpenHands CLI和SDK。", + "zh-TW": "購買至少10美元的積分，即可獲得OpenHands LLM金鑰，用於OpenHands CLI和SDK。", + "de": "Kaufen Sie mindestens $10 an Guthaben, um Zugang zum OpenHands LLM-Schlüssel für die Verwendung mit OpenHands CLI und SDK zu erhalten.", + "ko-KR": "OpenHands CLI 및 SDK에서 사용할 OpenHands LLM 키에 액세스하려면 최소 $10의 크레딧을 구매하세요.", + "no": "Kjøp minst $10 i kreditter for å få tilgang til OpenHands LLM-nøkkel for bruk med OpenHands CLI og SDK.", + "it": "Acquista almeno $10 di crediti per ottenere l'accesso alla chiave LLM OpenHands per l'uso con OpenHands CLI e SDK.", + "pt": "Compre pelo menos $10 em créditos para obter acesso à chave LLM OpenHands para uso com OpenHands CLI e SDK.", + "es": "Compre al menos $10 en créditos para obtener acceso a la clave LLM de OpenHands para usar con OpenHands CLI y SDK.", + "ar": "اشترِ ما لا يقل عن 10 دولارات من الرصيد للحصول على مفتاح LLM OpenHands لاستخدامه مع OpenHands CLI و SDK.", + "fr": "Achetez au moins 10$ de crédits pour accéder à la clé LLM OpenHands à utiliser avec OpenHands CLI et SDK.", + "tr": "OpenHands CLI ve SDK ile kullanmak için OpenHands LLM Anahtarına erişim sağlamak için en az 10$ kredi satın alın.", + "ja": "OpenHands CLIおよびSDKで使用するOpenHands LLMキーにアクセスするには、少なくとも$10のクレジットを購入してください。", + "uk": "Придбайте кредитів мінімум на $10, щоб отримати доступ до ключа LLM OpenHands для використання з OpenHands CLI та SDK." + }, + "SETTINGS$LLM_API_KEY_BUY_NOW": { + "en": "Buy Now", + "zh-CN": "立即购买", + "zh-TW": "立即購買", + "de": "Jetzt kaufen", + "ko-KR": "지금 구매", + "no": "Kjøp nå", + "it": "Acquista ora", + "pt": "Comprar agora", + "es": "Comprar ahora", + "ar": "اشترِ الآن", + "fr": "Acheter maintenant", + "tr": "Şimdi Satın Al", + "ja": "今すぐ購入", + "uk": "Купити зараз" + }, "SETTINGS$CONFIRMATION_MODE": { "en": "Enable Confirmation Mode", "de": "Bestätigungsmodus aktivieren",