Fix migration for byor key

- Created encrypt_legacy_model(), encrypt_legacy_kwargs(), and encrypt_legacy_value()
  functions in encrypt_utils.py as the inverse of decrypt_legacy_* functions
- Updated downgrade_user() in UserStore to encrypt sensitive fields
  (llm_api_key, llm_api_key_for_byor, search_api_key, sandbox_api_key)
  before merging UserSettings into the database

Co-authored-by: openhands <openhands@all-hands.dev>

enterprise/storage/encrypt_utils.py

@@ -98,6 +98,29 @@ def decrypt_legacy_value(value: str | SecretStr) -> str:
         return get_fernet().decrypt(b64decode(value.encode())).decode()
 
 
+def encrypt_legacy_model(encrypt_keys: list, model_instance) -> dict:
+    return encrypt_legacy_kwargs(encrypt_keys, model_to_kwargs(model_instance))
+
+
+def encrypt_legacy_kwargs(encrypt_keys: list, kwargs: dict) -> dict:
+    for key, value in kwargs.items():
+        if value is None:
+            continue
+        if key in encrypt_keys:
+            value = encrypt_legacy_value(value)
+            kwargs[key] = value
+    return kwargs
+
+
+def encrypt_legacy_value(value: str | SecretStr) -> str:
+    if isinstance(value, SecretStr):
+        return b64encode(
+            get_fernet().encrypt(value.get_secret_value().encode())
+        ).decode()
+    else:
+        return b64encode(get_fernet().encrypt(value.encode())).decode()
+
+
 def get_fernet():
     global _fernet
     if _fernet is None:

enterprise/storage/user_store.py

@@ -17,7 +17,10 @@ from server.logger import logger
 from sqlalchemy import select, text
 from sqlalchemy.orm import joinedload
 from storage.database import a_session_maker, session_maker
-from storage.encrypt_utils import decrypt_legacy_model
+from storage.encrypt_utils import (
+    decrypt_legacy_model,
+    encrypt_legacy_value,
+)
 from storage.org import Org
 from storage.org_member import OrgMember
 from storage.role_store import RoleStore
@@ -238,7 +241,6 @@ class UserStore:
             if not custom_settings:
                 del org_member_kwargs['llm_model']
                 del org_member_kwargs['llm_base_url']
-                del org_member_kwargs['llm_api_key_for_byor']
 
             org_member = OrgMember(
                 org_id=org.id,
@@ -423,20 +425,10 @@ class UserStore:
                 org_member = org_members[0]
                 is_new_signup = True
 
-                # Create a new user_settings entry from org_member data
+                # Create a new user_settings entry from OrgMember, User, and Org data
                 # This is needed for new sign-ups who don't have user_settings
-                user_settings = UserSettings(
-                    keycloak_user_id=user_id,
-                    llm_api_key=org_member.llm_api_key.get_secret_value()
-                    if org_member.llm_api_key
-                    else None,
-                    llm_api_key_for_byor=org_member.llm_api_key_for_byor.get_secret_value()
-                    if org_member.llm_api_key_for_byor
-                    else None,
-                    llm_model=org_member.llm_model,
-                    llm_base_url=org_member.llm_base_url,
-                    max_iterations=org_member.max_iterations,
-                    already_migrated=False,  # Will be set correctly below
+                user_settings = UserStore._create_user_settings_from_entities(
+                    user_id, org_member, user, org
                 )
                 session.add(user_settings)
                 session.flush()
@@ -565,8 +557,21 @@ class UserStore:
                 {'org_id': user_uuid},
             )
 
-            # Step 8: Set already_migrated=False on user_settings
+            # Step 8: Set already_migrated=False on user_settings and encrypt fields
             user_settings.already_migrated = False
+
+            # Re-encrypt the sensitive fields before storing in the DB
+            encrypt_keys = [
+                'llm_api_key',
+                'llm_api_key_for_byor',
+                'search_api_key',
+                'sandbox_api_key',
+            ]
+            for key in encrypt_keys:
+                value = getattr(user_settings, key, None)
+                if value is not None:
+                    setattr(user_settings, key, encrypt_legacy_value(value))
+
             session.merge(user_settings)
 
             session.commit()
@@ -767,6 +772,96 @@ class UserStore:
         }
         return kwargs
 
