do not block on CodyLLMSiteConfiguration (configOverwrites) fetch in …

…initial auth To authenticate a user and start showing the UI, we actually don't need to fetch the CodyLLMSiteConfiguration (`configOverwrites`). This reduces the number of HTTP requests (albeit they were parallelized somewhat) needed before the user can be authenticated and start seeing the Cody UI. This uncovered some possible "bugs" in the `anthropic.test.ts` autocomplete tests, where they were asserting that the `model` would be set but it actually would not be in real life. In `Provider.maybeFilterOutModel` it sees that `this.configOverwrites?.provider === 'sourcegraph'` in the test, but that should not be true in real life because it’s being passed an `authStatus` that is non-dotcom. The reason the test is bad is because that method gets the auth status from `currentAuthStatus()`, which is dotcom in testing only because of a prior test's mock value.
sourcegraph · Oct 5, 2024 · 00553ff · 00553ff
1 parent 6798a68
commit 00553ff
Show file tree

Hide file tree

Showing 26 changed files with 315 additions and 192 deletions.
diff --git a/...gs/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/AuthStatus.kt b/...gs/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/AuthStatus.kt
@@ -39,7 +39,6 @@ data class AuthenticatedAuthStatus(
   val requiresVerifiedEmail: Boolean? = null,
   val siteVersion: String,
   val codyApiVersion: Long,
-  val configOverwrites: CodyLLMSiteConfiguration? = null,
   val primaryEmail: String? = null,
   val displayName: String? = null,
   val avatarURL: String? = null,

diff --git a/lib/shared/src/auth/types.ts b/lib/shared/src/auth/types.ts
@@ -1,5 +1,4 @@
 import { isDotCom } from '../sourcegraph-api/environments'
-import type { CodyLLMSiteConfiguration } from '../sourcegraph-api/graphql/client'
 
 /**
  * The authentication status, which includes representing the state when authentication failed or
@@ -27,7 +26,6 @@ export interface AuthenticatedAuthStatus {
     requiresVerifiedEmail?: boolean
     siteVersion: string
     codyApiVersion: number
-    configOverwrites?: CodyLLMSiteConfiguration
 
     primaryEmail?: string
     displayName?: string
@@ -78,10 +76,6 @@ export const AUTH_STATUS_FIXTURE_UNAUTHED: AuthStatus & { authenticated: false }
 export const AUTH_STATUS_FIXTURE_AUTHED_DOTCOM: AuthenticatedAuthStatus = {
     ...AUTH_STATUS_FIXTURE_AUTHED,
     endpoint: 'https://sourcegraph.com',
-    configOverwrites: {
-        provider: 'sourcegraph',
-        completionModel: 'fireworks/starcoder-hybrid',
-    },
 }
 
 export function isCodyProUser(authStatus: AuthStatus): boolean {

diff --git a/lib/shared/src/index.ts b/lib/shared/src/index.ts
@@ -382,3 +382,4 @@ export * from './singletons'
 export * from './auth/authStatus'
 export { fetchLocalOllamaModels } from './llm-providers/ollama/utils'
 export * from './editor/editorState'
+export { configOverwrites } from './models/configOverwrites'
diff --git a/lib/shared/src/misc/observable.ts b/lib/shared/src/misc/observable.ts
@@ -546,13 +546,16 @@ export function pluck<T>(...keyPath: any[]): (input: ObservableLike<T>) => Obser
 }
 
 export function pick<T, K extends keyof T>(
-    key: K
+    ...keys: K[]
 ): (input: ObservableLike<T>) => Observable<Pick<T, K>> {
-    return map(
-        value =>
-            ({
-                [key]: value[key],
-            }) as Pick<T, K>
+    return map(value =>
+        keys.reduce(
+            (acc, key) => {
+                acc[key] = value[key]
+                return acc
+            },
+            {} as Pick<T, K>
+        )
     )
 }
 

diff --git a/lib/shared/src/models/configOverwrites.ts b/lib/shared/src/models/configOverwrites.ts
@@ -0,0 +1,52 @@
+import { Observable, map } from 'observable-fns'
+import { authStatus } from '../auth/authStatus'
+import { logError } from '../logger'
+import { distinctUntilChanged, pick, promiseFactoryToObservable } from '../misc/observable'
+import { pendingOperation, switchMapReplayOperation } from '../misc/observableOperation'
+import { type CodyLLMSiteConfiguration, graphqlClient } from '../sourcegraph-api/graphql/client'
+import { isError } from '../utils'
+
+/**
+ * Observe the model-related config overwrites on the server for the currently authenticated user.
+ */
+export const configOverwrites: Observable<CodyLLMSiteConfiguration | null | typeof pendingOperation> =
+    authStatus.pipe(
+        pick('authenticated', 'endpoint', 'pendingValidation'),
+        distinctUntilChanged(),
+        switchMapReplayOperation(
+            (
+                authStatus
+            ): Observable<CodyLLMSiteConfiguration | Error | null | typeof pendingOperation> => {
+                if (authStatus.pendingValidation) {
+                    return Observable.of(pendingOperation)
+                }
+
+                if (!authStatus.authenticated) {
+                    return Observable.of(null)
+                }
+
+                return promiseFactoryToObservable(signal =>
+                    graphqlClient.getCodyLLMConfiguration(signal)
+                ).pipe(
+                    map((result): CodyLLMSiteConfiguration | null | typeof pendingOperation => {
+                        if (isError(result)) {
+                            logError(
+                                'configOverwrites',
+                                `Failed to get Cody LLM configuration from ${authStatus.endpoint}: ${result}`
+                            )
+                            return null
+                        }
+                        return result ?? null
+                    })
+                )
+            }
+        ),
+        map(result => (isError(result) ? null : result)) // the operation catches its own errors, so errors will never get here
+    )
+
+// Subscribe so that other subscribers get the replayed value. There are no other permanent
+// subscribers to this value.
+//
+// TODO(sqs): This fixes an issue where switching accounts (`rtx exec node@18.17.1 -- pnpm run test
+// agent/src/auth.test.ts -t 'switches'`) took ~2.7s on Node 18.
+configOverwrites.subscribe({})
diff --git a/lib/shared/src/models/modelsService.ts b/lib/shared/src/models/modelsService.ts
@@ -21,6 +21,7 @@ import {
 } from '../misc/observableOperation'
 import { ClientConfigSingleton } from '../sourcegraph-api/clientConfig'
 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'
 import { ModelTag } from './tags'
@@ -277,6 +278,7 @@ export class ModelsService {
             distinctUntilChanged()
         ),
         authStatus,
+        configOverwrites,
         clientConfig: ClientConfigSingleton.getInstance().changes,
     })
 

diff --git a/lib/shared/src/models/sync.test.ts b/lib/shared/src/models/sync.test.ts
@@ -13,6 +13,7 @@ import {
 } from '../misc/observable'
 import { pendingOperation, skipPendingOperation } from '../misc/observableOperation'
 import type { CodyClientConfig } from '../sourcegraph-api/clientConfig'
+import type { CodyLLMSiteConfiguration } from '../sourcegraph-api/graphql/client'
 import type { PartialDeep } from '../utils'
 import {
     type Model,
@@ -198,6 +199,7 @@ describe('server sent models', async () => {
                 clientState: { modelPreferences: {} },
             } satisfies PartialDeep<ResolvedConfiguration> as ResolvedConfiguration),
             authStatus: Observable.of(AUTH_STATUS_FIXTURE_AUTHED),
+            configOverwrites: Observable.of(null),
             clientConfig: Observable.of({
                 modelsAPIEnabled: true,
             } satisfies Partial<CodyClientConfig> as CodyClientConfig),
@@ -243,13 +245,15 @@ describe('syncModels', () => {
                 (): Promise<ServerModelConfiguration | undefined> => Promise.resolve(undefined)
             )
             const authStatusSubject = new Subject<AuthStatus>()
+            const configOverwritesSubject = new Subject<CodyLLMSiteConfiguration | null>()
             const clientConfigSubject = new Subject<CodyClientConfig>()
             const syncModelsObservable = syncModels({
                 resolvedConfig: Observable.of({
                     configuration: {},
                     clientState: { modelPreferences: {} },
                 } satisfies PartialDeep<ResolvedConfiguration> as ResolvedConfiguration),
                 authStatus: authStatusSubject.pipe(shareReplay()),
+                configOverwrites: configOverwritesSubject.pipe(shareReplay()),
                 clientConfig: clientConfigSubject.pipe(shareReplay()),
                 fetchServerSideModels_: mockFetchServerSideModels,
             })
@@ -279,16 +283,15 @@ describe('syncModels', () => {
                 } satisfies Partial<ServerModel> as ServerModel
             }
 
-            // Emits when authStatus emits.
-            authStatusSubject.next({
-                ...AUTH_STATUS_FIXTURE_AUTHED,
-                configOverwrites: { chatModel: 'foo' },
-            })
+            // Emits when authStatus configOverwrites emits.
+            authStatusSubject.next(AUTH_STATUS_FIXTURE_AUTHED)
             await vi.advanceTimersByTimeAsync(0)
             clientConfigSubject.next({
                 modelsAPIEnabled: false,
             } satisfies Partial<CodyClientConfig> as CodyClientConfig)
             await vi.advanceTimersByTimeAsync(0)
+            configOverwritesSubject.next({ chatModel: 'foo' })
+            await vi.advanceTimersByTimeAsync(0)
             await vi.runOnlyPendingTimersAsync()
             expect(values).toStrictEqual<typeof values>([
                 pendingOperation,
@@ -304,15 +307,12 @@ describe('syncModels', () => {
             clearValues()
 
             // Emits immediately when the new data can be computed synchronously.
-            authStatusSubject.next({
-                ...AUTH_STATUS_FIXTURE_AUTHED,
-                configOverwrites: { chatModel: 'bar' },
-            })
-            await vi.advanceTimersByTimeAsync(0)
             clientConfigSubject.next({
                 modelsAPIEnabled: false,
             } satisfies Partial<CodyClientConfig> as CodyClientConfig)
             await vi.advanceTimersByTimeAsync(0)
+            configOverwritesSubject.next({ chatModel: 'bar' })
+            await vi.advanceTimersByTimeAsync(0)
             const result0: ModelsData = {
                 localModels: [],
                 primaryModels: [modelFixture('bar')],

diff --git a/lib/shared/src/models/sync.ts b/lib/shared/src/models/sync.ts
@@ -10,6 +10,7 @@ import { logDebug } from '../logger'
 import {
     combineLatest,
     distinctUntilChanged,
+    pick,
     promiseFactoryToObservable,
     shareReplay,
     startWith,
@@ -21,6 +22,7 @@ import { pendingOperation, switchMapReplayOperation } from '../misc/observableOp
 import { ANSWER_TOKENS } from '../prompt/constants'
 import type { CodyClientConfig } from '../sourcegraph-api/clientConfig'
 import { isDotCom } from '../sourcegraph-api/environments'
+import type { CodyLLMSiteConfiguration } from '../sourcegraph-api/graphql/client'
 import { RestClient } from '../sourcegraph-api/rest/client'
 import { CHAT_INPUT_TOKEN_BUDGET } from '../token/constants'
 import { isError } from '../utils'
@@ -39,6 +41,7 @@ const EMPTY_PREFERENCES: SitePreferences = { defaults: {}, selected: {} }
 export function syncModels({
     resolvedConfig,
     authStatus,
+    configOverwrites,
     clientConfig,
     fetchServerSideModels_ = fetchServerSideModels,
 }: {
@@ -50,6 +53,7 @@ export function syncModels({
         }>
     >
     authStatus: Observable<AuthStatus>
+    configOverwrites: Observable<CodyLLMSiteConfiguration | null | typeof pendingOperation>
     clientConfig: Observable<CodyClientConfig | undefined | typeof pendingOperation>
     fetchServerSideModels_?: typeof fetchServerSideModels
 }): Observable<ModelsData | typeof pendingOperation> {
@@ -109,11 +113,11 @@ export function syncModels({
             map(config => config.clientState.modelPreferences),
             distinctUntilChanged()
         ),
-        authStatus
+        authStatus.pipe(pick('endpoint'), distinctUntilChanged())
     ).pipe(
-        map(([modelPreferences, authStatus]) => {
+        map(([modelPreferences, { endpoint }]) => {
             // Deep clone so it's not readonly and we can mutate it, for convenience.
-            const prevPreferences = modelPreferences[authStatus.endpoint] as SitePreferences | undefined
+            const prevPreferences = modelPreferences[endpoint] as SitePreferences | undefined
             return deepClone(
                 (prevPreferences ?? EMPTY_PREFERENCES) satisfies SitePreferences as SitePreferences
             )
@@ -235,34 +239,41 @@ export function syncModels({
                         // BYOK customers to set a model of their choice without us having to map it to a known
                         // model on the client.
                         //
-                        // NOTE: If authStatus?.configOverwrites?.chatModel is empty, automatically fallback to
-                        // use the default model configured on the instance.
-                        if (authStatus?.configOverwrites?.chatModel) {
-                            return Observable.of<RemoteModelsData>({
-                                preferences: null,
-                                primaryModels: [
-                                    createModel({
-                                        id: authStatus.configOverwrites.chatModel,
-                                        // TODO (umpox) Add configOverwrites.editModel for separate edit support
-                                        usage: [ModelUsage.Chat, ModelUsage.Edit],
-                                        contextWindow: getEnterpriseContextWindow(
-                                            authStatus?.configOverwrites?.chatModel,
-                                            authStatus?.configOverwrites,
-                                            config.configuration
-                                        ),
-                                        tags: [ModelTag.Enterprise],
-                                    }),
-                                ],
-                            })
-                        }
+                        // NOTE: If configOverwrites?.chatModel is empty, automatically fallback to use the
+                        // default model configured on the instance.
+                        return configOverwrites.pipe(
+                            map((configOverwrites): RemoteModelsData | typeof pendingOperation => {
+                                if (configOverwrites === pendingOperation) {
+                                    return pendingOperation
+                                }
+                                if (configOverwrites?.chatModel) {
+                                    return {
+                                        preferences: null,
+                                        primaryModels: [
+                                            createModel({
+                                                id: configOverwrites.chatModel,
+                                                // TODO (umpox) Add configOverwrites.editModel for separate edit support
+                                                usage: [ModelUsage.Chat, ModelUsage.Edit],
+                                                contextWindow: getEnterpriseContextWindow(
+                                                    configOverwrites?.chatModel,
+                                                    configOverwrites,
+                                                    config.configuration
+                                                ),
+                                                tags: [ModelTag.Enterprise],
+                                            }),
+                                        ],
+                                    } satisfies RemoteModelsData
+                                }
 
-                        // If the enterprise instance didn't have any configuration data for Cody, clear the
-                        // models available in the modelsService. Otherwise there will be stale, defunct models
-                        // available.
-                        return Observable.of<RemoteModelsData>({
-                            preferences: null,
-                            primaryModels: [],
-                        })
+                                // If the enterprise instance didn't have any configuration data for Cody, clear the
+                                // models available in the modelsService. Otherwise there will be stale, defunct models
+                                // available.
+                                return {
+                                    preferences: null,
+                                    primaryModels: [],
+                                } satisfies RemoteModelsData
+                            })
+                        )
                     })
                 )
                 return serverModelsConfig

diff --git a/vscode/src/auth/auth.ts b/vscode/src/auth/auth.ts
@@ -417,12 +417,10 @@ export async function validateCredentials(
         clientState: config.clientState,
     })
 
-    const [{ enabled: siteHasCodyEnabled, version: siteVersion }, codyLLMConfiguration, userInfo] =
-        await Promise.all([
-            client.isCodyEnabled(signal),
-            client.getCodyLLMConfiguration(signal),
-            client.getCurrentUserInfo(signal),
-        ])
+    const [{ enabled: siteHasCodyEnabled, version: siteVersion }, userInfo] = await Promise.all([
+        client.isCodyEnabled(signal),
+        client.getCurrentUserInfo(signal),
+    ])
     signal?.throwIfAborted()
 
     if (isError(userInfo) && isNetworkLikeError(userInfo)) {
@@ -464,14 +462,11 @@ export async function validateCredentials(
 
     logDebug('auth', `Authentication succeeed to endpoint ${config.auth.serverEndpoint}`)
 
-    const configOverwrites = isError(codyLLMConfiguration) ? undefined : codyLLMConfiguration
-
     if (!isDotCom(config.auth.serverEndpoint)) {
         return newAuthStatus({
             ...userInfo,
             endpoint: config.auth.serverEndpoint,
             siteVersion,
-            configOverwrites,
             authenticated: true,
             hasVerifiedEmail: false,
             userCanUpgrade: false,
@@ -498,7 +493,6 @@ export async function validateCredentials(
         authenticated: true,
         endpoint: config.auth.serverEndpoint,
         siteVersion,
-        configOverwrites,
         userCanUpgrade: !isActiveProUser,
     })
 }
diff --git a/vscode/src/chat/utils.ts b/vscode/src/chat/utils.ts
@@ -11,7 +11,6 @@ type NewAuthStatusOptions = { endpoint: string } & (
           | 'authenticated'
           | 'username'
           | 'siteVersion'
-          | 'configOverwrites'
           | 'hasVerifiedEmail'
           | 'displayName'
           | 'avatarURL'

diff --git a/vscode/src/completions/create-inline-completion-item-provider.ts b/vscode/src/completions/create-inline-completion-item-provider.ts
@@ -6,6 +6,7 @@ import {
     NEVER,
     type PickResolvedConfiguration,
     type UnauthenticatedAuthStatus,
+    configOverwrites,
     createDisposables,
     promiseFactoryToObservable,
     skipPendingOperation,
@@ -22,9 +23,7 @@ import { registerAutocompleteTraceView } from './tracer/traceView'
 
 interface InlineCompletionItemProviderArgs {
     config: PickResolvedConfiguration<{ configuration: true }>
-    authStatus:
-        | UnauthenticatedAuthStatus
-        | Pick<AuthenticatedAuthStatus, 'authenticated' | 'endpoint' | 'configOverwrites'>
+    authStatus: UnauthenticatedAuthStatus | Pick<AuthenticatedAuthStatus, 'authenticated' | 'endpoint'>
     platform: Pick<PlatformContext, 'extensionClient'>
     statusBar: CodyStatusBar
 }
@@ -61,7 +60,7 @@ export function createInlineCompletionItemProvider({
         return await getInlineCompletionItemProviderFilters(configuration.autocompleteLanguages)
     }).pipe(
         switchMap(documentFilters =>
-            createProvider({ config: { configuration }, authStatus }).pipe(
+            createProvider({ config: { configuration }, authStatus, configOverwrites }).pipe(
                 skipPendingOperation(),
                 createDisposables(providerOrError => {
                     if (providerOrError instanceof Error) {