feat: Add BYOR export flag to org for LLM key access control (#12753)

Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: hieptl <hieptl.developer@gmail.com>

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')

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)
 

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')

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)
 

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."""

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<typeof useRefreshLlmApiKey>;
 }
 
+function LlmApiKeyPaywall() {
+  const { t } = useTranslation();
+  const navigate = useNavigate();
+
+  return (
+    <div className="border-b border-gray-200 pb-6 mb-6 flex flex-col gap-6">
+      <h3 className="text-xl font-medium text-white">
+        {t(I18nKey.SETTINGS$LLM_API_KEY)}
+      </h3>
+      <div className="bg-base-tertiary rounded-md p-4 flex flex-col gap-4">
+        <p className="text-sm text-gray-300">
+          {t(I18nKey.SETTINGS$LLM_API_KEY_PAYWALL_MESSAGE)}
+        </p>
+        <div>
+          <BrandButton
+            type="button"
+            variant="primary"
+            onClick={() => navigate("/settings/billing")}
+          >
+            {t(I18nKey.SETTINGS$LLM_API_KEY_BUY_NOW)}
+          </BrandButton>
+        </div>
+      </div>
+    </div>
+  );
+}
+
 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 <LlmApiKeyPaywall />;
+  }
+
   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<CreateApiKeyResponse | null>(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() {
         <LlmApiKeyManager
           llmApiKey={llmApiKey}
           isLoadingLlmKey={isLoadingLlmKey}
+          isPaymentRequired={isPaymentRequired}
           refreshLlmApiKey={refreshLlmApiKey}
         />
 

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,
+  };
 }

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",

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",