Edit: prep for the gpt-4o-mini edit a/b test

lib/shared/src/experimentation/FeatureFlagProvider.ts

@@ -54,6 +54,9 @@ export enum FeatureFlag {
     CodyAutocompleteContextExperimentVariant4 = 'cody-autocomplete-context-experiment-variant-4',
     CodyAutocompleteContextExperimentControl = 'cody-autocomplete-context-experiment-control',
 
+    // Enables gpt-4o-mini as a default Edit model
+    CodyEditDefaultToGpt4oMini = 'cody-edit-default-to-gpt-4o-mini',
+
     // use-ssc-for-cody-subscription is a feature flag that enables the use of SSC as the source of truth for Cody subscription data.
     UseSscForCodySubscription = 'use-ssc-for-cody-subscription',
 

lib/shared/src/models/modelsService.test.ts

@@ -2,6 +2,7 @@ import { Observable } from 'observable-fns'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { currentAuthStatus, mockAuthStatus } from '../auth/authStatus'
 import { AUTH_STATUS_FIXTURE_AUTHED, type AuthenticatedAuthStatus } from '../auth/types'
+import { FeatureFlag, featureFlagProvider } from '../experimentation/FeatureFlagProvider'
 import { firstValueFrom } from '../misc/observable'
 import { DOTCOM_URL } from '../sourcegraph-api/environments'
 import * as userProductSubscriptionModule from '../sourcegraph-api/userProductSubscription'
@@ -23,7 +24,8 @@ const EMPTY_MODELS_DATA: ModelsData = {
 describe('modelsService', () => {
     function modelsServiceWithModels(models: Model[]): ModelsService {
         const modelsService = new ModelsService()
-        modelsService.storage = new TestLocalStorageForModelPreferences()
+        storage = new TestLocalStorageForModelPreferences()
+        modelsService.setStorage(storage)
         // TODO(sqs)#observe: this only mocks tests that don't use modelsService.modelsChanges
         vi.spyOn(modelsService, 'models', 'get').mockReturnValue(models)
         return modelsService
@@ -52,6 +54,7 @@ describe('modelsService', () => {
 
     // Reset service
     let modelsService: ModelsService
+    let storage: TestLocalStorageForModelPreferences
     beforeEach(() => {
         modelsService = new ModelsService()
     })
@@ -157,7 +160,7 @@ describe('modelsService', () => {
 
         function modelsServiceWithModels(models: Model[]): ModelsService {
             const modelsService = new ModelsService()
-            modelsService.storage = new TestLocalStorageForModelPreferences()
+            modelsService.setStorage(new TestLocalStorageForModelPreferences())
             return modelsService
         }
 
@@ -181,8 +184,7 @@ describe('modelsService', () => {
                 Observable.of({
                     localModels: [],
                     primaryModels: [model2chat, model4edit],
-                    preferences:
-                        modelsService.storage?.getModelPreferences()[currentAuthStatus().endpoint]!,
+                    preferences: storage?.getModelPreferences()[currentAuthStatus().endpoint]!,
                 })
             )
             expect(await firstValueFrom(modelsService.getDefaultEditModel())).toBe(model4edit.id)
@@ -201,8 +203,7 @@ describe('modelsService', () => {
                 Observable.of({
                     ...EMPTY_MODELS_DATA,
                     primaryModels: [model1chat],
-                    preferences:
-                        modelsService.storage?.getModelPreferences()[currentAuthStatus().endpoint]!,
+                    preferences: storage?.getModelPreferences()[currentAuthStatus().endpoint]!,
                 })
             )
             vi.spyOn(modelsService, 'models', 'get').mockReturnValue([model1chat])
@@ -410,4 +411,83 @@ describe('modelsService', () => {
             expect(modelsService.models).toContain(uncategorizedModel)
         })
     })
+
+    describe('A/B test for default edit model', () => {
+        let modelsService: ModelsService
+        let storage: TestLocalStorageForModelPreferences
+        const gpt4oMiniModel = createModel({
+            id: 'gpt-4o-mini',
+            modelRef: 'openai::unknown::gpt-4o-mini',
+            usage: [ModelUsage.Edit],
+        })
+        const otherEditModel = createModel({
+            id: 'other-edit-model',
+            modelRef: 'other::unknown::other-edit-model',
+            usage: [ModelUsage.Edit],
+        })
+
+        beforeEach(() => {
+            mockAuthStatus(freeUserAuthStatus)
+            vi.spyOn(featureFlagProvider, 'evaluatedFeatureFlag').mockReturnValue(Observable.of(true))
+            vi.spyOn(userProductSubscriptionModule, 'userProductSubscription', 'get').mockReturnValue(
+                Observable.of(codyProSub)
+            )
+
+            storage = new TestLocalStorageForModelPreferences()
+        })
+
+        afterEach(() => {
+            vi.resetAllMocks()
+        })
+
+        it('sets gpt-4o-mini as default edit model for enrolled users in A/B test', async () => {
+            modelsService = new ModelsService(
+                Observable.of({
+                    primaryModels: [otherEditModel, gpt4oMiniModel],
+                    localModels: [],
+                    preferences: {
+                        defaults: { edit: otherEditModel.id },
+                        selected: { edit: otherEditModel.id },
+                    },
+                })
+            )
+            modelsService.setStorage(storage)
+
+            const defaultEditModelId = await firstValueFrom(modelsService.getDefaultEditModel())
+            expect(defaultEditModelId).toBe('gpt-4o-mini')
+
+            const prefsAfter = storage.getModelPreferences()
+            const selectedEditModel = prefsAfter[freeUserAuthStatus.endpoint]?.selected?.edit
+            expect(selectedEditModel).toBe('gpt-4o-mini')
+        })
+
+        it('does not overwrite user preferences if already enrolled', async () => {
+            storage.setModelPreferences({
+                [AUTH_STATUS_FIXTURE_AUTHED.endpoint]: {
+                    defaults: {},
+                    selected: { edit: 'other-edit-model' },
+                },
+            })
+            storage.getEnrollmentHistory(FeatureFlag.CodyEditDefaultToGpt4oMini)
+
+            modelsService = new ModelsService(
+                Observable.of({
+                    primaryModels: [otherEditModel, gpt4oMiniModel],
+                    localModels: [],
+                    preferences: {
+                        defaults: { edit: otherEditModel.id },
+                        selected: { edit: otherEditModel.id },
+                    },
+                })
+            )
+            modelsService.setStorage(storage)
+
+            const defaultEditModelId = await firstValueFrom(modelsService.getDefaultEditModel())
+            expect(defaultEditModelId).toBe('other-edit-model')
+
+            const prefsAfter = storage.getModelPreferences()
+            const selectedEditModel = prefsAfter[freeUserAuthStatus.endpoint]?.selected?.edit
+            expect(selectedEditModel).toBe('other-edit-model')
+        })
+    })
 })

lib/shared/src/models/modelsService.ts

@@ -1,9 +1,11 @@
 import { type Observable, map } from 'observable-fns'
+
 import { authStatus, currentAuthStatus } from '../auth/authStatus'
 import { mockAuthStatus } from '../auth/authStatus'
 import { type AuthStatus, isCodyProUser, isEnterpriseUser } from '../auth/types'
 import { AUTH_STATUS_FIXTURE_AUTHED_DOTCOM } from '../auth/types'
 import { type PickResolvedConfiguration, resolvedConfig } from '../configuration/resolver'
+import { FeatureFlag, featureFlagProvider } from '../experimentation/FeatureFlagProvider'
 import { logDebug } from '../logger'
 import {
     type StoredLastValue,
@@ -14,17 +16,14 @@ import {
     storeLastValue,
     tap,
 } from '../misc/observable'
-import {
-    firstResultFromOperation,
-    pendingOperation,
-    skipPendingOperation,
-} from '../misc/observableOperation'
+import { firstResultFromOperation, pendingOperation } from '../misc/observableOperation'
 import { ClientConfigSingleton } from '../sourcegraph-api/clientConfig'
 import {
     type UserProductSubscription,
     userProductSubscription,
 } from '../sourcegraph-api/userProductSubscription'
 import { CHAT_INPUT_TOKEN_BUDGET, CHAT_OUTPUT_TOKEN_BUDGET } from '../token/constants'
+
 import { configOverwrites } from './configOverwrites'
 import { type Model, type ServerModel, modelTier } from './model'
 import { syncModels } from './sync'
@@ -223,6 +222,7 @@ const EMPTY_MODELS_DATA: ModelsData = {
 }
 
 export interface LocalStorageForModelPreferences {
+    getEnrollmentHistory(featureName: string): boolean
     getModelPreferences(): DefaultsAndUserPreferencesByEndpoint
     setModelPreferences(preferences: DefaultsAndUserPreferencesByEndpoint): Promise<void>
 }
@@ -244,10 +244,10 @@ export interface ModelAvailabilityStatus {
  *      from this type.)
  */
 export class ModelsService {
-    public storage: LocalStorageForModelPreferences | undefined
+    private storage: LocalStorageForModelPreferences | undefined
 
     private storedValue: StoredLastValue<ModelsData>
-    private syncPreferencesSubscription: Unsubscribable
+    private syncPreferencesSubscription?: Unsubscribable
 
     constructor(testing__mockModelsChanges?: ModelsService['modelsChanges']) {
         if (testing__mockModelsChanges) {
@@ -256,27 +256,56 @@ export class ModelsService {
         this.storedValue = storeLastValue(
             this.modelsChanges.pipe(map(data => (data === pendingOperation ? EMPTY_MODELS_DATA : data)))
         )
+    }
+
+    public setStorage(storage: LocalStorageForModelPreferences): void {
+        this.storage = storage
 
-        this.syncPreferencesSubscription = this.modelsChanges
+        this.syncPreferencesSubscription = combineLatest(
+            this.modelsChanges,
+            featureFlagProvider.evaluatedFeatureFlag(FeatureFlag.CodyEditDefaultToGpt4oMini)
+        )
             .pipe(
-                skipPendingOperation(),
-                tap(data => {
-                    if (this.storage) {
-                        const allSitePrefs = this.storage.getModelPreferences()
-                        const updated: DefaultsAndUserPreferencesByEndpoint = {
-                            ...allSitePrefs,
-                            [currentAuthStatus().endpoint]: data.preferences,
-                        }
-                        this.storage?.setModelPreferences(updated)
+                tap(([data, shouldEditDefaultToGpt4oMini]) => {
+                    if (data === pendingOperation) {
+                        return
+                    }
+
+                    // Ensures we only change user preferences once
+                    // when they join the A/B test.
+                    const isEnrolled = this.storage?.getEnrollmentHistory(
+                        FeatureFlag.CodyEditDefaultToGpt4oMini
+                    )
+
+                    // Ensures that we have the gpt-4o-mini model
+                    // we want to default to in this A/B test.
+                    const gpt4oMini = data.primaryModels.find(
+                        model => model?.modelRef?.modelId === 'gpt-4o-mini'
+                    )
+
+                    const allSitePrefs = this.storage?.getModelPreferences()
+                    const currentAccountPrefs = { ...data.preferences }
+
+                    if (!isEnrolled && shouldEditDefaultToGpt4oMini && gpt4oMini) {
+                        // For users enrolled in the A/B test, we'll default
+                        // to the gpt-4-mini model when using the Edit command.
+                        // They still can switch back to other models if they want.
+                        currentAccountPrefs.selected.edit = gpt4oMini.id
                     }
+
+                    const updated: DefaultsAndUserPreferencesByEndpoint = {
+                        ...allSitePrefs,
+                        [currentAuthStatus().endpoint]: currentAccountPrefs,
+                    }
+                    this.storage?.setModelPreferences(updated)
                 })
             )
             .subscribe({})
     }
 
     public dispose(): void {
         this.storedValue.subscription.unsubscribe()
-        this.syncPreferencesSubscription.unsubscribe()
+        this.syncPreferencesSubscription?.unsubscribe()
     }
 
     /**
@@ -313,13 +342,13 @@ export class ModelsService {
 
     private getModelsByType(usage: ModelUsage): Observable<Model[] | typeof pendingOperation> {
         return this.modelsChanges.pipe(
-            map(models =>
-                models === pendingOperation
+            map(models => {
+                return models === pendingOperation
                     ? pendingOperation
                     : [...models.primaryModels, ...models.localModels].filter(model =>
                           model.usage.includes(usage)
                       )
-            ),
+            }),
             distinctUntilChanged()
         )
     }
@@ -575,6 +604,7 @@ interface MockModelsServiceResult {
 }
 
 export class TestLocalStorageForModelPreferences implements LocalStorageForModelPreferences {
+    private isEnrolled = false
     constructor(public data: DefaultsAndUserPreferencesByEndpoint | null = null) {}
 
     getModelPreferences(): DefaultsAndUserPreferencesByEndpoint {
@@ -584,6 +614,14 @@ export class TestLocalStorageForModelPreferences implements LocalStorageForModel
     async setModelPreferences(preferences: DefaultsAndUserPreferencesByEndpoint): Promise<void> {
         this.data = preferences
     }
+
+    getEnrollmentHistory(_featureName: string): boolean {
+        if (!this.isEnrolled) {
+            this.isEnrolled = true
+            return false
+        }
+        return true
+    }
 }
 
 export function mockModelsService({
@@ -595,7 +633,7 @@ export function mockModelsService({
     modelsService?: ModelsService
     storage?: TestLocalStorageForModelPreferences
 }): MockModelsServiceResult {
-    modelsService.storage = storage
+    modelsService.setStorage(storage)
     mockAuthStatus(authStatus)
     return { storage, modelsService }
 }

lib/shared/src/models/sync.test.ts

@@ -208,7 +208,7 @@ describe('server sent models', async () => {
         }).pipe(skipPendingOperation())
     )
     const storage = new TestLocalStorageForModelPreferences()
-    modelsService.storage = storage
+    modelsService.setStorage(storage)
     mockAuthStatus(AUTH_STATUS_FIXTURE_AUTHED)
     vi.spyOn(modelsService, 'modelsChanges', 'get').mockReturnValue(Observable.of(result))
 
@@ -504,7 +504,7 @@ describe('syncModels', () => {
         )
 
         const storage = new TestLocalStorageForModelPreferences()
-        modelsService.storage = storage
+        modelsService.setStorage(storage)
         mockAuthStatus(AUTH_STATUS_FIXTURE_AUTHED)
         expect(storage.data?.[AUTH_STATUS_FIXTURE_AUTHED.endpoint]!.selected.chat).toBe(undefined)
         vi.spyOn(modelsService, 'modelsChanges', 'get').mockReturnValue(Observable.of(result))

vscode/src/edit/manager.ts

@@ -122,7 +122,8 @@ export class EditManager implements vscode.Disposable {
         }
 
         // Set default edit configuration, if not provided
-        // It is possible that these values may be overriden later, e.g. if the user changes them in the edit input.
+        // It is possible that these values may be overridden later,
+        // e.g. if the user changes them in the edit input.
         const range = getEditLineSelection(document, proposedRange)
         const model =
             configuration.model || (await firstResultFromOperation(modelsService.getDefaultEditModel()))

vscode/src/edit/utils/edit-intent.ts

@@ -20,7 +20,7 @@ export function getEditIntent(
     proposedIntent?: EditIntent
 ): EditIntent {
     if (proposedIntent !== undefined && proposedIntent !== 'add') {
-        // Return provided intent that should not be overriden
+        // Return provided intent that should not be overridden
         return proposedIntent
     }
 

vscode/src/edit/utils/edit-models.ts

@@ -1,7 +1,7 @@
 import { type AuthStatus, type EditModel, isDotCom } from '@sourcegraph/cody-shared'
 import type { EditIntent } from '../types'
 
-export function getOverridenModelForIntent(
+export function getOverriddenModelForIntent(
     intent: EditIntent,
     currentModel: EditModel,
     authStatus: AuthStatus
@@ -21,7 +21,7 @@ export function getOverridenModelForIntent(
             // Issue: https://github.com/sourcegraph/cody/issues/3512
             return 'anthropic/claude-3-5-sonnet-20240620'
         case 'doc':
-            // Doc is a case where we can sacrifice LLM performnace for improved latency and get comparable results.
+            // Doc is a case where we can sacrifice LLM performance for improved latency and get comparable results.
             return 'anthropic/claude-3-haiku-20240307'
         case 'test':
         case 'add':

vscode/src/main.ts

@@ -325,7 +325,7 @@ async function initializeSingletons(
 ): Promise<void> {
     commandControllerInit(platform.createCommandsProvider?.(), platform.extensionClient.capabilities)
 
-    modelsService.storage = localStorage
+    modelsService.setStorage(localStorage)
 
     if (platform.otherInitialization) {
         disposables.push(platform.otherInitialization())