From b7e5c9d25bb8998ac4bfbc40bb3293cd1260617b Mon Sep 17 00:00:00 2001 From: chuckbutkus Date: Fri, 13 Mar 2026 18:39:07 -0400 Subject: [PATCH] Use a flag to indicate if new users should use V1 (#13393) Co-authored-by: openhands --- enterprise/server/constants.py | 3 + enterprise/storage/org_store.py | 3 + enterprise/storage/user_store.py | 7 ++ enterprise/tests/unit/test_org_store.py | 80 +++++++++++++++++++++++ enterprise/tests/unit/test_user_store.py | 82 ++++++++++++++++++++++++ 5 files changed, 175 insertions(+) diff --git a/enterprise/server/constants.py b/enterprise/server/constants.py index 670f28c34a..4e9734c770 100644 --- a/enterprise/server/constants.py +++ b/enterprise/server/constants.py @@ -77,6 +77,9 @@ PERMITTED_CORS_ORIGINS = [ ) ] +# Controls whether new orgs/users default to V1 API (env: DEFAULT_V1_ENABLED) +DEFAULT_V1_ENABLED = os.getenv('DEFAULT_V1_ENABLED', '1').lower() in ('1', 'true') + def build_litellm_proxy_model_path(model_name: str) -> str: """Build the LiteLLM proxy model path based on model name. diff --git a/enterprise/storage/org_store.py b/enterprise/storage/org_store.py index 3332024a1c..23015ea5be 100644 --- a/enterprise/storage/org_store.py +++ b/enterprise/storage/org_store.py @@ -6,6 +6,7 @@ from typing import Optional from uuid import UUID from server.constants import ( + DEFAULT_V1_ENABLED, LITE_LLM_API_URL, ORG_SETTINGS_VERSION, get_default_litellm_model, @@ -36,6 +37,8 @@ class OrgStore: org = Org(**kwargs) org.org_version = ORG_SETTINGS_VERSION org.default_llm_model = get_default_litellm_model() + if org.v1_enabled is None: + org.v1_enabled = DEFAULT_V1_ENABLED session.add(org) await session.commit() await session.refresh(org) diff --git a/enterprise/storage/user_store.py b/enterprise/storage/user_store.py index 9d363d10c4..ad42449285 100644 --- a/enterprise/storage/user_store.py +++ b/enterprise/storage/user_store.py @@ -7,6 +7,7 @@ from uuid import UUID from server.auth.token_manager import TokenManager from server.constants import ( + DEFAULT_V1_ENABLED, LITE_LLM_API_URL, ORG_SETTINGS_VERSION, PERSONAL_WORKSPACE_VERSION_TO_MODEL, @@ -241,6 +242,10 @@ class UserStore: if hasattr(org, key): setattr(org, key, value) + # Apply DEFAULT_V1_ENABLED for migrated orgs if v1_enabled was not set + if org.v1_enabled is None: + org.v1_enabled = DEFAULT_V1_ENABLED + user_kwargs = UserStore.get_kwargs_from_user_settings( decrypted_user_settings ) @@ -892,6 +897,8 @@ class UserStore: language='en', enable_proactive_conversation_starters=True ) + default_settings.v1_enabled = DEFAULT_V1_ENABLED + from storage.lite_llm_manager import LiteLlmManager settings = await LiteLlmManager.create_entries( diff --git a/enterprise/tests/unit/test_org_store.py b/enterprise/tests/unit/test_org_store.py index 7bada3b7b1..0c8a8173f9 100644 --- a/enterprise/tests/unit/test_org_store.py +++ b/enterprise/tests/unit/test_org_store.py @@ -144,6 +144,86 @@ async def test_create_org(async_session_maker, mock_litellm_api): assert org.id is not None +@pytest.mark.asyncio +async def test_create_org_v1_enabled_defaults_to_true_when_default_is_true( + async_session_maker, mock_litellm_api +): + """ + GIVEN: DEFAULT_V1_ENABLED is True and org.v1_enabled is not specified (None) + WHEN: create_org is called + THEN: org.v1_enabled should be set to True + """ + with ( + patch('storage.org_store.a_session_maker', async_session_maker), + patch('storage.org_store.DEFAULT_V1_ENABLED', True), + ): + org = await OrgStore.create_org(kwargs={'name': 'test-org-v1-default-true'}) + + assert org is not None + assert org.v1_enabled is True + + +@pytest.mark.asyncio +async def test_create_org_v1_enabled_defaults_to_false_when_default_is_false( + async_session_maker, mock_litellm_api +): + """ + GIVEN: DEFAULT_V1_ENABLED is False and org.v1_enabled is not specified (None) + WHEN: create_org is called + THEN: org.v1_enabled should be set to False + """ + with ( + patch('storage.org_store.a_session_maker', async_session_maker), + patch('storage.org_store.DEFAULT_V1_ENABLED', False), + ): + org = await OrgStore.create_org(kwargs={'name': 'test-org-v1-default-false'}) + + assert org is not None + assert org.v1_enabled is False + + +@pytest.mark.asyncio +async def test_create_org_v1_enabled_explicit_false_overrides_default_true( + async_session_maker, mock_litellm_api +): + """ + GIVEN: DEFAULT_V1_ENABLED is True but org.v1_enabled is explicitly set to False + WHEN: create_org is called + THEN: org.v1_enabled should stay False (explicit value wins over default) + """ + with ( + patch('storage.org_store.a_session_maker', async_session_maker), + patch('storage.org_store.DEFAULT_V1_ENABLED', True), + ): + org = await OrgStore.create_org( + kwargs={'name': 'test-org-v1-explicit-false', 'v1_enabled': False} + ) + + assert org is not None + assert org.v1_enabled is False + + +@pytest.mark.asyncio +async def test_create_org_v1_enabled_explicit_true_overrides_default_false( + async_session_maker, mock_litellm_api +): + """ + GIVEN: DEFAULT_V1_ENABLED is False but org.v1_enabled is explicitly set to True + WHEN: create_org is called + THEN: org.v1_enabled should stay True (explicit value wins over default) + """ + with ( + patch('storage.org_store.a_session_maker', async_session_maker), + patch('storage.org_store.DEFAULT_V1_ENABLED', False), + ): + org = await OrgStore.create_org( + kwargs={'name': 'test-org-v1-explicit-true', 'v1_enabled': True} + ) + + assert org is not None + assert org.v1_enabled is True + + @pytest.mark.asyncio async def test_get_org_by_name(async_session_maker, mock_litellm_api): # Test getting org by name diff --git a/enterprise/tests/unit/test_user_store.py b/enterprise/tests/unit/test_user_store.py index 61ea471883..e5c4262711 100644 --- a/enterprise/tests/unit/test_user_store.py +++ b/enterprise/tests/unit/test_user_store.py @@ -101,6 +101,72 @@ async def test_create_default_settings_with_litellm(mock_litellm_api): assert settings.llm_base_url == 'http://test.url' +@pytest.mark.asyncio +async def test_create_default_settings_v1_enabled_true_when_default_is_true( + mock_litellm_api, +): + """ + GIVEN: DEFAULT_V1_ENABLED is True + WHEN: create_default_settings is called + THEN: The default_settings.v1_enabled should be set to True + """ + org_id = str(uuid.uuid4()) + user_id = str(uuid.uuid4()) + + # Track the settings passed to LiteLlmManager.create_entries + captured_settings = None + + async def capture_create_entries(_org_id, _user_id, settings, _create_user): + nonlocal captured_settings + captured_settings = settings + return settings + + with ( + patch('storage.user_store.DEFAULT_V1_ENABLED', True), + patch( + 'storage.lite_llm_manager.LiteLlmManager.create_entries', + side_effect=capture_create_entries, + ), + ): + await UserStore.create_default_settings(org_id, user_id) + + assert captured_settings is not None + assert captured_settings.v1_enabled is True + + +@pytest.mark.asyncio +async def test_create_default_settings_v1_enabled_false_when_default_is_false( + mock_litellm_api, +): + """ + GIVEN: DEFAULT_V1_ENABLED is False + WHEN: create_default_settings is called + THEN: The default_settings.v1_enabled should be set to False + """ + org_id = str(uuid.uuid4()) + user_id = str(uuid.uuid4()) + + # Track the settings passed to LiteLlmManager.create_entries + captured_settings = None + + async def capture_create_entries(_org_id, _user_id, settings, _create_user): + nonlocal captured_settings + captured_settings = settings + return settings + + with ( + patch('storage.user_store.DEFAULT_V1_ENABLED', False), + patch( + 'storage.lite_llm_manager.LiteLlmManager.create_entries', + side_effect=capture_create_entries, + ), + ): + await UserStore.create_default_settings(org_id, user_id) + + assert captured_settings is not None + assert captured_settings.v1_enabled is False + + # --- Tests for get_user_by_id --- @@ -1243,3 +1309,19 @@ async def test_migrate_user_sql_multiple_conversations(async_session_maker): assert ( row.org_id == user_uuid_str ), f'org_id should match: {row.org_id} vs {user_uuid_str}' + + +# Note: The v1_enabled logic in migrate_user follows the same pattern as OrgStore.create_org: +# if org.v1_enabled is None: +# org.v1_enabled = DEFAULT_V1_ENABLED +# +# This behavior is tested in test_org_store.py via: +# - test_create_org_v1_enabled_defaults_to_true_when_default_is_true +# - test_create_org_v1_enabled_defaults_to_false_when_default_is_false +# - test_create_org_v1_enabled_explicit_false_overrides_default_true +# - test_create_org_v1_enabled_explicit_true_overrides_default_false +# +# Testing migrate_user directly is impractical due to its complex raw SQL migration +# statements that have SQLite/UUID compatibility issues in the test environment. +# The SQL migration tests above (test_migrate_user_sql_type_handling, etc.) verify +# the SQL operations work correctly with proper type handling.