Reset feature flags when switching through logins (#2182)

Fixes #2181 Fixes #2154 This PR attempts to clean up the auth flow and makes it so that when the user account is switched, we never accidentely use feature flags of a wrong server URL and also ensures we reconfigure Autocomplete so it works. It does this by removing the relationship between the `AuthProvider` and a specific webview. This is necessary because: - With features like Autocomplete, we have things that depend on the AuthProvider unrelated on wether a web view is enabled or not - We have more than one web view right now and this doesn't make sense Instead, we make it so that the web views subscribe to the auth state like every other service. In untangling this, the biggest issue was undocumented dependencies into how events were passed. What I noticed was that the `ContextProvider` plays a major role in initializing the web views. This is something we should clean up in a follow up PR (IMO the web views should subscribe to the ContextProvider and not the other way around). I also noticed that the `cody.auth.sync` action isn't really necessary as we can subscribe to the `AuthProvider` via `addChangeListener` anyways (which most of the services did). Also: Apparently there were two ways on when we reconfigure the `Configuration` object. I now made sure both ways at least reconfigure the same components by extracting a `onConfigurationChange` function inside `main.ts`. Two unrelated fixes: - The code to update the upgrade state when the editor was focused would also run on enterprise instances causing a bunch of errors in the console. Fixed. - Added the remote URL to the completion loggers we can easily see where the request is sent to. ⚠️ This PR has a potential to break areas depending on auth, so please take some time to test it. I've been testing it as thorough as I could though! ## Test plan The canonical test case was: - Be signed in to dotcom and have the starcoder-hybrid feature flag enable - Now switch account to an enterprise instance - Ensure that requests are made to the enterprise instance and that we do not use starcoder model strings https://github.com/sourcegraph/cody/assets/458591/ab5a19a6-a906-42e6-8f76-ae61b500e30c
sourcegraph · Dec 8, 2023 · b46bb30 · b46bb30
1 parent 18b56f5
commit b46bb30
Show file tree

Hide file tree

Showing 14 changed files with 67 additions and 73 deletions.
diff --git a/agent/src/agent.ts b/agent/src/agent.ts
@@ -530,7 +530,6 @@ export class Agent extends MessageHandler {
 
     private async reloadAuth(): Promise<void> {
         await vscode_shim.commands.executeCommand('agent.auth.reload')
-        await vscode_shim.commands.executeCommand('cody.auth.sync')
     }
 
     /**

diff --git a/lib/shared/src/experimentation/FeatureFlagProvider.ts b/lib/shared/src/experimentation/FeatureFlagProvider.ts
@@ -28,43 +28,51 @@ export enum FeatureFlag {
 const ONE_HOUR = 60 * 60 * 1000
 
 export class FeatureFlagProvider {
-    private featureFlags: Record<string, boolean> = {}
+    // The first key maps to the endpoint so that we do never cache the wrong flag for different
+    // endpoints
+    private featureFlags: Record<string, Record<string, boolean>> = {}
     private lastUpdated = 0
 
     constructor(private apiClient: SourcegraphGraphQLAPIClient) {}
 
-    private getFromCache(flagName: FeatureFlag): boolean | undefined {
+    private getFromCache(flagName: FeatureFlag, endpoint: string): boolean | undefined {
         const now = Date.now()
         if (now - this.lastUpdated > ONE_HOUR) {
             // Cache expired, refresh
             void this.refreshFeatureFlags()
         }
 
-        return this.featureFlags[flagName]
+        return this.featureFlags[endpoint]?.[flagName]
     }
 
     public async evaluateFeatureFlag(flagName: FeatureFlag): Promise<boolean> {
+        const endpoint = this.apiClient.endpoint
         if (process.env.BENCHMARK_DISABLE_FEATURE_FLAGS) {
             return false
         }
 
-        const cachedValue = this.getFromCache(flagName)
+        const cachedValue = this.getFromCache(flagName, endpoint)
         if (cachedValue !== undefined) {
             return cachedValue
         }
 
         const value = await this.apiClient.evaluateFeatureFlag(flagName)
-        this.featureFlags[flagName] = value === null || isError(value) ? false : value
-        return this.featureFlags[flagName]
+        if (!this.featureFlags[endpoint]) {
+            this.featureFlags[endpoint] = {}
+        }
+        this.featureFlags[endpoint][flagName] = value === null || isError(value) ? false : value
+        return this.featureFlags[endpoint][flagName]
     }
 
     public syncAuthStatus(): void {
+        this.featureFlags = {}
         void this.refreshFeatureFlags()
     }
 
     private async refreshFeatureFlags(): Promise<void> {
+        const endpoint = this.apiClient.endpoint
         const data = await this.apiClient.getEvaluatedFeatureFlags()
-        this.featureFlags = isError(data) ? {} : data
+        this.featureFlags[endpoint] = isError(data) ? {} : data
         this.lastUpdated = Date.now()
     }
 }

diff --git a/lib/shared/src/sourcegraph-api/completions/client.ts b/lib/shared/src/sourcegraph-api/completions/client.ts
@@ -3,7 +3,10 @@ import { ConfigurationWithAccessToken } from '../../configuration'
 import { CompletionCallbacks, CompletionParameters, CompletionResponse, Event } from './types'
 
 export interface CompletionLogger {
-    startCompletion(params: CompletionParameters | {}):
+    startCompletion(
+        params: CompletionParameters | {},
+        endpoint: string
+    ):
         | undefined
         | {
               onError: (error: string) => void

diff --git a/lib/shared/src/sourcegraph-api/completions/nodeClient.ts b/lib/shared/src/sourcegraph-api/completions/nodeClient.ts
@@ -12,7 +12,7 @@ import { CompletionCallbacks, CompletionParameters } from './types'
 
 export class SourcegraphNodeCompletionsClient extends SourcegraphCompletionsClient {
     public stream(params: CompletionParameters, cb: CompletionCallbacks): () => void {
-        const log = this.logger?.startCompletion(params)
+        const log = this.logger?.startCompletion(params, this.completionsEndpoint)
 
         const abortController = new AbortController()
         const abortSignal = abortController.signal

diff --git a/lib/shared/src/sourcegraph-api/graphql/queries.ts b/lib/shared/src/sourcegraph-api/graphql/queries.ts
@@ -28,14 +28,6 @@ query SiteGraphQLFields {
     }
 }`
 
-export const CURRENT_USER_ID_AND_VERIFIED_EMAIL_QUERY = `
-query CurrentUser {
-    currentUser {
-        id
-        hasVerifiedEmail
-    }
-}`
-
 export const CURRENT_USER_ID_AND_VERIFIED_EMAIL_AND_CODY_PRO_QUERY = `
 query CurrentUser {
     currentUser {

diff --git a/vscode/CHANGELOG.md b/vscode/CHANGELOG.md
@@ -8,6 +8,8 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 
 ### Fixed
 
+- Autocomplete: Fixes an issue where changing user accounts caused some configuration issues. [pull/2182](https://github.com/sourcegraph/cody/pull/2182)
+- Fixes an issue where focusing the VS Code extension window caused unexpected errors when connected to an Enterprise instance. [pull/2182](https://github.com/sourcegraph/cody/pull/2182)
 - Chat: You can @-mention files on Windows without generating an error. [pull/2197](https://github.com/sourcegraph/cody/pull/2197)
 
 ### Changed

diff --git a/vscode/src/chat/chat-view/ChatManager.ts b/vscode/src/chat/chat-view/ChatManager.ts
@@ -226,7 +226,6 @@ export class ChatManager implements vscode.Disposable {
 
     private disposeChatPanelsManager(): void {
         this.options.contextProvider.webview = this.sidebarChat.webview
-        this.options.authProvider.webview = this.sidebarChat.webview
         this.chatPanelsManager?.dispose()
     }
 

diff --git a/vscode/src/chat/chat-view/ChatPanelProvider.ts b/vscode/src/chat/chat-view/ChatPanelProvider.ts
@@ -61,9 +61,6 @@ export class ChatPanelProvider extends MessageProvider {
     private async onDidReceiveMessage(message: WebviewMessage): Promise<void> {
         switch (message.command) {
             case 'ready':
-                // The web view is ready to receive events. We need to make sure that it has an up
-                // to date config, even if it was already published
-                await this.authProvider.announceNewAuthStatus()
                 this.handleWebviewContext()
                 break
             case 'initialized':
@@ -448,7 +445,6 @@ export class ChatPanelProvider extends MessageProvider {
         // should not overwrite the context provider's webview, or it
         // will break the previous chat panel.
         this.contextProvider.webview = panel.webview
-        this.authProvider.webview = panel.webview
         this.postEnhancedContextStatusToWebview()
 
         // Dispose panel when the panel is closed

diff --git a/vscode/src/chat/chat-view/SidebarChatProvider.ts b/vscode/src/chat/chat-view/SidebarChatProvider.ts
@@ -46,9 +46,7 @@ export class SidebarChatProvider extends MessageProvider implements vscode.Webvi
     private async onDidReceiveMessage(message: WebviewMessage): Promise<void> {
         switch (message.command) {
             case 'ready':
-                // The web view is ready to receive events. We need to make sure that it has an up
-                // to date config, even if it was already published
-                await this.authProvider.announceNewAuthStatus()
+                await this.contextProvider.syncAuthStatus()
                 break
             case 'initialized':
                 logDebug('SidebarChatProvider:onDidReceiveMessage', 'initialized')
@@ -363,7 +361,6 @@ export class SidebarChatProvider extends MessageProvider implements vscode.Webvi
         _token: vscode.CancellationToken
     ): Promise<void> {
         this.webview = webviewView.webview
-        this.authProvider.webview = webviewView.webview
         this.contextProvider.webview = webviewView.webview
 
         const webviewPath = vscode.Uri.joinPath(this.extensionUri, 'dist', 'webviews')

diff --git a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
@@ -299,9 +299,6 @@ export class SimpleChatPanelProvider implements vscode.Disposable, IChatPanelPro
     private async onDidReceiveMessage(message: WebviewMessage): Promise<void> {
         switch (message.command) {
             case 'ready':
-                // The web view is ready to receive events. We need to make sure that it has an up
-                // to date config, even if it was already published
-                await this.authProvider.announceNewAuthStatus()
                 await this.postViewConfig()
                 break
             case 'initialized':

diff --git a/vscode/src/completions/client.ts b/vscode/src/completions/client.ts
@@ -60,7 +60,8 @@ export function createClient(config: CompletionsClientConfig, logger?: Completio
         onPartialResponse?: (incompleteResponse: CompletionResponse) => void,
         signal?: AbortSignal
     ): Promise<CompletionResponse> {
-        const log = logger?.startCompletion(params)
+        const url = getCodeCompletionsEndpoint()
+        const log = logger?.startCompletion(params, url)
 
         const tracingFlagEnabled = await featureFlagProvider.evaluateFeatureFlag(FeatureFlag.CodyAutocompleteTracing)
 
@@ -86,7 +87,6 @@ export function createClient(config: CompletionsClientConfig, logger?: Completio
         const isNode = typeof process !== 'undefined'
         const enableStreaming = !!isNode
 
-        const url = getCodeCompletionsEndpoint()
         const response: Response = await fetch(url, {
             method: 'POST',
             body: JSON.stringify({

diff --git a/vscode/src/log.ts b/vscode/src/log.ts
@@ -90,7 +90,7 @@ function log(level: 'debug' | 'error', filterLabel: string, text: string, ...arg
 }
 
 export const logger: CompletionLogger = {
-    startCompletion(params: CompletionParameters | {}) {
+    startCompletion(params: CompletionParameters | {}, endpoint: string) {
         const workspaceConfig = vscode.workspace.getConfiguration()
         const config = getConfiguration(workspaceConfig)
 
@@ -112,6 +112,7 @@ export const logger: CompletionLogger = {
                 'CompletionLogger:onError',
                 JSON.stringify({
                     type,
+                    endpoint,
                     status: 'error',
                     duration: Date.now() - start,
                     err,
@@ -130,6 +131,7 @@ export const logger: CompletionLogger = {
                 'CompletionLogger:onComplete',
                 JSON.stringify({
                     type,
+                    endpoint,
                     status: 'success',
                     duration: Date.now() - start,
                 }),

diff --git a/vscode/src/main.ts b/vscode/src/main.ts
@@ -6,6 +6,7 @@ import { ConfigurationWithAccessToken } from '@sourcegraph/cody-shared/src/confi
 import { FixupIntent } from '@sourcegraph/cody-shared/src/editor'
 import { FeatureFlag, featureFlagProvider } from '@sourcegraph/cody-shared/src/experimentation/FeatureFlagProvider'
 import { newPromptMixin, PromptMixin } from '@sourcegraph/cody-shared/src/prompt/prompt-mixin'
+import { isDotCom } from '@sourcegraph/cody-shared/src/sourcegraph-api/environments'
 import { graphqlClient } from '@sourcegraph/cody-shared/src/sourcegraph-api/graphql'
 
 import { CachedRemoteEmbeddingsClient } from './chat/CachedRemoteEmbeddingsClient'
@@ -184,6 +185,25 @@ const register = async (
 
     disposables.push(new CodeActionProvider({ contextProvider }))
 
+    let oldConfig = JSON.stringify(initialConfig)
+    function onConfigurationChange(newConfig: ConfigurationWithAccessToken): void {
+        if (oldConfig === JSON.stringify(newConfig)) {
+            return
+        }
+        oldConfig = JSON.stringify(newConfig)
+
+        featureFlagProvider.syncAuthStatus()
+        graphqlClient.onConfigurationChange(newConfig)
+        contextProvider.onConfigurationChange(newConfig)
+        externalServicesOnDidConfigurationChange(newConfig)
+        void configureEventsInfra(newConfig, isExtensionModeDevOrTest)
+        platform.onConfigurationChange?.(newConfig)
+        symfRunner?.setSourcegraphAuth(newConfig.serverEndpoint, newConfig.accessToken)
+        void localEmbeddings?.setAccessToken(newConfig.serverEndpoint, newConfig.accessToken)
+        embeddingsClient.updateConfiguration(newConfig)
+        setupAutocomplete()
+    }
+
     // Register tree views
     disposables.push(
         chatManager,
@@ -193,8 +213,7 @@ const register = async (
         // Update external services when configurationChangeEvent is fired by chatProvider
         contextProvider.configurationChangeEvent.event(async () => {
             const newConfig = await getFullConfig()
-            externalServicesOnDidConfigurationChange(newConfig)
-            await configureEventsInfra(newConfig, isExtensionModeDevOrTest)
+            onConfigurationChange(newConfig)
         })
     )
 
@@ -210,8 +229,13 @@ const register = async (
     }
 
     // Adds a change listener to the auth provider that syncs the auth status
-    authProvider.addChangeListener((authStatus: AuthStatus) => {
-        void chatManager.syncAuthStatus(authStatus)
+    authProvider.addChangeListener(async (authStatus: AuthStatus) => {
+        // Update context provider first since it will also update the configuration
+        await contextProvider.syncAuthStatus()
+
+        featureFlagProvider.syncAuthStatus()
+        await chatManager.syncAuthStatus(authStatus)
+
         if (symfRunner && authStatus.isLoggedIn) {
             getAccessToken()
                 .then(token => {
@@ -222,6 +246,8 @@ const register = async (
         } else {
             symfRunner?.setSourcegraphAuth(null, null)
         }
+
+        setupAutocomplete()
     })
     // Sync initial auth status
     await chatManager.syncAuthStatus(authProvider.getAuthStatus())
@@ -286,13 +312,6 @@ const register = async (
         vscode.commands.registerCommand('cody.auth.signin', () => authProvider.signinMenu()),
         vscode.commands.registerCommand('cody.auth.signout', () => authProvider.signoutMenu()),
         vscode.commands.registerCommand('cody.auth.support', () => showFeedbackSupportQuickPick()),
-        vscode.commands.registerCommand('cody.auth.sync', () => {
-            const result = contextProvider.syncAuthStatus()
-            void featureFlagProvider.syncAuthStatus()
-            // Important that we return a promise here to allow `AuthProvider`
-            // to `await` on the auth config changes to propagate.
-            return result
-        }),
         // Commands
         vscode.commands.registerCommand('cody.chat.restart', async () => {
             const confirmation = await vscode.window.showWarningMessage(
@@ -465,7 +484,8 @@ const register = async (
     updateAuthStatusBarIndicator()
 
     vscode.window.onDidChangeWindowState(async ws => {
-        if (ws.focused) {
+        const endpoint = authProvider.getAuthStatus().endpoint
+        if (ws.focused && endpoint && isDotCom(endpoint)) {
             const res = await graphqlClient.getCurrentUserIdAndVerifiedEmailAndCodyPro()
             if (res instanceof Error) {
                 console.error(res)
@@ -484,7 +504,7 @@ const register = async (
     let completionsProvider: vscode.Disposable | null = null
     let setupAutocompleteQueue = Promise.resolve() // Create a promise chain to avoid parallel execution
     disposables.push({ dispose: () => completionsProvider?.dispose() })
-    const setupAutocomplete = (): void => {
+    function setupAutocomplete(): void {
         setupAutocompleteQueue = setupAutocompleteQueue
             .then(async () => {
                 const config = getConfiguration(vscode.workspace.getConfiguration())
@@ -520,15 +540,6 @@ const register = async (
             })
     }
 
-    // Reload autocomplete if either the configuration changes or the auth status is updated
-    vscode.workspace.onDidChangeConfiguration(event => {
-        if (event.affectsConfiguration('cody.autocomplete')) {
-            setupAutocomplete()
-        }
-    })
-    authProvider.addChangeListener(() => {
-        setupAutocomplete()
-    })
     setupAutocomplete()
 
     if (initialConfig.experimentalGuardrails) {
@@ -559,16 +570,7 @@ const register = async (
 
     return {
         disposable: vscode.Disposable.from(...disposables),
-        onConfigurationChange: newConfig => {
-            graphqlClient.onConfigurationChange(newConfig)
-            contextProvider.onConfigurationChange(newConfig)
-            externalServicesOnDidConfigurationChange(newConfig)
-            void configureEventsInfra(newConfig, isExtensionModeDevOrTest)
-            platform.onConfigurationChange?.(newConfig)
-            symfRunner?.setSourcegraphAuth(newConfig.serverEndpoint, newConfig.accessToken)
-            void localEmbeddings?.setAccessToken(newConfig.serverEndpoint, newConfig.accessToken)
-            embeddingsClient.updateConfiguration(newConfig)
-        },
+        onConfigurationChange,
     }
 }
 

diff --git a/vscode/src/services/AuthProvider.ts b/vscode/src/services/AuthProvider.ts
@@ -6,7 +6,6 @@ import { SourcegraphGraphQLAPIClient } from '@sourcegraph/cody-shared/src/source
 import { isError } from '@sourcegraph/cody-shared/src/utils'
 
 import { CodyChatPanelViewType } from '../chat/chat-view/ChatManager'
-import { SidebarChatWebview } from '../chat/chat-view/SidebarChatProvider'
 import {
     AuthStatus,
     defaultAuthStatus,
@@ -33,7 +32,6 @@ export class AuthProvider {
     private client: SourcegraphGraphQLAPIClient | null = null
 
     private authStatus: AuthStatus = defaultAuthStatus
-    public webview?: SidebarChatWebview
     private listeners: Set<Listener> = new Set()
 
     constructor(
@@ -237,7 +235,7 @@ export class AuthProvider {
         const isLoggedIn = isAuthed(authStatus)
         authStatus.isLoggedIn = isLoggedIn
         await this.storeAuthInfo(endpoint, token)
-        await this.syncAuthStatus(authStatus)
+        this.syncAuthStatus(authStatus)
         await vscode.commands.executeCommand('setContext', 'cody.activated', isLoggedIn)
         return { authStatus, isLoggedIn }
     }
@@ -249,23 +247,22 @@ export class AuthProvider {
     }
 
     // Set auth status and share it with chatview
-    private async syncAuthStatus(authStatus: AuthStatus): Promise<void> {
+    private syncAuthStatus(authStatus: AuthStatus): void {
         if (this.authStatus === authStatus) {
             return
         }
         this.authStatus = authStatus
-        await this.announceNewAuthStatus()
+        this.announceNewAuthStatus()
     }
 
-    public async announceNewAuthStatus(): Promise<void> {
-        if (this.authStatus.endpoint === 'init' || !this.webview) {
+    public announceNewAuthStatus(): void {
+        if (this.authStatus.endpoint === 'init') {
             return
         }
         const authStatus = this.getAuthStatus()
         for (const listener of this.listeners) {
             listener(authStatus)
         }
-        await vscode.commands.executeCommand('cody.auth.sync')
     }
 
     // Register URI Handler (vscode://sourcegraph.cody-ai) for resolving token