add user eligibility condition for autoedits

sourcegraph · Jan 6, 2025 · 5eebf42 · 5eebf42
1 parent 81b0302
commit 5eebf42
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 54 deletions.
diff --git a/vscode/src/autoedits/create-autoedits-provider.ts b/vscode/src/autoedits/create-autoedits-provider.ts
@@ -2,28 +2,34 @@ import { type Observable, map } from 'observable-fns'
 import * as vscode from 'vscode'
 
 import {
-    type AuthenticatedAuthStatus,
+    type AuthStatus,
     type ChatClient,
     NEVER,
     type PickResolvedConfiguration,
-    type UnauthenticatedAuthStatus,
+    type UserProductSubscription,
+    combineLatest,
     createDisposables,
+    currentUserProductSubscription,
     promiseFactoryToObservable,
     skipPendingOperation,
 } from '@sourcegraph/cody-shared'
+import { isFreeUser } from '@sourcegraph/cody-shared/src/auth/types'
+import { isRunningInsideAgent } from '../jsonrpc/isRunningInsideAgent'
 import { AutoeditsProvider } from './autoedits-provider'
 import { autoeditsOutputChannelLogger } from './output-channel-logger'
 
 interface AutoeditsItemProviderArgs {
     config: PickResolvedConfiguration<{ configuration: true }>
-    authStatus: UnauthenticatedAuthStatus | Pick<AuthenticatedAuthStatus, 'authenticated' | 'endpoint'>
+    authStatus: AuthStatus
     chatClient: ChatClient
+    autoeditsFeatureFlagEnabled: boolean
 }
 
 export function createAutoEditsProvider({
     config: { configuration },
     authStatus,
     chatClient,
+    autoeditsFeatureFlagEnabled,
 }: AutoeditsItemProviderArgs): Observable<void> {
     if (!configuration.experimentalAutoeditsEnabled) {
         return NEVER
@@ -36,11 +42,24 @@ export function createAutoEditsProvider({
         return NEVER
     }
 
-    return promiseFactoryToObservable(async () => {
-        return await getAutoeditsProviderDocumentFilters()
-    }).pipe(
+    return combineLatest(
+        promiseFactoryToObservable(async () => await currentUserProductSubscription()),
+        promiseFactoryToObservable(async () => await getAutoeditsProviderDocumentFilters())
+    ).pipe(
         skipPendingOperation(),
-        createDisposables(documentFilters => {
+        createDisposables(([userProductSubscription, documentFilters]) => {
+            if (
+                !isUserEligibleForAutoeditsFeature(
+                    autoeditsFeatureFlagEnabled,
+                    authStatus,
+                    userProductSubscription
+                )
+            ) {
+                throw new Error(
+                    'User is not eligible for auto-edits. Can not initialize auto-edits provider.'
+                )
+            }
+
             const provider = new AutoeditsProvider(chatClient)
             return [
                 vscode.commands.registerCommand('cody.command.autoedits-manual-trigger', async () => {
@@ -58,3 +77,20 @@ export function createAutoEditsProvider({
 export async function getAutoeditsProviderDocumentFilters(): Promise<vscode.DocumentFilter[]> {
     return [{ scheme: 'file', language: '*' }, { notebookType: '*' }]
 }
+
+export function isUserEligibleForAutoeditsFeature(
+    autoeditsFeatureFlagEnabled: boolean,
+    authStatus: AuthStatus,
+    productSubsubscription: UserProductSubscription | null
+): boolean {
+    // Editors other than vscode are not eligible for auto-edits
+    if (isRunningInsideAgent()) {
+        return false
+    }
+    // Free users are not eligible for auto-edits
+    if (isFreeUser(authStatus, productSubsubscription)) {
+        return false
+    }
+    // Users with autoedits feature flag enabled are eligible for auto-edits
+    return autoeditsFeatureFlagEnabled
+}
diff --git a/vscode/src/main.ts b/vscode/src/main.ts
@@ -1,7 +1,6 @@
 import * as vscode from 'vscode'
 
 import {
-    type AuthStatus,
     type ChatClient,
     ClientConfigSingleton,
     type ConfigurationInput,
@@ -25,7 +24,6 @@ import {
     fromVSCodeEvent,
     graphqlClient,
     isDotCom,
-    isS2,
     modelsService,
     resolvedConfig,
     setClientCapabilities,
@@ -714,14 +712,12 @@ function registerAutoEdits(chatClient: ChatClient, disposables: vscode.Disposabl
                             isEqual(a[2], b[2])
                         )
                     }),
-                    switchMap(([config, authStatus, autoeditEnabled]) => {
-                        if (!shouldEnableExperimentalAutoedits(config, autoeditEnabled, authStatus)) {
-                            return NEVER
-                        }
+                    switchMap(([config, authStatus, autoeditsFeatureFlagEnabled]) => {
                         return createAutoEditsProvider({
                             config,
                             authStatus,
                             chatClient,
+                            autoeditsFeatureFlagEnabled,
                         })
                     }),
                     catchError(error => {
@@ -734,23 +730,6 @@ function registerAutoEdits(chatClient: ChatClient, disposables: vscode.Disposabl
     )
 }
 
-function shouldEnableExperimentalAutoedits(
-    config: ResolvedConfiguration,
-    autoeditFeatureFlagEnabled: boolean,
-    authStatus: AuthStatus
-): boolean {
-    // If running inside agent don't enable experimental autoedits
-    if (isRunningInsideAgent()) {
-        return false
-    }
-    // If the config is explicitly set in the vscode settings, use the setting instead of the feature flag.
-    if (config.configuration.experimentalAutoeditsEnabled !== undefined) {
-        return config.configuration.experimentalAutoeditsEnabled
-    }
-    // Feature flag should only control S2, use the flag instead of the config.
-    return autoeditFeatureFlagEnabled && isS2(authStatus)
-}
-
 /**
  * Registers autocomplete functionality.
  */
@@ -771,35 +750,14 @@ function registerAutocomplete(
 
     disposables.push(
         subscriptionDisposable(
-            combineLatest(
-                resolvedConfig,
-                authStatus,
-                featureFlagProvider.evaluatedFeatureFlag(
-                    FeatureFlag.CodyAutoeditExperimentEnabledFeatureFlag
-                )
-            )
+            combineLatest(resolvedConfig, authStatus)
                 .pipe(
                     //TODO(@rnauta -> @sqs): It feels yuk to handle the invalidation outside of
                     //where the state is picked. It's also very tedious
                     distinctUntilChanged((a, b) => {
-                        return (
-                            isEqual(a[0].configuration, b[0].configuration) &&
-                            isEqual(a[1], b[1]) &&
-                            isEqual(a[2], b[2])
-                        )
+                        return isEqual(a[0].configuration, b[0].configuration) && isEqual(a[1], b[1])
                     }),
-                    switchMap(([config, authStatus, autoeditFeatureFlagEnabled]) => {
-                        // If the auto-edit experiment is enabled, we don't need to load the completion provider
-                        if (
-                            shouldEnableExperimentalAutoedits(
-                                config,
-                                autoeditFeatureFlagEnabled,
-                                authStatus
-                            )
-                        ) {
-                            finishLoading()
-                            return NEVER
-                        }
+                    switchMap(([config, authStatus]) => {
                         if (!authStatus.pendingValidation && !statusBarLoader) {
                             statusBarLoader = statusBar.addLoader({
                                 title: 'Completion Provider is starting',