diff --git a/enterprise/server/routes/auth.py b/enterprise/server/routes/auth.py index 3c0596677d..3511310c85 100644 --- a/enterprise/server/routes/auth.py +++ b/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"]}') diff --git a/enterprise/server/routes/integration/slack.py b/enterprise/server/routes/integration/slack.py index 532b487367..b8103a33af 100644 --- a/enterprise/server/routes/integration/slack.py +++ b/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 OpenHands Cloud. Then try installing the OpenHands Slack App 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 OpenHands Cloud. Then try installing the OpenHands Slack App again', + status_code=400, + ) # These tokens are offline access tokens - store them! await token_manager.store_offline_token(keycloak_user_id, keycloak_refresh_token) diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py index 30058915b5..5c77c8c18d 100644 --- a/enterprise/storage/lite_llm_manager.py +++ b/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, diff --git a/enterprise/storage/saas_settings_store.py b/enterprise/storage/saas_settings_store.py index cbacfc2415..19faeaa0a0 100644 --- a/enterprise/storage/saas_settings_store.py +++ b/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 diff --git a/enterprise/storage/user_store.py b/enterprise/storage/user_store.py index 14fb1e8e04..da042396ec 100644 --- a/enterprise/storage/user_store.py +++ b/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]: