Refactor register and move auth/config propagation into controllers (…

…#4859)

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/AuthStatus.kt

@@ -3,7 +3,7 @@ package com.sourcegraph.cody.protocol_generated;
 
 data class AuthStatus(
   val username: String,
-  val endpoint: String? = null,
+  val endpoint: String,
   val isDotCom: Boolean,
   val isLoggedIn: Boolean,
   val isFireworksTracingEnabled: Boolean,

agent/src/cli/command-bench/command-bench.ts

@@ -314,7 +314,7 @@ export const benchCommand = new commander.Command('bench')
         )
 
         // Required to use `PromptString`.
-        graphqlClient.onConfigurationChange({
+        graphqlClient.setConfig({
             accessToken: options.srcAccessToken,
             serverEndpoint: options.srcEndpoint,
             customHeaders: {},

lib/shared/src/auth/types.ts

@@ -6,7 +6,7 @@ import type { CodyLLMSiteConfiguration } from '../sourcegraph-api/graphql/client
  */
 export interface AuthStatus {
     username: string
-    endpoint: string | null
+    endpoint: string
     isDotCom: boolean
     isLoggedIn: boolean
     /**

lib/shared/src/experimentation/FeatureFlagProvider.test.ts

@@ -26,7 +26,7 @@ describe('FeatureFlagProvider', () => {
         }
 
         const provider = new FeatureFlagProvider(apiClient as unknown as SourcegraphGraphQLAPIClient)
-        await provider.syncAuthStatus()
+        await provider.refresh()
 
         // Wait for the async initialization
         await nextTick()
@@ -64,7 +64,7 @@ describe('FeatureFlagProvider', () => {
             [FeatureFlag.TestFlagDoNotUse]: false,
         })
 
-        await provider.syncAuthStatus()
+        await provider.refresh()
 
         // Wait for the async reload
         await nextTick()
@@ -84,7 +84,7 @@ describe('FeatureFlagProvider', () => {
             }
 
             const provider = new FeatureFlagProvider(apiClient as unknown as SourcegraphGraphQLAPIClient)
-            await provider.syncAuthStatus()
+            await provider.refresh()
 
             // Wait for the async initialization
             await nextTick()

lib/shared/src/experimentation/FeatureFlagProvider.ts

@@ -161,13 +161,13 @@ export class FeatureFlagProvider {
         })
     }
 
-    public async syncAuthStatus(): Promise<void> {
+    public async refresh(): Promise<void> {
         this.exposedFeatureFlags = {}
         this.unexposedFeatureFlags = {}
         await this.refreshFeatureFlags()
     }
 
-    public async refreshFeatureFlags(): Promise<void> {
+    private async refreshFeatureFlags(): Promise<void> {
         return wrapInActiveSpan('FeatureFlagProvider.refreshFeatureFlags', async () => {
             const endpoint = this.apiClient.endpoint
             const data = process.env.DISABLE_FEATURE_FLAGS

lib/shared/src/sourcegraph-api/graphql/client.ts

@@ -510,7 +510,7 @@ export class SourcegraphGraphQLAPIClient {
         this._config = config
     }
 
-    public onConfigurationChange(newConfig: GraphQLAPIClientConfig): void {
+    public setConfig(newConfig: GraphQLAPIClientConfig): void {
         this._config = newConfig
     }
 
@@ -1342,7 +1342,7 @@ export class ClientConfigSingleton {
         return ClientConfigSingleton.instance
     }
 
-    public async syncAuthStatus(authStatus: AuthStatus): Promise<void> {
+    public async setAuthStatus(authStatus: AuthStatus): Promise<void> {
         this.isSignedIn = authStatus.authenticated && authStatus.isLoggedIn
         if (this.isSignedIn) {
             await this.refreshConfig()

vscode/src/chat/chat-view/ChatController.ts

@@ -529,7 +529,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
     // #region top-level view action handlers
     // =======================================================================
 
-    public syncAuthStatus(): void {
+    public setAuthStatus(authStatus: AuthStatus): void {
         // Run this async because this method may be called during initialization
         // and awaiting on this.postMessage may result in a deadlock
         void this.sendConfig()

vscode/src/chat/chat-view/ChatsController.ts

@@ -27,6 +27,7 @@ import type { ExecuteChatArguments } from '../../commands/execute/ask'
 import { getConfiguration } from '../../configuration'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
+import type { AuthProvider } from '../../services/AuthProvider'
 import {
     ChatController,
     type ChatSession,
@@ -63,6 +64,7 @@ export class ChatsController implements vscode.Disposable {
     constructor(
         private options: Options,
         private chatClient: ChatClient,
+        private authProvider: AuthProvider,
         private readonly enterpriseContext: EnterpriseContextFactory,
         private readonly localEmbeddings: LocalEmbeddingsController | null,
         private readonly contextRanking: ContextRankingController | null,
@@ -71,9 +73,14 @@ export class ChatsController implements vscode.Disposable {
     ) {
         logDebug('ChatsController:constructor', 'init')
         this.panel = this.createChatController()
+        this.disposables.push(
+            this.authProvider.onChange(authStatus => this.setAuthStatus(authStatus), {
+                runImmediately: true,
+            })
+        )
     }
 
-    public async syncAuthStatus(authStatus: AuthStatus): Promise<void> {
+    private async setAuthStatus(authStatus: AuthStatus): Promise<void> {
         const hasLoggedOut = !authStatus.isLoggedIn
         const hasSwitchedAccount =
             this.currentAuthAccount && this.currentAuthAccount.endpoint !== authStatus.endpoint
@@ -88,9 +95,8 @@ export class ChatsController implements vscode.Disposable {
             username: authStatus.username,
         }
 
-        this.supportTreeViewProvider.syncAuthStatus(authStatus)
-
-        this.panel.syncAuthStatus()
+        this.supportTreeViewProvider.setAuthStatus(authStatus)
+        this.panel.setAuthStatus(authStatus)
     }
 
     public async restoreToPanel(panel: vscode.WebviewPanel, chatID: string): Promise<void> {

vscode/src/commands/GhostHintDecorator.ts

@@ -159,7 +159,12 @@ const GHOST_TEXT_THROTTLE = 250
 const TELEMETRY_THROTTLE = 30 * 1000 // 30 Seconds
 
 export class GhostHintDecorator implements vscode.Disposable {
-    private disposables: vscode.Disposable[] = []
+    // permanentDisposables are disposed when this instance is disposed.
+    private permanentDisposables: vscode.Disposable[] = []
+
+    // activeDisposables are disposed when the ghost hint is inactive (e.g., due to sign out)
+    private activeDisposables: vscode.Disposable[] = []
+
     private isActive = false
     private activeDecorationRange: vscode.Range | null = null
     private setThrottledGhostText: DebouncedFunc<typeof this.setGhostText>
@@ -196,10 +201,14 @@ export class GhostHintDecorator implements vscode.Disposable {
         this.updateEnablement(initialAuth)
 
         // Listen to authentication changes
-        authProvider.addChangeListener(authStatus => this.updateEnablement(authStatus))
+        this.permanentDisposables.push(
+            authProvider.onChange(authStatus => this.updateEnablement(authStatus), {
+                runImmediately: true,
+            })
+        )
 
         // Listen to configuration changes (e.g. if the setting is disabled)
-        this.disposables.push(
+        this.permanentDisposables.push(
             vscode.workspace.onDidChangeConfiguration(e => {
                 if (e.affectsConfiguration('cody')) {
                     this.updateEnablement(authProvider.getAuthStatus())
@@ -219,7 +228,7 @@ export class GhostHintDecorator implements vscode.Disposable {
     }
 
     private init(enabledFeatures: EnabledFeatures): void {
-        this.disposables.push(
+        this.activeDisposables.push(
             vscode.window.onDidChangeTextEditorSelection(
                 async (event: vscode.TextEditorSelectionChangeEvent) => {
                     const editor = event.textEditor
@@ -317,7 +326,7 @@ export class GhostHintDecorator implements vscode.Disposable {
         )
 
         if (enabledFeatures.Generate) {
-            this.disposables.push(
+            this.activeDisposables.push(
                 vscode.window.onDidChangeActiveTextEditor((editor?: vscode.TextEditor) => {
                     if (!editor) {
                         return
@@ -409,7 +418,7 @@ export class GhostHintDecorator implements vscode.Disposable {
             !authStatus.isLoggedIn ||
             !(featureEnablement.Document || featureEnablement.EditOrChat || featureEnablement.Generate)
         ) {
-            this.dispose()
+            this.disposeActive()
             return
         }
 
@@ -421,16 +430,24 @@ export class GhostHintDecorator implements vscode.Disposable {
     }
 
     public dispose(): void {
+        this.disposeActive()
+        for (const d of this.permanentDisposables) {
+            d.dispose()
+        }
+        this.permanentDisposables = []
+    }
+
+    private disposeActive(): void {
         this.isActive = false
 
         // Clear any existing ghost text
         if (vscode.window.activeTextEditor) {
             this.clearGhostText(vscode.window.activeTextEditor)
         }
 
-        for (const disposable of this.disposables) {
+        for (const disposable of this.activeDisposables) {
             disposable.dispose()
         }
-        this.disposables = []
+        this.activeDisposables = []
     }
 }

vscode/src/commands/scm/source-control.ts

@@ -231,7 +231,7 @@ export class CodySourceControl implements vscode.Disposable {
         this.abortController = abortController
     }
 
-    public syncAuthStatus(authStatus: AuthStatus): void {
+    public setAuthStatus(authStatus: AuthStatus): void {
         const models = ModelsService.getModels(ModelUsage.Chat, authStatus)
         const preferredModel = models?.find(p => p.model.includes('claude-3-haiku'))
         this.model = preferredModel ?? models[0]

vscode/src/completions/inline-completion-item-provider.test.ts

@@ -60,7 +60,7 @@ const DUMMY_AUTH_STATUS: AuthStatus = {
     codyApiVersion: 0,
 }
 
-graphqlClient.onConfigurationChange({} as unknown as GraphQLAPIClientConfig)
+graphqlClient.setConfig({} as unknown as GraphQLAPIClientConfig)
 
 class MockableInlineCompletionItemProvider extends InlineCompletionItemProvider {
     constructor(

vscode/src/completions/providers/create-provider.test.ts

@@ -35,7 +35,7 @@ const dummyCodeCompletionsClient: CodeCompletionsClient = {
 
 const dummyAuthStatus: AuthStatus = defaultAuthStatus
 
-graphqlClient.onConfigurationChange({} as unknown as GraphQLAPIClientConfig)
+graphqlClient.setConfig({} as unknown as GraphQLAPIClientConfig)
 
 describe('createProviderConfig', () => {
     describe('if completions provider fields are defined in VSCode settings', () => {

vscode/src/configwatcher.ts

@@ -1,53 +1,62 @@
+import type { ConfigurationWithAccessToken } from '@sourcegraph/cody-shared'
 import * as vscode from 'vscode'
+import { getFullConfig } from './configuration'
+import type { AuthProvider } from './services/AuthProvider'
+
+interface OnChangeOptions {
+    runImmediately: boolean
+}
 
 /**
  * A wrapper around a configuration source that lets the client retrieve the current config and watch for changes.
  */
 export interface ConfigWatcher<C> extends vscode.Disposable {
     get(): C
 
-    // NOTE(beyang): should remove this
-    set(c: C): void
-
     /*
      * Register a callback that is called only when Cody's configuration is changed.
-     * Appends to the disposable array methods that unregister the callback
-     */
-    onChange(callback: (config: C) => Promise<void>, disposables: vscode.Disposable[]): void
-
-    /**
-     * Same behavior as onChange, but fires the callback once immediately for initialization.
+     * Appends to the disposable array methods that unregister the callback.
+     *
+     * If `runImmediately` is true, the callback is called immediately and the returned
+     * Promise is that of the callback. If false (the default), then the return value
+     * is a resolved Promise.
      */
-    initAndOnChange(
+    onChange(
         callback: (config: C) => Promise<void>,
-        disposables: vscode.Disposable[]
+        disposables: vscode.Disposable[],
+        options?: OnChangeOptions
     ): Promise<void>
 }
 
-export class BaseConfigWatcher<C> implements ConfigWatcher<C> {
-    private currentConfig: C
+export class BaseConfigWatcher implements ConfigWatcher<ConfigurationWithAccessToken> {
+    private currentConfig: ConfigurationWithAccessToken
     private disposables: vscode.Disposable[] = []
-    private configChangeEvent: vscode.EventEmitter<C>
+    private configChangeEvent: vscode.EventEmitter<ConfigurationWithAccessToken>
 
-    public static async create<C>(
-        getConfig: () => Promise<C>,
+    public static async create(
+        authProvider: AuthProvider,
         disposables: vscode.Disposable[]
-    ): Promise<ConfigWatcher<C>> {
-        const w = new BaseConfigWatcher(await getConfig())
+    ): Promise<ConfigWatcher<ConfigurationWithAccessToken>> {
+        const w = new BaseConfigWatcher(await getFullConfig())
+        disposables.push(w)
         disposables.push(
             vscode.workspace.onDidChangeConfiguration(async event => {
                 if (!event.affectsConfiguration('cody')) {
                     return
                 }
-                w.set(await getConfig())
+                w.set(await getFullConfig())
+            })
+        )
+        disposables.push(
+            authProvider.onChange(async () => {
+                w.set(await getFullConfig())
             })
         )
-        disposables.push(w)
 
         return w
     }
 
-    constructor(initialConfig: C) {
+    constructor(initialConfig: ConfigurationWithAccessToken) {
         this.currentConfig = initialConfig
         this.configChangeEvent = new vscode.EventEmitter()
         this.disposables.push(this.configChangeEvent)
@@ -60,7 +69,22 @@ export class BaseConfigWatcher<C> implements ConfigWatcher<C> {
         this.disposables = []
     }
 
-    public set(config: C): void {
+    public get(): ConfigurationWithAccessToken {
+        return this.currentConfig
+    }
+
+    public async onChange(
+        callback: (config: ConfigurationWithAccessToken) => Promise<void>,
+        disposables: vscode.Disposable[],
+        { runImmediately }: OnChangeOptions = { runImmediately: false }
+    ): Promise<void> {
+        disposables.push(this.configChangeEvent.event(callback))
+        if (runImmediately) {
+            await callback(this.currentConfig)
+        }
+    }
+
+    private set(config: ConfigurationWithAccessToken): void {
         const oldConfig = JSON.stringify(this.currentConfig)
         const newConfig = JSON.stringify(config)
         if (oldConfig === newConfig) {
@@ -70,20 +94,4 @@ export class BaseConfigWatcher<C> implements ConfigWatcher<C> {
         this.currentConfig = config
         this.configChangeEvent.fire(config)
     }
-
-    get(): C {
-        return this.currentConfig
-    }
-
-    async initAndOnChange(
-        callback: (config: C) => Promise<void>,
-        disposables: vscode.Disposable[]
-    ): Promise<void> {
-        await callback(this.currentConfig)
-        this.onChange(callback, disposables)
-    }
-
-    public onChange(callback: (config: C) => Promise<void>, disposables: vscode.Disposable[]): void {
-        disposables.push(this.configChangeEvent.event(callback))
-    }
 }

vscode/src/extension.node.ts

@@ -23,7 +23,7 @@ import {
     createLocalEmbeddingsController,
 } from './local-context/local-embeddings'
 import { SymfRunner } from './local-context/symf'
-import { authProvider } from './services/AuthProvider'
+import { AuthProvider } from './services/AuthProvider'
 import { localStorage } from './services/LocalStorageProvider'
 import { OpenTelemetryService } from './services/open-telemetry/OpenTelemetryService.node'
 import { getExtensionDetails } from './services/telemetry-v2'
@@ -84,7 +84,7 @@ export function activate(
 // The vscode API is not available in the post-uninstall script.
 export async function deactivate(): Promise<void> {
     const config = localStorage.getConfig() ?? (await getFullConfig())
-    const authStatus = authProvider?.getAuthStatus() ?? defaultAuthStatus
+    const authStatus = AuthProvider.instance?.getAuthStatus() ?? defaultAuthStatus
     const { anonymousUserID } = await localStorage.anonymousUserID()
     serializeConfigSnapshot({
         config,