From 3ec999e88a13aabfae02bfcf57a6434de3384337 Mon Sep 17 00:00:00 2001 From: Tim O'Farrell Date: Fri, 6 Mar 2026 21:48:19 -0700 Subject: [PATCH] Fix LiteLLM key management and user migration SQL queries (#13279) Co-authored-by: openhands --- enterprise/storage/lite_llm_manager.py | 14 +- enterprise/storage/user_store.py | 25 +- .../tests/unit/test_lite_llm_manager.py | 6 +- enterprise/tests/unit/test_user_store.py | 410 ++++++++++++++++++ 4 files changed, 438 insertions(+), 17 deletions(-) diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py index 49af669359..a7e240c7d7 100644 --- a/enterprise/storage/lite_llm_manager.py +++ b/enterprise/storage/lite_llm_manager.py @@ -137,11 +137,23 @@ class LiteLlmManager: client, keycloak_user_id, org_id, team_budget ) + # We delete the key if it already exists. In environments where multiple + # installations are using the same keycloak and litellm instance, this + # will mean other installations will have their key invalidated. + key_alias = get_openhands_cloud_key_alias(keycloak_user_id, org_id) + try: + await LiteLlmManager._delete_key_by_alias(client, key_alias) + except httpx.HTTPStatusError as ex: + if ex.status_code == 404: + logger.debug(f'Key "{key_alias}" did not exist - continuing') + else: + raise + key = await LiteLlmManager._generate_key( client, keycloak_user_id, org_id, - get_openhands_cloud_key_alias(keycloak_user_id, org_id), + key_alias, None, ) diff --git a/enterprise/storage/user_store.py b/enterprise/storage/user_store.py index f5d1c9d27a..9d363d10c4 100644 --- a/enterprise/storage/user_store.py +++ b/enterprise/storage/user_store.py @@ -295,29 +295,28 @@ class UserStore: extra={'user_id': user_id}, ) + user_uuid = uuid.UUID(user_id) + # need to migrate conversation metadata await session.execute( text(""" INSERT INTO conversation_metadata_saas (conversation_id, user_id, org_id) SELECT conversation_id, - :user_id, - :user_id + :user_uuid, + :user_uuid FROM conversation_metadata - WHERE user_id = :user_id + WHERE user_id = :user_id_text """), - {'user_id': user_id}, + {'user_uuid': user_uuid, 'user_id_text': user_id}, ) - # Update org_id for tables that had org_id added - user_uuid = uuid.UUID(user_id) - # Update stripe_customers await session.execute( text( 'UPDATE stripe_customers SET org_id = :org_id WHERE keycloak_user_id = :user_id' ), - {'org_id': user_uuid, 'user_id': user_uuid}, + {'org_id': user_uuid, 'user_id': user_id}, ) # Update slack_users @@ -325,7 +324,7 @@ class UserStore: text( 'UPDATE slack_users SET org_id = :org_id WHERE keycloak_user_id = :user_id' ), - {'org_id': user_uuid, 'user_id': user_uuid}, + {'org_id': user_uuid, 'user_id': user_id}, ) # Update slack_conversation @@ -333,13 +332,13 @@ class UserStore: text( 'UPDATE slack_conversation SET org_id = :org_id WHERE keycloak_user_id = :user_id' ), - {'org_id': user_uuid, 'user_id': user_uuid}, + {'org_id': user_uuid, 'user_id': user_id}, ) # Update api_keys await session.execute( text('UPDATE api_keys SET org_id = :org_id WHERE user_id = :user_id'), - {'org_id': user_uuid, 'user_id': user_uuid}, + {'org_id': user_uuid, 'user_id': user_id}, ) # Update custom_secrets @@ -347,7 +346,7 @@ class UserStore: text( 'UPDATE custom_secrets SET org_id = :org_id WHERE keycloak_user_id = :user_id' ), - {'org_id': user_uuid, 'user_id': user_uuid}, + {'org_id': user_uuid, 'user_id': user_id}, ) # Update billing_sessions @@ -355,7 +354,7 @@ class UserStore: text( 'UPDATE billing_sessions SET org_id = :org_id WHERE user_id = :user_id' ), - {'org_id': user_uuid, 'user_id': user_uuid}, + {'org_id': user_uuid, 'user_id': user_id}, ) await session.commit() diff --git a/enterprise/tests/unit/test_lite_llm_manager.py b/enterprise/tests/unit/test_lite_llm_manager.py index 04ca2347fc..8c26709aa3 100644 --- a/enterprise/tests/unit/test_lite_llm_manager.py +++ b/enterprise/tests/unit/test_lite_llm_manager.py @@ -180,11 +180,11 @@ class TestLiteLlmManager: assert result.llm_api_key.get_secret_value() == 'test-api-key' assert result.llm_base_url == 'http://test.com' - # Verify API calls were made (get_team + 3 posts) + # Verify API calls were made (get_team + 4 posts) assert mock_client.get.call_count == 1 # get_team assert ( - mock_client.post.call_count == 3 - ) # create_team, add_user_to_team, generate_key + mock_client.post.call_count == 4 + ) # create_team, add_user_to_team, delete_key_by_alias, generate_key @pytest.mark.asyncio async def test_create_entries_inherits_existing_team_budget( diff --git a/enterprise/tests/unit/test_user_store.py b/enterprise/tests/unit/test_user_store.py index 6a2ecb41ac..61ea471883 100644 --- a/enterprise/tests/unit/test_user_store.py +++ b/enterprise/tests/unit/test_user_store.py @@ -833,3 +833,413 @@ async def test_release_user_creation_lock_released(): assert result is True mock_redis.delete.assert_called_once() + + +# --- Tests for migrate_user SQL parameter type handling --- + + +@pytest.mark.asyncio +async def test_migrate_user_sql_type_handling(async_session_maker): + """Test that migrate_user correctly handles UUID vs string types in SQL queries. + + This test verifies the fixes for SQL parameter binding issues in _migrate_personal_data + where UUID and string parameters need to be correctly matched to their column types. + + Note: SQLite doesn't natively support UUID types, so we use string representations. + The key verification is that: + 1. String user_ids in WHERE clauses match source tables correctly + 2. UUID values are inserted into target UUID columns correctly + 3. The migration queries don't fail due to type mismatches + """ + from sqlalchemy import text + + user_id = str(uuid.uuid4()) + user_uuid = uuid.UUID(user_id) + # For SQLite raw SQL, use string representation of UUID + user_uuid_str = str(user_uuid) + + # Set up legacy data with string user_ids (as in the old schema) + async with async_session_maker() as session: + # First, add conversation_metadata with user_id as string column + # The current model doesn't have user_id, but the real DB did before migration + # We use raw SQL to add the column and insert test data + await session.execute( + text('ALTER TABLE conversation_metadata ADD COLUMN user_id VARCHAR') + ) + await session.execute( + text( + """ + INSERT INTO conversation_metadata (conversation_id, user_id, conversation_version, created_at, last_updated_at) + VALUES (:conv_id, :user_id, 'V0', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + """ + ), + {'conv_id': 'test-conv-1', 'user_id': user_id}, + ) + + # Create org first (needed for foreign keys) + org = Org(id=user_uuid, name=f'user_{user_id}_org') + session.add(org) + + # Create user (needed for foreign keys) + user = User(id=user_uuid, current_org_id=user_uuid) + session.add(user) + await session.commit() + + # Add stripe_customers with keycloak_user_id as string + from storage.stripe_customer import StripeCustomer + + stripe_customer = StripeCustomer( + keycloak_user_id=user_id, stripe_customer_id='stripe_123' + ) + session.add(stripe_customer) + + # Add slack_users with keycloak_user_id as string + from storage.slack_user import SlackUser + + slack_user = SlackUser( + keycloak_user_id=user_id, + slack_user_id='slack_user_123', + slack_display_name='Test User', + ) + session.add(slack_user) + + # Add slack_conversation with keycloak_user_id as string + from storage.slack_conversation import SlackConversation + + slack_conv = SlackConversation( + conversation_id='slack-conv-1', + channel_id='channel_123', + keycloak_user_id=user_id, + ) + session.add(slack_conv) + + # Add api_keys with user_id as string + from storage.api_key import ApiKey + + api_key = ApiKey(key='api_key_123', user_id=user_id, name='Test API Key') + session.add(api_key) + + # Add custom_secrets with keycloak_user_id as string + from storage.stored_custom_secrets import StoredCustomSecrets + + custom_secret = StoredCustomSecrets( + keycloak_user_id=user_id, + secret_name='test_secret', + secret_value='secret_value', + ) + session.add(custom_secret) + + # Add billing_sessions with user_id as string + from storage.billing_session import BillingSession + + billing_session = BillingSession( + id='billing-session-1', + user_id=user_id, + status='completed', + price=10, + price_code='USD', + ) + session.add(billing_session) + + await session.commit() + + # Now execute the migration SQL statements with the correct parameter types + # This tests the fix: using user_uuid for UUID columns and user_id for string columns + # Note: For SQLite, we use string representation of UUID + + # Test 1: conversation_metadata to conversation_metadata_saas migration + # The fix uses user_uuid (UUID) for inserting into user_id/org_id (UUID columns) + # and user_id_text (string) for comparing with user_id in conversation_metadata (string column) + await session.execute( + text( + """ + INSERT INTO conversation_metadata_saas (conversation_id, user_id, org_id) + SELECT + conversation_id, + :user_uuid, + :user_uuid + FROM conversation_metadata + WHERE user_id = :user_id_text + """ + ), + {'user_uuid': user_uuid_str, 'user_id_text': user_id}, + ) + + # Test 2: Update stripe_customers - org_id is UUID, keycloak_user_id is string + await session.execute( + text( + 'UPDATE stripe_customers SET org_id = :org_id WHERE keycloak_user_id = :user_id' + ), + {'org_id': user_uuid_str, 'user_id': user_id}, + ) + + # Test 3: Update slack_users - org_id is UUID, keycloak_user_id is string + await session.execute( + text( + 'UPDATE slack_users SET org_id = :org_id WHERE keycloak_user_id = :user_id' + ), + {'org_id': user_uuid_str, 'user_id': user_id}, + ) + + # Test 4: Update slack_conversation - org_id is UUID, keycloak_user_id is string + await session.execute( + text( + 'UPDATE slack_conversation SET org_id = :org_id WHERE keycloak_user_id = :user_id' + ), + {'org_id': user_uuid_str, 'user_id': user_id}, + ) + + # Test 5: Update api_keys - org_id is UUID, user_id is string + await session.execute( + text('UPDATE api_keys SET org_id = :org_id WHERE user_id = :user_id'), + {'org_id': user_uuid_str, 'user_id': user_id}, + ) + + # Test 6: Update custom_secrets - org_id is UUID, keycloak_user_id is string + await session.execute( + text( + 'UPDATE custom_secrets SET org_id = :org_id WHERE keycloak_user_id = :user_id' + ), + {'org_id': user_uuid_str, 'user_id': user_id}, + ) + + # Test 7: Update billing_sessions - org_id is UUID, user_id is string + await session.execute( + text( + 'UPDATE billing_sessions SET org_id = :org_id WHERE user_id = :user_id' + ), + {'org_id': user_uuid_str, 'user_id': user_id}, + ) + + await session.commit() + + # Verify the data was migrated correctly + from storage.stored_conversation_metadata_saas import ( + StoredConversationMetadataSaas, + ) + + # Verify conversation_metadata_saas + result = await session.execute( + select(StoredConversationMetadataSaas).filter( + StoredConversationMetadataSaas.conversation_id == 'test-conv-1' + ) + ) + saas_metadata = result.scalars().first() + assert ( + saas_metadata is not None + ), 'conversation_metadata_saas record should exist' + assert saas_metadata.user_id == user_uuid, 'user_id should be UUID type' + assert saas_metadata.org_id == user_uuid, 'org_id should be UUID type' + + # Verify stripe_customers org_id was set + result = await session.execute( + select(StripeCustomer).filter(StripeCustomer.keycloak_user_id == user_id) + ) + stripe_record = result.scalars().first() + assert stripe_record is not None + assert ( + stripe_record.org_id == user_uuid + ), 'stripe_customers.org_id should be UUID' + + # Verify slack_users org_id was set + result = await session.execute( + select(SlackUser).filter(SlackUser.keycloak_user_id == user_id) + ) + slack_user_record = result.scalars().first() + assert slack_user_record is not None + assert ( + slack_user_record.org_id == user_uuid + ), 'slack_users.org_id should be UUID' + + # Verify slack_conversation org_id was set + result = await session.execute( + select(SlackConversation).filter( + SlackConversation.keycloak_user_id == user_id + ) + ) + slack_conv_record = result.scalars().first() + assert slack_conv_record is not None + assert ( + slack_conv_record.org_id == user_uuid + ), 'slack_conversation.org_id should be UUID' + + # Verify api_keys org_id was set + result = await session.execute(select(ApiKey).filter(ApiKey.user_id == user_id)) + api_key_record = result.scalars().first() + assert api_key_record is not None + assert api_key_record.org_id == user_uuid, 'api_keys.org_id should be UUID' + + # Verify custom_secrets org_id was set + result = await session.execute( + select(StoredCustomSecrets).filter( + StoredCustomSecrets.keycloak_user_id == user_id + ) + ) + custom_secret_record = result.scalars().first() + assert custom_secret_record is not None + assert ( + custom_secret_record.org_id == user_uuid + ), 'custom_secrets.org_id should be UUID' + + # Verify billing_sessions org_id was set + result = await session.execute( + select(BillingSession).filter(BillingSession.user_id == user_id) + ) + billing_record = result.scalars().first() + assert billing_record is not None + assert ( + billing_record.org_id == user_uuid + ), 'billing_sessions.org_id should be UUID' + + +@pytest.mark.asyncio +async def test_migrate_user_sql_no_matching_records(async_session_maker): + """Test that migration SQL handles the case where no records match the user_id. + + This verifies that the SQL queries don't fail when there are no matching records. + """ + from sqlalchemy import text + + user_id = str(uuid.uuid4()) + user_uuid = uuid.UUID(user_id) + user_uuid_str = str(user_uuid) + other_user_id = str(uuid.uuid4()) + + # Set up data for a different user + async with async_session_maker() as session: + # Add conversation_metadata with user_id column for a different user + await session.execute( + text('ALTER TABLE conversation_metadata ADD COLUMN user_id VARCHAR') + ) + await session.execute( + text( + """ + INSERT INTO conversation_metadata (conversation_id, user_id, conversation_version, created_at, last_updated_at) + VALUES (:conv_id, :user_id, 'V0', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + """ + ), + {'conv_id': 'other-conv-1', 'user_id': other_user_id}, + ) + + # Create org and user for our test user + org = Org(id=user_uuid, name=f'user_{user_id}_org') + session.add(org) + user = User(id=user_uuid, current_org_id=user_uuid) + session.add(user) + await session.commit() + + # Execute migration SQL for our user (no data should match) + await session.execute( + text( + """ + INSERT INTO conversation_metadata_saas (conversation_id, user_id, org_id) + SELECT + conversation_id, + :user_uuid, + :user_uuid + FROM conversation_metadata + WHERE user_id = :user_id_text + """ + ), + {'user_uuid': user_uuid_str, 'user_id_text': user_id}, + ) + await session.commit() + + # Verify no records were created for our user + from storage.stored_conversation_metadata_saas import ( + StoredConversationMetadataSaas, + ) + + result = await session.execute( + select(StoredConversationMetadataSaas).filter( + StoredConversationMetadataSaas.user_id == user_uuid + ) + ) + records = result.scalars().all() + assert ( + len(records) == 0 + ), 'No records should be created for non-matching user_id' + + +@pytest.mark.asyncio +async def test_migrate_user_sql_multiple_conversations(async_session_maker): + """Test that migration SQL correctly handles multiple conversations for a user.""" + from sqlalchemy import text + + user_id = str(uuid.uuid4()) + user_uuid = uuid.UUID(user_id) + user_uuid_str = str(user_uuid) + + async with async_session_maker() as session: + # Create org and user FIRST (needed for foreign keys) + org = Org(id=user_uuid, name=f'user_{user_id}_org') + session.add(org) + user = User(id=user_uuid, current_org_id=user_uuid) + session.add(user) + await session.commit() + + # Add conversation_metadata with user_id column + await session.execute( + text('ALTER TABLE conversation_metadata ADD COLUMN user_id VARCHAR') + ) + await session.commit() + + # Insert multiple conversations for the same user + for i in range(3): + await session.execute( + text( + """ + INSERT INTO conversation_metadata (conversation_id, user_id, conversation_version, created_at, last_updated_at) + VALUES (:conv_id, :user_id, 'V0', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + """ + ), + {'conv_id': f'test-conv-{i}', 'user_id': user_id}, + ) + + await session.commit() + + # Verify that conversation_metadata was inserted + result = await session.execute( + text('SELECT conversation_id, user_id FROM conversation_metadata') + ) + conv_rows = result.fetchall() + assert ( + len(conv_rows) == 3 + ), f'Expected 3 conversation_metadata rows, got {len(conv_rows)}' + + # Execute migration SQL + await session.execute( + text( + """ + INSERT INTO conversation_metadata_saas (conversation_id, user_id, org_id) + SELECT + conversation_id, + :user_uuid, + :user_uuid + FROM conversation_metadata + WHERE user_id = :user_id_text + """ + ), + {'user_uuid': user_uuid_str, 'user_id_text': user_id}, + ) + await session.commit() + + # Verify all conversations were migrated using raw SQL + # (SQLite stores UUIDs as strings, ORM comparison may differ) + result = await session.execute( + text( + 'SELECT conversation_id, user_id, org_id FROM conversation_metadata_saas WHERE user_id = :user_uuid' + ), + {'user_uuid': user_uuid_str}, + ) + saas_rows = result.fetchall() + assert len(saas_rows) == 3, 'All 3 conversations should be migrated' + + # Verify the user_id and org_id values + for row in saas_rows: + assert ( + row.user_id == user_uuid_str + ), f'user_id should match: {row.user_id} vs {user_uuid_str}' + assert ( + row.org_id == user_uuid_str + ), f'org_id should match: {row.org_id} vs {user_uuid_str}'