+    @staticmethod
+    def _create_user_settings_from_entities(
+        user_id: str, org_member: OrgMember, user: User, org: Org
+    ) -> UserSettings:
+        """Create UserSettings from OrgMember, User, and Org data.
+
+        Uses OrgMember values first. If an OrgMember field is None and there's
+        a corresponding "default_" field in Org, use the Org value.
+        Also pulls relevant fields from User.
+
+        Args:
+            user_id: The Keycloak user ID
+            org_member: The OrgMember entity
+            user: The User entity
+            org: The Org entity
+
+        Returns:
+            A new UserSettings object populated from the entities
+        """
+        # Mapping from OrgMember fields to corresponding Org "default_" fields
+        org_member_to_org_default = {
+            'llm_model': 'default_llm_model',
+            'llm_base_url': 'default_llm_base_url',
+            'max_iterations': 'default_max_iterations',
+        }
+
+        def get_value_with_org_fallback(field_name: str, org_member_value):
+            """Get value from OrgMember, falling back to Org default if None."""
+            if org_member_value is not None:
+                return org_member_value
+            org_default_field = org_member_to_org_default.get(field_name)
+            if org_default_field and hasattr(org, org_default_field):
+                return getattr(org, org_default_field)
+            return None
+
+        # Get values from OrgMember with Org fallback for fields with default_ prefix
+        llm_model = get_value_with_org_fallback('llm_model', org_member.llm_model)
+        llm_base_url = get_value_with_org_fallback(
+            'llm_base_url', org_member.llm_base_url
+        )
+        max_iterations = get_value_with_org_fallback(
+            'max_iterations', org_member.max_iterations
+        )
+
+        return UserSettings(
+            keycloak_user_id=user_id,
+            # OrgMember fields
+            llm_api_key=org_member.llm_api_key.get_secret_value()
+            if org_member.llm_api_key
+            else None,
+            llm_api_key_for_byor=org_member.llm_api_key_for_byor.get_secret_value()
+            if org_member.llm_api_key_for_byor
+            else None,
+            llm_model=llm_model,
+            llm_base_url=llm_base_url,
+            max_iterations=max_iterations,
+            # User fields
+            accepted_tos=user.accepted_tos,
+            enable_sound_notifications=user.enable_sound_notifications,
+            language=user.language,
+            user_consents_to_analytics=user.user_consents_to_analytics,
+            email=user.email,
+            email_verified=user.email_verified,
+            git_user_name=user.git_user_name,
+            git_user_email=user.git_user_email,
+            # Org fields
+            agent=org.agent,
+            security_analyzer=org.security_analyzer,
+            confirmation_mode=org.confirmation_mode,
+            remote_runtime_resource_factor=org.remote_runtime_resource_factor,
+            enable_default_condenser=org.enable_default_condenser,
+            billing_margin=org.billing_margin,
+            enable_proactive_conversation_starters=org.enable_proactive_conversation_starters,
+            sandbox_base_container_image=org.sandbox_base_container_image,
+            sandbox_runtime_container_image=org.sandbox_runtime_container_image,
+            user_version=org.org_version,
+            mcp_config=org.mcp_config,
+            search_api_key=org.search_api_key.get_secret_value()
+            if org.search_api_key
+            else None,
+            sandbox_api_key=org.sandbox_api_key.get_secret_value()
+            if org.sandbox_api_key
+            else None,
+            max_budget_per_task=org.max_budget_per_task,
+            enable_solvability_analysis=org.enable_solvability_analysis,
+            v1_enabled=org.v1_enabled,
+            condenser_max_size=org.condenser_max_size,
+            already_migrated=False,
+        )
+
     @staticmethod
     def _has_custom_settings(
         user_settings: UserSettings, old_user_version: int | None