Fix migrations

enterprise/server/routes/auth.py

@@ -145,18 +145,7 @@ async def keycloak_callback(
     user_id = user_info['sub']
     user = UserStore.get_user_by_id(user_id)
     if not user:
-        user_settings = None
-        with session_maker() as session:
-            user_settings = (
-                session.query(UserSettings)
-                .filter(UserSettings.keycloak_user_id == user_id)
-                .first()
-            )
-        if user_settings:
-            user = await UserStore.migrate_user(user_id, user_settings, user_info)
-        else:
-            # new user
-            user = await UserStore.create_user(user_id, user_info)
+        user = await UserStore.create_user(user_id, user_info)
 
     if not user:
         logger.error(f'Failed to authenticate user {user_info["preferred_username"]}')

enterprise/server/routes/integration/slack.py

@@ -199,20 +199,11 @@ async def keycloak_callback(
     keycloak_user_id = user_info['sub']
     user = UserStore.get_user_by_id(keycloak_user_id)
     if not user:
-        user_settings = None
-        with session_maker() as session:
-            user_settings = (
-                session.query(UserSettings)
-                .filter(UserSettings.keycloak_user_id == keycloak_user_id)
-                .first()
-            )
-        if not user_settings:
-            return _html_response(
-                title='Failed to authenticate.',
-                description=f'Please re-login into <a href="{HOST_URL}" style="color:#ecedee;text-decoration:underline;">OpenHands Cloud</a>. Then try <a href="https://docs.all-hands.dev/usage/cloud/slack-installation" style="color:#ecedee;text-decoration:underline;">installing the OpenHands Slack App</a> again',
-                status_code=400,
-            )
-        user = await UserStore.migrate_user(keycloak_user_id, user_settings, user_info)
+        return _html_response(
+            title='Failed to authenticate.',
+            description=f'Please re-login into <a href="{HOST_URL}" style="color:#ecedee;text-decoration:underline;">OpenHands Cloud</a>. Then try <a href="https://docs.all-hands.dev/usage/cloud/slack-installation" style="color:#ecedee;text-decoration:underline;">installing the OpenHands Slack App</a> again',
+            status_code=400,
+        )
 
     # These tokens are offline access tokens - store them!
     await token_manager.store_offline_token(keycloak_user_id, keycloak_refresh_token)

enterprise/storage/lite_llm_manager.py

@@ -85,6 +85,7 @@ class LiteLlmManager:
         org_id: str,
         keycloak_user_id: str,
         user_settings: UserSettings,
+        keycloak_user_info: dict,
     ) -> UserSettings | None:
         logger.info(
             'SettingsStore:umigrate_lite_llm_entries:start',
@@ -97,11 +98,6 @@ class LiteLlmManager:
         key = LITE_LLM_API_KEY
         if not local_deploy:
             # Get user info to add to litellm
-            token_manager = TokenManager()
-            keycloak_user_info = (
-                await token_manager.get_user_info_from_user_id(keycloak_user_id) or {}
-            )
-
             async with httpx.AsyncClient(
                 headers={
                     'x-goog-api-key': LITE_LLM_API_KEY,

enterprise/storage/saas_settings_store.py

@@ -68,22 +68,8 @@ class SaasSettingsStore(OssSettingsStore):
     async def load(self) -> Settings | None:
         user = UserStore.get_user_by_id(self.user_id)
         if not user:
-            # Check if we need to migrate from user_settings
-            user_settings = None
-            with session_maker() as session:
-                user_settings = (
-                    session.query(UserSettings)
-                    .filter(
-                        UserSettings.keycloak_user_id == self.user_id,
-                        UserSettings.already_migrated.is_(False),
-                    )
-                    .first()
-                )
-            if user_settings:
-                user = await UserStore.migrate_user(self.user_id, user_settings)
-            else:
-                logger.error(f'User not found for ID {self.user_id}')
-                return None
+            logger.error(f'User not found for ID {self.user_id}')
+            return None
 
         org_id = user.current_org_id
         org_member: OrgMember = None

enterprise/storage/user_store.py

@@ -5,9 +5,11 @@ Store class for managing users.
 import uuid
 from typing import Optional
 
-from server.logger import logger
+from openhands.utils.async_utils import GENERAL_TIMEOUT, call_async_from_sync
 from sqlalchemy import text
 from sqlalchemy.orm import joinedload
+
+from server.logger import logger
 from storage.database import session_maker
 from storage.encrypt_utils import decrypt_legacy_model
 from storage.org import Org
@@ -82,15 +84,11 @@ class UserStore:
     async def migrate_user(
         user_id: str,
         user_settings: UserSettings,
-        user_info: dict | None = None,
+        user_info: dict,
     ) -> User:
         if not user_id or not user_settings:
             return None
 
-        # Check if user is already migrated to prevent double migration
-        if user_settings.already_migrated is True:
-            logger.warning(f'User {user_id} already migrated, skipping')
-            return UserStore.get_user_by_id(user_id)
         kwargs = decrypt_legacy_model(
             [
                 'llm_api_key',
@@ -103,26 +101,18 @@ class UserStore:
         decrypted_user_settings = UserSettings(**kwargs)
         with session_maker() as session:
             # create personal org
-            contact_name = (
-                user_info['preferred_username']
-                if user_info
-                else decrypted_user_settings.email.split('@')[0]
-            )
-            contact_email = (
-                user_info['email'] if user_info else decrypted_user_settings.email
-            )
             org = Org(
                 id=uuid.UUID(user_id),
                 name=f'user_{user_id}_org',
-                contact_name=contact_name,
-                contact_email=contact_email,
+                contact_name=user_info['username'],
+                contact_email=user_info['email'],
             )
             session.add(org)
 
             from storage.lite_llm_manager import LiteLlmManager
 
             await LiteLlmManager.migrate_entries(
-                str(org.id), user_id, decrypted_user_settings
+                str(org.id), user_id, decrypted_user_settings, user_info
             )
 
             # avoids circular reference. This migrate method is temprorary until all users are migrated.
@@ -243,12 +233,44 @@ class UserStore:
     def get_user_by_id(user_id: str) -> Optional[User]:
         """Get user by Keycloak user ID."""
         with session_maker() as session:
-            return (
+            user = (
                 session.query(User)
                 .options(joinedload(User.org_members))
                 .filter(User.id == uuid.UUID(user_id))
                 .first()
             )
+            if user:
+                return user
+
+            # Check if we need to migrate from user_settings
+            user_settings = (
+                session.query(UserSettings)
+                .filter(
+                    UserSettings.keycloak_user_id == user_id,
+                    UserSettings.already_migrated.is_(False),
+                )
+                .first()
+            )
+            if user_settings:
+                from server.auth.token_manager import TokenManager
+
+                token_manager = TokenManager()
+                token_manager.get_user_info_from_user_id(user_id)
+                user_info = call_async_from_sync(
+                    token_manager.get_user_info_from_user_id,
+                    GENERAL_TIMEOUT,
+                    user_id,
+                )
+                user = call_async_from_sync(
+                    UserStore.migrate_user,
+                    GENERAL_TIMEOUT,
+                    user_id,
+                    user_settings,
+                    user_info,
+                )
+                return user
+            else:
+                return None
 
     @staticmethod
     def list_users() -> list[User]: