From a60019159f17178508ff37f0b7a701b4864a37ac Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Wed, 20 Mar 2024 16:13:28 -0700
Subject: [PATCH 01/14] Chat: sync token limit at model import

---
 lib/shared/src/models/dotcom.ts               | 12 ++++++
 lib/shared/src/models/index.ts                | 20 +++++++--
 vscode/src/chat/chat-view/ChatManager.ts      | 14 ++-----
 .../src/chat/chat-view/ChatPanelsManager.ts   | 14 ++++---
 .../chat/chat-view/SimpleChatPanelProvider.ts | 17 ++------
 vscode/src/edit/provider.ts                   |  7 +---
 vscode/src/edit/utils/edit-models.ts          | 11 +----
 vscode/src/models/utilts.ts                   | 42 +++++++------------
 8 files changed, 62 insertions(+), 75 deletions(-)

diff --git a/lib/shared/src/models/dotcom.ts b/lib/shared/src/models/dotcom.ts
index 03fc6ed929e..b085ef9aabd 100644
--- a/lib/shared/src/models/dotcom.ts
+++ b/lib/shared/src/models/dotcom.ts
@@ -1,6 +1,9 @@
 import type { ModelProvider } from '.'
 import { ModelUsage } from './types'
 
+export const DEFAULT_CHAT_MODEL_TOKEN_LIMIT = 28000 // 7000 tokens * 4 bytes per token
+export const DEFAULT_FAST_MODEL_TOKEN_LIMIT = 10000 // 4,096 tokens * 4 bytes per token
+
 // The models must first be added to the custom chat models list in https://sourcegraph.com/github.com/sourcegraph/sourcegraph/-/blob/internal/completions/httpapi/chat.go?L48-51
 export const DEFAULT_DOT_COM_MODELS = [
     {
@@ -10,6 +13,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: true,
         codyProOnly: false,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 2.1',
@@ -18,6 +22,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude Instant',
@@ -26,6 +31,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 3 Haiku',
@@ -34,6 +40,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 3 Sonnet',
@@ -42,6 +49,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 3 Opus',
@@ -50,6 +58,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'GPT-3.5 Turbo',
@@ -58,6 +67,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'GPT-4 Turbo Preview',
@@ -66,6 +76,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Mixtral 8x7B',
@@ -75,5 +86,6 @@ export const DEFAULT_DOT_COM_MODELS = [
         codyProOnly: true,
         // TODO: Improve prompt for Mixtral + Edit to see if we can use it there too.
         usage: [ModelUsage.Chat],
+        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
 ] as const satisfies ModelProvider[]
diff --git a/lib/shared/src/models/index.ts b/lib/shared/src/models/index.ts
index 653ba899594..768da4339a3 100644
--- a/lib/shared/src/models/index.ts
+++ b/lib/shared/src/models/index.ts
@@ -1,7 +1,11 @@
 import { logError } from '../logger'
 import { OLLAMA_DEFAULT_URL } from '../ollama'
 import { isDotCom } from '../sourcegraph-api/environments'
-import { DEFAULT_DOT_COM_MODELS } from './dotcom'
+import {
+    DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+    DEFAULT_DOT_COM_MODELS,
+    DEFAULT_FAST_MODEL_TOKEN_LIMIT,
+} from './dotcom'
 import { ModelUsage } from './types'
 import { getProviderName } from './utils'
 
@@ -15,16 +19,18 @@ export class ModelProvider {
     public codyProOnly = false
     public provider: string
     public readonly title: string
+    public readonly contextWindow: number
 
     constructor(
         public readonly model: string,
         public readonly usage: ModelUsage[],
-        isDefaultModel = true
+        tokenLimit?: number
     ) {
         const splittedModel = model.split('/')
         this.provider = getProviderName(splittedModel[0])
         this.title = splittedModel[1]?.replaceAll('-', ' ')
-        this.default = isDefaultModel
+        this.default = true
+        this.contextWindow = tokenLimit ? tokenLimit * 4 : DEFAULT_FAST_MODEL_TOKEN_LIMIT
     }
 
     // Providers available for non-dotcom instances
@@ -120,4 +126,12 @@ export class ModelProvider {
             }
         })
     }
+
+    public static getContextWindow(modelID: string): number {
+        return (
+            ModelProvider.privateProviders.get(modelID)?.contextWindow ||
+            ModelProvider.dotComProviders.find(model => model.model === modelID)?.contextWindow ||
+            DEFAULT_CHAT_MODEL_TOKEN_LIMIT
+        )
+    }
 }
diff --git a/vscode/src/chat/chat-view/ChatManager.ts b/vscode/src/chat/chat-view/ChatManager.ts
index ecb227ea31f..2186723bc06 100644
--- a/vscode/src/chat/chat-view/ChatManager.ts
+++ b/vscode/src/chat/chat-view/ChatManager.ts
@@ -7,7 +7,6 @@ import {
     CODY_PASSTHROUGH_VSCODE_OPEN_COMMAND_ID,
     type ChatClient,
     type Guardrails,
-    ModelProvider,
 } from '@sourcegraph/cody-shared'
 
 import type { View } from '../../../webviews/NavBar'
@@ -19,10 +18,10 @@ import { localStorage } from '../../services/LocalStorageProvider'
 import { telemetryService } from '../../services/telemetry'
 import { telemetryRecorder } from '../../services/telemetry-v2'
 
-import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import type { ExecuteChatArguments } from '../../commands/execute/ask'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
+import { addEnterpriseChatModel } from '../../models/utilts'
 import { ChatPanelsManager } from './ChatPanelsManager'
 import { SidebarViewController, type SidebarViewOptions } from './SidebarViewController'
 import type { ChatSession, SimpleChatPanelProvider } from './SimpleChatPanelProvider'
@@ -100,15 +99,8 @@ export class ChatManager implements vscode.Disposable {
     }
 
     public async syncAuthStatus(authStatus: AuthStatus): Promise<void> {
-        if (authStatus?.configOverwrites?.chatModel) {
-            ModelProvider.add(
-                new ModelProvider(authStatus.configOverwrites.chatModel, [
-                    ModelUsage.Chat,
-                    // TODO: Add configOverwrites.editModel for separate edit support
-                    ModelUsage.Edit,
-                ])
-            )
-        }
+        addEnterpriseChatModel(authStatus)
+
         await this.chatPanelsManager.syncAuthStatus(authStatus)
     }
 
diff --git a/vscode/src/chat/chat-view/ChatPanelsManager.ts b/vscode/src/chat/chat-view/ChatPanelsManager.ts
index 3d18b8bb0a7..68480efd383 100644
--- a/vscode/src/chat/chat-view/ChatPanelsManager.ts
+++ b/vscode/src/chat/chat-view/ChatPanelsManager.ts
@@ -217,11 +217,15 @@ export class ChatPanelsManager implements vscode.Disposable {
         const authStatus = authProvider.getAuthStatus()
         if (authStatus?.configOverwrites?.chatModel) {
             ModelProvider.add(
-                new ModelProvider(authStatus.configOverwrites.chatModel, [
-                    ModelUsage.Chat,
-                    // TODO: Add configOverwrites.editModel for separate edit support
-                    ModelUsage.Edit,
-                ])
+                new ModelProvider(
+                    authStatus.configOverwrites.chatModel,
+                    [
+                        ModelUsage.Chat,
+                        // TODO: Add configOverwrites.editModel for separate edit support
+                        ModelUsage.Edit,
+                    ],
+                    authStatus.configOverwrites?.chatModelMaxTokens
+                )
             )
         }
         const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint)
diff --git a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
index 90229770577..421c0422f36 100644
--- a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
+++ b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
@@ -58,7 +58,7 @@ import type { Repo } from '../../context/repo-fetcher'
 import type { RemoteRepoPicker } from '../../context/repo-picker'
 import type { ContextRankingController } from '../../local-context/context-ranking'
 import { chatModel } from '../../models'
-import { getContextWindowForModel } from '../../models/utilts'
+import { addEnterpriseChatModel } from '../../models/utilts'
 import { recordExposedExperimentsToSpan } from '../../services/open-telemetry/utils'
 import type { MessageErrorType } from '../MessageProvider'
 import { getChatContextItemsForMention } from '../context/chatContext'
@@ -727,15 +727,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
         if (!authStatus?.isLoggedIn) {
             return
         }
-        if (authStatus?.configOverwrites?.chatModel) {
-            ModelProvider.add(
-                new ModelProvider(authStatus.configOverwrites.chatModel, [
-                    ModelUsage.Chat,
-                    // TODO: Add configOverwrites.editModel for separate edit support
-                    ModelUsage.Edit,
-                ])
-            )
-        }
+        addEnterpriseChatModel(authStatus)
         const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint, this.chatModel.modelID)
 
         void this.postMessage({
@@ -786,10 +778,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
         prompter: IPrompter,
         sendTelemetry?: (contextSummary: any) => void
     ): Promise<Message[]> {
-        const maxChars = getContextWindowForModel(
-            this.authProvider.getAuthStatus(),
-            this.chatModel.modelID
-        )
+        const maxChars = ModelProvider.getContextWindow(this.chatModel.modelID)
         const { prompt, newContextUsed } = await prompter.makePrompt(this.chatModel, maxChars)
 
         // Update UI based on prompt construction
diff --git a/vscode/src/edit/provider.ts b/vscode/src/edit/provider.ts
index a931df58ef3..6fbfedc4eda 100644
--- a/vscode/src/edit/provider.ts
+++ b/vscode/src/edit/provider.ts
@@ -2,6 +2,7 @@ import { Utils } from 'vscode-uri'
 
 import {
     BotResponseMultiplexer,
+    ModelProvider,
     Typewriter,
     isAbortError,
     isDotCom,
@@ -17,7 +18,6 @@ import { isNetworkError } from '../services/AuthProvider'
 
 import { workspace } from 'vscode'
 import { doesFileExist } from '../commands/utils/workspace-files'
-import { getContextWindowForModel } from '../models/utilts'
 import { CodyTaskState } from '../non-stop/utils'
 import { telemetryService } from '../services/telemetry'
 import { telemetryRecorder } from '../services/telemetry-v2'
@@ -43,10 +43,7 @@ export class EditProvider {
         return wrapInActiveSpan('command.edit.start', async span => {
             this.config.controller.startTask(this.config.task)
             const model = this.config.task.model
-            const contextWindow = getContextWindowForModel(
-                this.config.authProvider.getAuthStatus(),
-                model
-            )
+            const contextWindow = ModelProvider.getContextWindow(model)
             const {
                 messages,
                 stopSequences,
diff --git a/vscode/src/edit/utils/edit-models.ts b/vscode/src/edit/utils/edit-models.ts
index 6ca4430395e..570a1269b4c 100644
--- a/vscode/src/edit/utils/edit-models.ts
+++ b/vscode/src/edit/utils/edit-models.ts
@@ -1,17 +1,10 @@
 import { type AuthStatus, ModelProvider } from '@sourcegraph/cody-shared'
 import { type EditModel, ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
+import { addEnterpriseChatModel } from '../../models/utilts'
 import type { EditIntent } from '../types'
 
 export function getEditModelsForUser(authStatus: AuthStatus): ModelProvider[] {
-    if (authStatus?.configOverwrites?.chatModel) {
-        ModelProvider.add(
-            new ModelProvider(authStatus.configOverwrites.chatModel, [
-                ModelUsage.Chat,
-                // TODO: Add configOverwrites.editModel for separate edit support
-                ModelUsage.Edit,
-            ])
-        )
-    }
+    addEnterpriseChatModel(authStatus)
     return ModelProvider.get(ModelUsage.Edit, authStatus.endpoint)
 }
 
diff --git a/vscode/src/models/utilts.ts b/vscode/src/models/utilts.ts
index f0493fd02bb..ba3cf4736fd 100644
--- a/vscode/src/models/utilts.ts
+++ b/vscode/src/models/utilts.ts
@@ -1,39 +1,25 @@
+import { type AuthStatus, ModelProvider } from '@sourcegraph/cody-shared'
+import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import * as vscode from 'vscode'
 
-import { type AuthStatus, type ChatModel, type EditModel, isDotCom } from '@sourcegraph/cody-shared'
-
-export function getContextWindowForModel(
-    authStatus: AuthStatus,
-    modelID: EditModel | ChatModel
-): number {
+export function addEnterpriseChatModel(authStatus: AuthStatus): void {
     // In enterprise mode, we let the sg instance dictate the token limits and allow users to
     // overwrite it locally (for debugging purposes).
     //
     // This is similiar to the behavior we had before introducing the new chat and allows BYOK
     // customers to set a model of their choice without us having to map it to a known model on
     // the client.
-    if (authStatus.endpoint && !isDotCom(authStatus.endpoint)) {
+    if (authStatus?.configOverwrites?.chatModel) {
         const codyConfig = vscode.workspace.getConfiguration('cody')
-        const tokenLimit = codyConfig.get<number>('provider.limit.prompt')
-        if (tokenLimit) {
-            return tokenLimit * 4 // bytes per token
-        }
-
-        if (authStatus.configOverwrites?.chatModelMaxTokens) {
-            return authStatus.configOverwrites.chatModelMaxTokens * 4 // bytes per token
-        }
-
-        return 28000 // 7000 tokens * 4 bytes per token
-    }
-
-    if (modelID === 'openai/gpt-4-turbo-preview') {
-        return 28000 // 7000 tokens * 4 bytes per token
-    }
-    if (modelID === 'openai/gpt-3.5-turbo') {
-        return 10000 // 4,096 tokens * < 4 bytes per token
-    }
-    if (modelID === 'fireworks/accounts/fireworks/models/mixtral-8x7b-instruct') {
-        return 28000 // 7000 tokens * 4 bytes per token
+        const tokenLimitConfig = codyConfig.get<number>('provider.limit.prompt')
+        const tokenLimit = tokenLimitConfig ?? authStatus.configOverwrites?.chatModelMaxTokens
+        ModelProvider.add(
+            new ModelProvider(
+                authStatus.configOverwrites.chatModel,
+                // TODO: Add configOverwrites.editModel for separate edit support
+                [ModelUsage.Chat, ModelUsage.Edit],
+                tokenLimit
+            )
+        )
     }
-    return 28000 // assume default to Claude-2-like model
 }

From ed4ec1e96d4b972b3114c6eb4cb214be285d2717 Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Wed, 20 Mar 2024 16:24:30 -0700
Subject: [PATCH 02/14] clean up

---
 vscode/src/chat/chat-view/ChatPanelsManager.ts | 15 ++-------------
 1 file changed, 2 insertions(+), 13 deletions(-)

diff --git a/vscode/src/chat/chat-view/ChatPanelsManager.ts b/vscode/src/chat/chat-view/ChatPanelsManager.ts
index 68480efd383..36aae8d01ee 100644
--- a/vscode/src/chat/chat-view/ChatPanelsManager.ts
+++ b/vscode/src/chat/chat-view/ChatPanelsManager.ts
@@ -23,6 +23,7 @@ import type { ExtensionMessage } from '../protocol'
 import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
+import { addEnterpriseChatModel } from '../../models/utilts'
 import { chatHistory } from './ChatHistoryManager'
 import { CodyChatPanelViewType } from './ChatManager'
 import type { SidebarViewOptions } from './SidebarViewController'
@@ -215,19 +216,7 @@ export class ChatPanelsManager implements vscode.Disposable {
     private createProvider(): SimpleChatPanelProvider {
         const authProvider = this.options.authProvider
         const authStatus = authProvider.getAuthStatus()
-        if (authStatus?.configOverwrites?.chatModel) {
-            ModelProvider.add(
-                new ModelProvider(
-                    authStatus.configOverwrites.chatModel,
-                    [
-                        ModelUsage.Chat,
-                        // TODO: Add configOverwrites.editModel for separate edit support
-                        ModelUsage.Edit,
-                    ],
-                    authStatus.configOverwrites?.chatModelMaxTokens
-                )
-            )
-        }
+        addEnterpriseChatModel(authStatus)
         const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint)
         const isConsumer = authProvider.getAuthStatus().isDotCom
 

From 23ee6717ab1e595db58c6a52f6f3f0941aa4fc3e Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Wed, 20 Mar 2024 18:40:52 -0700
Subject: [PATCH 03/14] Re-generate Kotlin bindings

---
 .../com/sourcegraph/cody/protocol_generated/ModelProvider.kt     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
index a0ddc23e815..6ffa86d39ba 100644
--- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
+++ b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
@@ -6,6 +6,7 @@ data class ModelProvider(
   val codyProOnly: Boolean? = null,
   val provider: String? = null,
   val title: String? = null,
+  val contextWindow: Int? = null,
   val privateProviders: Map<String, ModelProvider>? = null,
   val dotComProviders: List<ModelProvider>? = null,
   val ollamaProvidersEnabled: Boolean? = null,

From e34d780b8b72b0d494c7cb82106573d30972acd7 Mon Sep 17 00:00:00 2001
From: Beatrix <68532117+abeatrix@users.noreply.github.com>
Date: Thu, 21 Mar 2024 07:34:56 -0700
Subject: [PATCH 04/14] Update
 agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt

Co-authored-by: Dominic Cooney <dominic.cooney@sourcegraph.com>
---
 .../com/sourcegraph/cody/protocol_generated/ModelProvider.kt    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
index 6ffa86d39ba..5bd030effd7 100644
--- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
+++ b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
@@ -6,7 +6,7 @@ data class ModelProvider(
   val codyProOnly: Boolean? = null,
   val provider: String? = null,
   val title: String? = null,
-  val contextWindow: Int? = null,
+  val contextWindowSizeChars: Int? = null,
   val privateProviders: Map<String, ModelProvider>? = null,
   val dotComProviders: List<ModelProvider>? = null,
   val ollamaProvidersEnabled: Boolean? = null,

From ad53e9817515341abbb7b6419e85f4f6e591fdf3 Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 09:39:50 -0700
Subject: [PATCH 05/14] clean up and add tests

---
 lib/shared/src/models/dotcom.ts               | 22 +++--
 lib/shared/src/models/index.test.ts           | 35 ++++++++
 lib/shared/src/models/index.ts                | 82 ++++++++++---------
 lib/shared/src/models/utils.test.ts           | 36 ++++++++
 lib/shared/src/models/utils.ts                | 15 ++++
 lib/shared/src/prompt/constants.ts            | 23 +++++-
 .../chat/chat-view/SimpleChatPanelProvider.ts |  2 +-
 vscode/src/edit/provider.ts                   |  2 +-
 8 files changed, 164 insertions(+), 53 deletions(-)
 create mode 100644 lib/shared/src/models/index.test.ts
 create mode 100644 lib/shared/src/models/utils.test.ts

diff --git a/lib/shared/src/models/dotcom.ts b/lib/shared/src/models/dotcom.ts
index b085ef9aabd..6f705364499 100644
--- a/lib/shared/src/models/dotcom.ts
+++ b/lib/shared/src/models/dotcom.ts
@@ -1,9 +1,7 @@
 import type { ModelProvider } from '.'
+import { DEFAULT_CHAT_MODEL_TOKEN_LIMIT, DEFAULT_FAST_MODEL_TOKEN_LIMIT } from '../prompt/constants'
 import { ModelUsage } from './types'
 
-export const DEFAULT_CHAT_MODEL_TOKEN_LIMIT = 28000 // 7000 tokens * 4 bytes per token
-export const DEFAULT_FAST_MODEL_TOKEN_LIMIT = 10000 // 4,096 tokens * 4 bytes per token
-
 // The models must first be added to the custom chat models list in https://sourcegraph.com/github.com/sourcegraph/sourcegraph/-/blob/internal/completions/httpapi/chat.go?L48-51
 export const DEFAULT_DOT_COM_MODELS = [
     {
@@ -13,7 +11,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: true,
         codyProOnly: false,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 2.1',
@@ -22,7 +20,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude Instant',
@@ -31,7 +29,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 3 Haiku',
@@ -40,7 +38,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 3 Sonnet',
@@ -49,7 +47,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Claude 3 Opus',
@@ -58,7 +56,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'GPT-3.5 Turbo',
@@ -67,7 +65,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'GPT-4 Turbo Preview',
@@ -76,7 +74,7 @@ export const DEFAULT_DOT_COM_MODELS = [
         default: false,
         codyProOnly: true,
         usage: [ModelUsage.Chat, ModelUsage.Edit],
-        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
     {
         title: 'Mixtral 8x7B',
@@ -86,6 +84,6 @@ export const DEFAULT_DOT_COM_MODELS = [
         codyProOnly: true,
         // TODO: Improve prompt for Mixtral + Edit to see if we can use it there too.
         usage: [ModelUsage.Chat],
-        contextWindow: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
+        maxToken: DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
     },
 ] as const satisfies ModelProvider[]
diff --git a/lib/shared/src/models/index.test.ts b/lib/shared/src/models/index.test.ts
new file mode 100644
index 00000000000..a91fa81a704
--- /dev/null
+++ b/lib/shared/src/models/index.test.ts
@@ -0,0 +1,35 @@
+import { beforeAll, describe, expect, it } from 'vitest'
+import { ModelProvider } from '../models/index'
+import { DEFAULT_FAST_MODEL_TOKEN_LIMIT, tokensToChars } from '../prompt/constants'
+import { DOTCOM_URL } from '../sourcegraph-api/environments'
+import { DEFAULT_DOT_COM_MODELS } from './dotcom'
+import { ModelUsage } from './types'
+
+describe('getMaxCharsByModel', () => {
+    beforeAll(() => {
+        ModelProvider.get(ModelUsage.Chat, DOTCOM_URL.toString())
+    })
+
+    it('returns default token limit for unknown model', () => {
+        const maxChars = ModelProvider.getMaxCharsByModel('unknown-model')
+        expect(maxChars).toEqual(tokensToChars(DEFAULT_FAST_MODEL_TOKEN_LIMIT))
+    })
+
+    it('returns max token limit for known chat model', () => {
+        const maxChars = ModelProvider.getMaxCharsByModel(DEFAULT_DOT_COM_MODELS[0].model)
+        expect(maxChars).toEqual(tokensToChars(DEFAULT_DOT_COM_MODELS[0].maxToken))
+    })
+
+    it('returns default token limit for unknown model - Enterprise user', () => {
+        ModelProvider.get(ModelUsage.Chat, 'https://example.com')
+        const maxChars = ModelProvider.getMaxCharsByModel('unknown-model')
+        expect(maxChars).toEqual(tokensToChars(DEFAULT_FAST_MODEL_TOKEN_LIMIT))
+    })
+
+    it('returns max token limit for known model - Enterprise user', () => {
+        ModelProvider.get(ModelUsage.Chat, 'https://example.com')
+        ModelProvider.add(new ModelProvider('model-with-limit', [ModelUsage.Chat], 200))
+        const maxChars = ModelProvider.getMaxCharsByModel('model-with-limit')
+        expect(maxChars).toEqual(tokensToChars(200))
+    })
+})
diff --git a/lib/shared/src/models/index.ts b/lib/shared/src/models/index.ts
index 768da4339a3..f617067889a 100644
--- a/lib/shared/src/models/index.ts
+++ b/lib/shared/src/models/index.ts
@@ -1,13 +1,14 @@
 import { logError } from '../logger'
 import { OLLAMA_DEFAULT_URL } from '../ollama'
-import { isDotCom } from '../sourcegraph-api/environments'
 import {
-    DEFAULT_CHAT_MODEL_TOKEN_LIMIT,
-    DEFAULT_DOT_COM_MODELS,
+    DEFAULT_FAST_MODEL_CHARS_LIMIT,
     DEFAULT_FAST_MODEL_TOKEN_LIMIT,
-} from './dotcom'
+    tokensToChars,
+} from '../prompt/constants'
+import { isDotCom } from '../sourcegraph-api/environments'
+import { DEFAULT_DOT_COM_MODELS } from './dotcom'
 import { ModelUsage } from './types'
-import { getProviderName } from './utils'
+import { getModelInfo } from './utils'
 
 /**
  * ModelProvider manages available chat and edit models.
@@ -17,27 +18,34 @@ import { getProviderName } from './utils'
 export class ModelProvider {
     public default = false
     public codyProOnly = false
+    // The name of the provider of the model, e.g. "Anthropic"
     public provider: string
+    // The title of the model, e.g. "Claude 2.0"
     public readonly title: string
-    public readonly contextWindow: number
 
     constructor(
+        // The model id that includes the provider name & the model name,
+        // e.g. "anthropic/claude-2.0"
         public readonly model: string,
+        // The usage of the model, e.g. chat or edit.
         public readonly usage: ModelUsage[],
-        tokenLimit?: number
+        // The maximum number of tokens that can be processed by the model in a single request.
+        // NOTE: A token is equivalent to 4 characters/bytes.
+        public readonly maxToken: number = DEFAULT_FAST_MODEL_TOKEN_LIMIT
     ) {
-        const splittedModel = model.split('/')
-        this.provider = getProviderName(splittedModel[0])
-        this.title = splittedModel[1]?.replaceAll('-', ' ')
+        const { provider, title } = getModelInfo(model)
+        this.provider = provider
+        this.title = title
         this.default = true
-        this.contextWindow = tokenLimit ? tokenLimit * 4 : DEFAULT_FAST_MODEL_TOKEN_LIMIT
     }
 
-    // Providers available for non-dotcom instances
-    private static privateProviders: Map<string, ModelProvider> = new Map()
-    // Providers available for dotcom instances
-    private static dotComProviders: ModelProvider[] = DEFAULT_DOT_COM_MODELS
-    // Providers available from local ollama instances
+    /**
+     * Providers available on the user's instance
+     */
+    private static primaryProviders: ModelProvider[] = DEFAULT_DOT_COM_MODELS
+    /**
+     * Providers available from local ollama instances
+     */
     private static ollamaProvidersEnabled = false
     private static ollamaProviders: ModelProvider[] = []
 
@@ -68,7 +76,11 @@ export class ModelProvider {
                     const models = new Set<ModelProvider>()
                     for (const model of data.models) {
                         const name = `ollama/${model.model}`
-                        const newModel = new ModelProvider(name, [ModelUsage.Chat, ModelUsage.Edit])
+                        const newModel = new ModelProvider(
+                            name,
+                            [ModelUsage.Chat, ModelUsage.Edit],
+                            DEFAULT_FAST_MODEL_CHARS_LIMIT
+                        )
                         models.add(newModel)
                     }
                     ModelProvider.ollamaProviders = Array.from(models)
@@ -88,11 +100,8 @@ export class ModelProvider {
      * made available for use.
      */
     public static add(provider: ModelProvider): void {
-        // private instances can only support 1 provider atm
-        if (ModelProvider.privateProviders.size) {
-            ModelProvider.privateProviders.clear()
-        }
-        ModelProvider.privateProviders.set(provider.model.trim(), provider)
+        // NOTE: private instances can only support 1 provider atm
+        ModelProvider.primaryProviders = [provider]
     }
 
     /**
@@ -106,18 +115,13 @@ export class ModelProvider {
         currentModel?: string
     ): ModelProvider[] {
         const isDotComUser = !endpoint || (endpoint && isDotCom(endpoint))
-        const models = (
-            isDotComUser
-                ? ModelProvider.dotComProviders
-                : Array.from(ModelProvider.privateProviders.values())
-        )
+        if (isDotComUser) {
+            ModelProvider.primaryProviders = DEFAULT_DOT_COM_MODELS
+        }
+        const models = ModelProvider.primaryProviders
             .concat(ModelProvider.ollamaProviders)
             .filter(model => model.usage.includes(type))
 
-        if (!isDotComUser) {
-            return models
-        }
-
         // Set the current model as default
         return models.map(model => {
             return {
@@ -127,11 +131,15 @@ export class ModelProvider {
         })
     }
 
-    public static getContextWindow(modelID: string): number {
-        return (
-            ModelProvider.privateProviders.get(modelID)?.contextWindow ||
-            ModelProvider.dotComProviders.find(model => model.model === modelID)?.contextWindow ||
-            DEFAULT_CHAT_MODEL_TOKEN_LIMIT
-        )
+    /**
+     * Finds the model provider with the given model ID and returns its characters limit.
+     * The limit is calculated based on the max number of tokens the model can process.
+     * E.g. 7000 tokens * 4 characters/token = 28000 characters
+     */
+    public static getMaxCharsByModel(modelID: string): number {
+        const model = ModelProvider.primaryProviders
+            .concat(ModelProvider.ollamaProviders)
+            .find(m => m.model === modelID)
+        return tokensToChars(model?.maxToken || DEFAULT_FAST_MODEL_TOKEN_LIMIT)
     }
 }
diff --git a/lib/shared/src/models/utils.test.ts b/lib/shared/src/models/utils.test.ts
new file mode 100644
index 00000000000..a7d0b519192
--- /dev/null
+++ b/lib/shared/src/models/utils.test.ts
@@ -0,0 +1,36 @@
+import { describe, expect, it } from 'vitest'
+import { getModelInfo } from './utils'
+
+describe('getModelInfo', () => {
+    it('splits model ID and returns provider and title', () => {
+        const result = getModelInfo('Anthropic/Claude 2.0')
+        expect(result).toEqual({
+            provider: 'Anthropic',
+            title: 'Claude 2.0',
+        })
+    })
+
+    it('handles model ID without title', () => {
+        const result = getModelInfo('Anthropic/')
+        expect(result).toEqual({
+            provider: 'Anthropic',
+            title: '',
+        })
+    })
+
+    it('replaces dashes in title with spaces', () => {
+        const result = getModelInfo('Example/model-with-dashes')
+        expect(result).toEqual({
+            provider: 'Example',
+            title: 'model with dashes',
+        })
+    })
+
+    it('handles model ID with multiple dashes', () => {
+        const result = getModelInfo('fireworks/accounts/fireworks/models/mixtral-8x7b-instruct')
+        expect(result).toEqual({
+            provider: 'fireworks',
+            title: 'mixtral 8x7b instruct',
+        })
+    })
+})
diff --git a/lib/shared/src/models/utils.ts b/lib/shared/src/models/utils.ts
index 6f41cfb6099..c81b2753b94 100644
--- a/lib/shared/src/models/utils.ts
+++ b/lib/shared/src/models/utils.ts
@@ -15,3 +15,18 @@ export function getProviderName(name: string): string {
 export function supportsFastPath(model: string): boolean {
     return model?.startsWith('anthropic/claude-3')
 }
+
+/**
+ * Gets the provider and title from a model ID string.
+ */
+export function getModelInfo(modelID: string): {
+    provider: string
+    title: string
+} {
+    const splittedModel = modelID.split('/')
+    // The name of provider of the model, e.g. "Anthropic"
+    const provider = getProviderName(splittedModel[0])
+    // The title/name of the model, e.g. "Claude 2.0"
+    const title = (splittedModel.pop() || splittedModel[1]).replaceAll('-', ' ')
+    return { provider, title }
+}
diff --git a/lib/shared/src/prompt/constants.ts b/lib/shared/src/prompt/constants.ts
index 135cb4fff8c..e8b11f7d9ff 100644
--- a/lib/shared/src/prompt/constants.ts
+++ b/lib/shared/src/prompt/constants.ts
@@ -1,6 +1,9 @@
 export const ANSWER_TOKENS = 1000
 export const MAX_HUMAN_INPUT_TOKENS = 1000
 export const MAX_CURRENT_FILE_TOKENS = 1000
+/**
+ * A token is equivalent to 4 characters/bytes.
+ */
 export const CHARS_PER_TOKEN = 4
 export const SURROUNDING_LINES = 50
 export const NUM_CODE_RESULTS = 12
@@ -8,6 +11,22 @@ export const NUM_TEXT_RESULTS = 3
 
 export const MAX_BYTES_PER_FILE = 4096
 
-export function tokensToChars(tokens: number): number {
-    return tokens * CHARS_PER_TOKEN
+// CHAT MODEL TOKEN LIMITS
+export const DEFAULT_CHAT_MODEL_TOKEN_LIMIT = 7000
+export const DEFAULT_FAST_MODEL_TOKEN_LIMIT = 4096
+export const DEFAULT_CHAT_MODEL_CHARS_LIMIT = tokensToChars(DEFAULT_CHAT_MODEL_TOKEN_LIMIT)
+export const DEFAULT_FAST_MODEL_CHARS_LIMIT = tokensToChars(DEFAULT_FAST_MODEL_TOKEN_LIMIT)
+
+/**
+ * Calculate the number of characters from the number of tokens.
+ */
+export function tokensToChars(tokenCount: number): number {
+    return tokenCount * CHARS_PER_TOKEN
+}
+
+/**
+ * Calculate the number of token from the number of characters.
+ */
+export function charsToTokens(charsCount: number): number {
+    return charsCount / CHARS_PER_TOKEN
 }
diff --git a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
index d1f651ffa16..e5d0a4a63db 100644
--- a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
+++ b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
@@ -781,7 +781,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
         prompter: IPrompter,
         sendTelemetry?: (contextSummary: any) => void
     ): Promise<Message[]> {
-        const maxChars = ModelProvider.getContextWindow(this.chatModel.modelID)
+        const maxChars = ModelProvider.getMaxCharsByModel(this.chatModel.modelID)
         const { prompt, newContextUsed } = await prompter.makePrompt(this.chatModel, maxChars)
 
         // Update UI based on prompt construction
diff --git a/vscode/src/edit/provider.ts b/vscode/src/edit/provider.ts
index 8ab54d0e181..e6ba83ccd3d 100644
--- a/vscode/src/edit/provider.ts
+++ b/vscode/src/edit/provider.ts
@@ -46,7 +46,7 @@ export class EditProvider {
         return wrapInActiveSpan('command.edit.start', async span => {
             this.config.controller.startTask(this.config.task)
             const model = this.config.task.model
-            const contextWindow = ModelProvider.getContextWindow(model)
+            const contextWindow = ModelProvider.getMaxCharsByModel(model)
             const {
                 messages,
                 stopSequences,

From 50d36301daf7e617014f5829b485c197b8651b5a Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 09:46:05 -0700
Subject: [PATCH 06/14] update bindings

---
 .../com/sourcegraph/cody/protocol_generated/ModelProvider.kt  | 4 +---
 lib/shared/src/models/utils.test.ts                           | 4 ++--
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
index 5bd030effd7..7e448a9e25d 100644
--- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
+++ b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
@@ -6,9 +6,7 @@ data class ModelProvider(
   val codyProOnly: Boolean? = null,
   val provider: String? = null,
   val title: String? = null,
-  val contextWindowSizeChars: Int? = null,
-  val privateProviders: Map<String, ModelProvider>? = null,
-  val dotComProviders: List<ModelProvider>? = null,
+  val primaryProviders: List<ModelProvider>? = null,
   val ollamaProvidersEnabled: Boolean? = null,
   val ollamaProviders: List<ModelProvider>? = null,
 )
diff --git a/lib/shared/src/models/utils.test.ts b/lib/shared/src/models/utils.test.ts
index a7d0b519192..1d47899e27f 100644
--- a/lib/shared/src/models/utils.test.ts
+++ b/lib/shared/src/models/utils.test.ts
@@ -19,9 +19,9 @@ describe('getModelInfo', () => {
     })
 
     it('replaces dashes in title with spaces', () => {
-        const result = getModelInfo('Example/model-with-dashes')
+        const result = getModelInfo('example/model-with-dashes')
         expect(result).toEqual({
-            provider: 'Example',
+            provider: 'example',
             title: 'model with dashes',
         })
     })

From 48bad9afb9790ee388777939a36a26b5bfa26d5f Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 09:51:18 -0700
Subject: [PATCH 07/14] rename function

---
 vscode/src/chat/chat-view/ChatManager.ts             | 4 ++--
 vscode/src/chat/chat-view/ChatPanelsManager.ts       | 4 ++--
 vscode/src/chat/chat-view/SimpleChatPanelProvider.ts | 4 ++--
 vscode/src/edit/utils/edit-models.ts                 | 4 ++--
 vscode/src/models/utilts.ts                          | 9 ++++++++-
 5 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/vscode/src/chat/chat-view/ChatManager.ts b/vscode/src/chat/chat-view/ChatManager.ts
index 2186723bc06..ef0ce668862 100644
--- a/vscode/src/chat/chat-view/ChatManager.ts
+++ b/vscode/src/chat/chat-view/ChatManager.ts
@@ -21,7 +21,7 @@ import { telemetryRecorder } from '../../services/telemetry-v2'
 import type { ExecuteChatArguments } from '../../commands/execute/ask'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
-import { addEnterpriseChatModel } from '../../models/utilts'
+import { setModelProviders } from '../../models/utilts'
 import { ChatPanelsManager } from './ChatPanelsManager'
 import { SidebarViewController, type SidebarViewOptions } from './SidebarViewController'
 import type { ChatSession, SimpleChatPanelProvider } from './SimpleChatPanelProvider'
@@ -99,7 +99,7 @@ export class ChatManager implements vscode.Disposable {
     }
 
     public async syncAuthStatus(authStatus: AuthStatus): Promise<void> {
-        addEnterpriseChatModel(authStatus)
+        setModelProviders(authStatus)
 
         await this.chatPanelsManager.syncAuthStatus(authStatus)
     }
diff --git a/vscode/src/chat/chat-view/ChatPanelsManager.ts b/vscode/src/chat/chat-view/ChatPanelsManager.ts
index 36aae8d01ee..56eb445e6bf 100644
--- a/vscode/src/chat/chat-view/ChatPanelsManager.ts
+++ b/vscode/src/chat/chat-view/ChatPanelsManager.ts
@@ -23,7 +23,7 @@ import type { ExtensionMessage } from '../protocol'
 import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
-import { addEnterpriseChatModel } from '../../models/utilts'
+import { setModelProviders } from '../../models/utilts'
 import { chatHistory } from './ChatHistoryManager'
 import { CodyChatPanelViewType } from './ChatManager'
 import type { SidebarViewOptions } from './SidebarViewController'
@@ -216,7 +216,7 @@ export class ChatPanelsManager implements vscode.Disposable {
     private createProvider(): SimpleChatPanelProvider {
         const authProvider = this.options.authProvider
         const authStatus = authProvider.getAuthStatus()
-        addEnterpriseChatModel(authStatus)
+        setModelProviders(authStatus)
         const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint)
         const isConsumer = authProvider.getAuthStatus().isDotCom
 
diff --git a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
index e5d0a4a63db..d3cbfc3eb6c 100644
--- a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
+++ b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
@@ -58,7 +58,7 @@ import type { Repo } from '../../context/repo-fetcher'
 import type { RemoteRepoPicker } from '../../context/repo-picker'
 import type { ContextRankingController } from '../../local-context/context-ranking'
 import { chatModel } from '../../models'
-import { addEnterpriseChatModel } from '../../models/utilts'
+import { setModelProviders } from '../../models/utilts'
 import { recordExposedExperimentsToSpan } from '../../services/open-telemetry/utils'
 import type { MessageErrorType } from '../MessageProvider'
 import { getChatContextItemsForMention } from '../context/chatContext'
@@ -730,7 +730,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
         if (!authStatus?.isLoggedIn) {
             return
         }
-        addEnterpriseChatModel(authStatus)
+        setModelProviders(authStatus)
         const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint, this.chatModel.modelID)
 
         void this.postMessage({
diff --git a/vscode/src/edit/utils/edit-models.ts b/vscode/src/edit/utils/edit-models.ts
index 570a1269b4c..8908f0162c7 100644
--- a/vscode/src/edit/utils/edit-models.ts
+++ b/vscode/src/edit/utils/edit-models.ts
@@ -1,10 +1,10 @@
 import { type AuthStatus, ModelProvider } from '@sourcegraph/cody-shared'
 import { type EditModel, ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
-import { addEnterpriseChatModel } from '../../models/utilts'
+import { setModelProviders } from '../../models/utilts'
 import type { EditIntent } from '../types'
 
 export function getEditModelsForUser(authStatus: AuthStatus): ModelProvider[] {
-    addEnterpriseChatModel(authStatus)
+    setModelProviders(authStatus)
     return ModelProvider.get(ModelUsage.Edit, authStatus.endpoint)
 }
 
diff --git a/vscode/src/models/utilts.ts b/vscode/src/models/utilts.ts
index ba3cf4736fd..1a77c53ae37 100644
--- a/vscode/src/models/utilts.ts
+++ b/vscode/src/models/utilts.ts
@@ -2,7 +2,14 @@ import { type AuthStatus, ModelProvider } from '@sourcegraph/cody-shared'
 import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import * as vscode from 'vscode'
 
-export function addEnterpriseChatModel(authStatus: AuthStatus): void {
+/**
+ * Sets the model providers based on the authentication status.
+ *
+ * If a chat model is configured to overwrite, it will add a model provider for that model.
+ * The token limit for the provider will use the configured limit,
+ * or fallback to the limit from the authentication status if not configured.
+ */
+export function setModelProviders(authStatus: AuthStatus): void {
     // In enterprise mode, we let the sg instance dictate the token limits and allow users to
     // overwrite it locally (for debugging purposes).
     //

From 424eac95175bf0f35d8ac6778996ad8506094675 Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 10:03:27 -0700
Subject: [PATCH 08/14] clean up

---
 lib/shared/src/models/utils.ts | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/lib/shared/src/models/utils.ts b/lib/shared/src/models/utils.ts
index c81b2753b94..e1099eb62cd 100644
--- a/lib/shared/src/models/utils.ts
+++ b/lib/shared/src/models/utils.ts
@@ -23,10 +23,8 @@ export function getModelInfo(modelID: string): {
     provider: string
     title: string
 } {
-    const splittedModel = modelID.split('/')
-    // The name of provider of the model, e.g. "Anthropic"
-    const provider = getProviderName(splittedModel[0])
-    // The title/name of the model, e.g. "Claude 2.0"
-    const title = (splittedModel.pop() || splittedModel[1]).replaceAll('-', ' ')
+    const [providerID, ...rest] = modelID.split('/')
+    const provider = getProviderName(providerID)
+    const title = (rest.at(-1) || '').replace(/-/g, ' ')
     return { provider, title }
 }

From fc5c9ef6cfc6ee50ca8b1a7d7f81925ada7eb83a Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 23:28:48 -0700
Subject: [PATCH 09/14] simplify sync logic

---
 agent/src/agent.ts                            |  2 +-
 lib/shared/src/models/index.test.ts           |  8 +++---
 lib/shared/src/models/index.ts                | 26 +++++--------------
 vscode/src/chat/chat-view/ChatManager.ts      |  2 +-
 .../src/chat/chat-view/ChatPanelsManager.ts   |  5 +---
 .../chat/chat-view/SimpleChatPanelProvider.ts |  4 +--
 vscode/src/edit/utils/edit-models.ts          |  4 +--
 vscode/src/models/{utilts.ts => utils.ts}     | 13 +++++++---
 .../ChatModelDropdownMenu.story.tsx           |  2 +-
 9 files changed, 26 insertions(+), 40 deletions(-)
 rename vscode/src/models/{utilts.ts => utils.ts} (78%)

diff --git a/agent/src/agent.ts b/agent/src/agent.ts
index 8318dc69c2c..7a1b1d55ee2 100644
--- a/agent/src/agent.ts
+++ b/agent/src/agent.ts
@@ -821,7 +821,7 @@ export class Agent extends MessageHandler implements ExtensionClient {
         })
 
         this.registerAuthenticatedRequest('chat/restore', async ({ modelID, messages, chatID }) => {
-            const theModel = modelID ? modelID : ModelProvider.get(ModelUsage.Chat).at(0)?.model
+            const theModel = modelID ? modelID : ModelProvider.getProviders(ModelUsage.Chat).at(0)?.model
             if (!theModel) {
                 throw new Error('No default chat model found')
             }
diff --git a/lib/shared/src/models/index.test.ts b/lib/shared/src/models/index.test.ts
index a91fa81a704..15b3311bf54 100644
--- a/lib/shared/src/models/index.test.ts
+++ b/lib/shared/src/models/index.test.ts
@@ -7,7 +7,7 @@ import { ModelUsage } from './types'
 
 describe('getMaxCharsByModel', () => {
     beforeAll(() => {
-        ModelProvider.get(ModelUsage.Chat, DOTCOM_URL.toString())
+        ModelProvider.getProviders(ModelUsage.Chat, DOTCOM_URL.toString())
     })
 
     it('returns default token limit for unknown model', () => {
@@ -21,14 +21,14 @@ describe('getMaxCharsByModel', () => {
     })
 
     it('returns default token limit for unknown model - Enterprise user', () => {
-        ModelProvider.get(ModelUsage.Chat, 'https://example.com')
+        ModelProvider.getProviders(ModelUsage.Chat, 'https://example.com')
         const maxChars = ModelProvider.getMaxCharsByModel('unknown-model')
         expect(maxChars).toEqual(tokensToChars(DEFAULT_FAST_MODEL_TOKEN_LIMIT))
     })
 
     it('returns max token limit for known model - Enterprise user', () => {
-        ModelProvider.get(ModelUsage.Chat, 'https://example.com')
-        ModelProvider.add(new ModelProvider('model-with-limit', [ModelUsage.Chat], 200))
+        ModelProvider.getProviders(ModelUsage.Chat, 'https://example.com')
+        ModelProvider.setProviders([new ModelProvider('model-with-limit', [ModelUsage.Chat], 200)])
         const maxChars = ModelProvider.getMaxCharsByModel('model-with-limit')
         expect(maxChars).toEqual(tokensToChars(200))
     })
diff --git a/lib/shared/src/models/index.ts b/lib/shared/src/models/index.ts
index f617067889a..05091ba2e09 100644
--- a/lib/shared/src/models/index.ts
+++ b/lib/shared/src/models/index.ts
@@ -5,7 +5,6 @@ import {
     DEFAULT_FAST_MODEL_TOKEN_LIMIT,
     tokensToChars,
 } from '../prompt/constants'
-import { isDotCom } from '../sourcegraph-api/environments'
 import { DEFAULT_DOT_COM_MODELS } from './dotcom'
 import { ModelUsage } from './types'
 import { getModelInfo } from './utils'
@@ -95,35 +94,24 @@ export class ModelProvider {
     }
 
     /**
-     * Adds a new model provider, instantiated from the given model string,
-     * to the internal providers set. This allows new models to be added and
-     * made available for use.
+     * Sets the primary model providers.
+     * NOTE: private instances can only support 1 provider atm
      */
-    public static add(provider: ModelProvider): void {
-        // NOTE: private instances can only support 1 provider atm
-        ModelProvider.primaryProviders = [provider]
+    public static setProviders(providers: ModelProvider[]): void {
+        ModelProvider.primaryProviders = providers
     }
 
     /**
-     * Gets the model providers based on the endpoint and current model.
-     * If endpoint is a dotcom endpoint, returns dotComProviders with ollama providers.
+     * Get the list of the primary models providers with local models.
      * If currentModel is provided, sets it as the default model.
      */
-    public static get(
-        type: ModelUsage,
-        endpoint?: string | null,
-        currentModel?: string
-    ): ModelProvider[] {
-        const isDotComUser = !endpoint || (endpoint && isDotCom(endpoint))
-        if (isDotComUser) {
-            ModelProvider.primaryProviders = DEFAULT_DOT_COM_MODELS
-        }
+    public static getProviders(type: ModelUsage, currentModel?: string): ModelProvider[] {
         const models = ModelProvider.primaryProviders
             .concat(ModelProvider.ollamaProviders)
             .filter(model => model.usage.includes(type))
 
         // Set the current model as default
-        return models.map(model => {
+        return models?.map(model => {
             return {
                 ...model,
                 default: model.model === currentModel,
diff --git a/vscode/src/chat/chat-view/ChatManager.ts b/vscode/src/chat/chat-view/ChatManager.ts
index ef0ce668862..002dfb9a932 100644
--- a/vscode/src/chat/chat-view/ChatManager.ts
+++ b/vscode/src/chat/chat-view/ChatManager.ts
@@ -21,7 +21,7 @@ import { telemetryRecorder } from '../../services/telemetry-v2'
 import type { ExecuteChatArguments } from '../../commands/execute/ask'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
-import { setModelProviders } from '../../models/utilts'
+import { setModelProviders } from '../../models/utils'
 import { ChatPanelsManager } from './ChatPanelsManager'
 import { SidebarViewController, type SidebarViewOptions } from './SidebarViewController'
 import type { ChatSession, SimpleChatPanelProvider } from './SimpleChatPanelProvider'
diff --git a/vscode/src/chat/chat-view/ChatPanelsManager.ts b/vscode/src/chat/chat-view/ChatPanelsManager.ts
index 56eb445e6bf..b0e3d254f00 100644
--- a/vscode/src/chat/chat-view/ChatPanelsManager.ts
+++ b/vscode/src/chat/chat-view/ChatPanelsManager.ts
@@ -23,7 +23,6 @@ import type { ExtensionMessage } from '../protocol'
 import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
-import { setModelProviders } from '../../models/utilts'
 import { chatHistory } from './ChatHistoryManager'
 import { CodyChatPanelViewType } from './ChatManager'
 import type { SidebarViewOptions } from './SidebarViewController'
@@ -215,9 +214,7 @@ export class ChatPanelsManager implements vscode.Disposable {
      */
     private createProvider(): SimpleChatPanelProvider {
         const authProvider = this.options.authProvider
-        const authStatus = authProvider.getAuthStatus()
-        setModelProviders(authStatus)
-        const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint)
+        const models = ModelProvider.getProviders(ModelUsage.Chat)
         const isConsumer = authProvider.getAuthStatus().isDotCom
 
         return new SimpleChatPanelProvider({
diff --git a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
index d3cbfc3eb6c..9cb0ca8f270 100644
--- a/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
+++ b/vscode/src/chat/chat-view/SimpleChatPanelProvider.ts
@@ -58,7 +58,6 @@ import type { Repo } from '../../context/repo-fetcher'
 import type { RemoteRepoPicker } from '../../context/repo-picker'
 import type { ContextRankingController } from '../../local-context/context-ranking'
 import { chatModel } from '../../models'
-import { setModelProviders } from '../../models/utilts'
 import { recordExposedExperimentsToSpan } from '../../services/open-telemetry/utils'
 import type { MessageErrorType } from '../MessageProvider'
 import { getChatContextItemsForMention } from '../context/chatContext'
@@ -730,8 +729,7 @@ export class SimpleChatPanelProvider implements vscode.Disposable, ChatSession {
         if (!authStatus?.isLoggedIn) {
             return
         }
-        setModelProviders(authStatus)
-        const models = ModelProvider.get(ModelUsage.Chat, authStatus.endpoint, this.chatModel.modelID)
+        const models = ModelProvider.getProviders(ModelUsage.Chat, this.chatModel.modelID)
 
         void this.postMessage({
             type: 'chatModels',
diff --git a/vscode/src/edit/utils/edit-models.ts b/vscode/src/edit/utils/edit-models.ts
index 8908f0162c7..b20292716b1 100644
--- a/vscode/src/edit/utils/edit-models.ts
+++ b/vscode/src/edit/utils/edit-models.ts
@@ -1,11 +1,9 @@
 import { type AuthStatus, ModelProvider } from '@sourcegraph/cody-shared'
 import { type EditModel, ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
-import { setModelProviders } from '../../models/utilts'
 import type { EditIntent } from '../types'
 
 export function getEditModelsForUser(authStatus: AuthStatus): ModelProvider[] {
-    setModelProviders(authStatus)
-    return ModelProvider.get(ModelUsage.Edit, authStatus.endpoint)
+    return ModelProvider.getProviders(ModelUsage.Edit)
 }
 
 export function getOverridenModelForIntent(intent: EditIntent, currentModel: EditModel): EditModel {
diff --git a/vscode/src/models/utilts.ts b/vscode/src/models/utils.ts
similarity index 78%
rename from vscode/src/models/utilts.ts
rename to vscode/src/models/utils.ts
index 1a77c53ae37..b3425662659 100644
--- a/vscode/src/models/utilts.ts
+++ b/vscode/src/models/utils.ts
@@ -1,4 +1,5 @@
-import { type AuthStatus, ModelProvider } from '@sourcegraph/cody-shared'
+import { type AuthStatus, ModelProvider, isDotCom } from '@sourcegraph/cody-shared'
+import { DEFAULT_DOT_COM_MODELS } from '@sourcegraph/cody-shared/src/models/dotcom'
 import { ModelUsage } from '@sourcegraph/cody-shared/src/models/types'
 import * as vscode from 'vscode'
 
@@ -10,6 +11,9 @@ import * as vscode from 'vscode'
  * or fallback to the limit from the authentication status if not configured.
  */
 export function setModelProviders(authStatus: AuthStatus): void {
+    if (authStatus.endpoint && isDotCom(authStatus.endpoint)) {
+        ModelProvider.setProviders(DEFAULT_DOT_COM_MODELS)
+    }
     // In enterprise mode, we let the sg instance dictate the token limits and allow users to
     // overwrite it locally (for debugging purposes).
     //
@@ -20,13 +24,14 @@ export function setModelProviders(authStatus: AuthStatus): void {
         const codyConfig = vscode.workspace.getConfiguration('cody')
         const tokenLimitConfig = codyConfig.get<number>('provider.limit.prompt')
         const tokenLimit = tokenLimitConfig ?? authStatus.configOverwrites?.chatModelMaxTokens
-        ModelProvider.add(
+        ModelProvider.setProviders([
             new ModelProvider(
                 authStatus.configOverwrites.chatModel,
                 // TODO: Add configOverwrites.editModel for separate edit support
                 [ModelUsage.Chat, ModelUsage.Edit],
                 tokenLimit
-            )
-        )
+            ),
+        ])
+        return
     }
 }
diff --git a/vscode/webviews/Components/ChatModelDropdownMenu.story.tsx b/vscode/webviews/Components/ChatModelDropdownMenu.story.tsx
index 924f45b743f..21061bae7a4 100644
--- a/vscode/webviews/Components/ChatModelDropdownMenu.story.tsx
+++ b/vscode/webviews/Components/ChatModelDropdownMenu.story.tsx
@@ -12,7 +12,7 @@ const meta: Meta<typeof ChatModelDropdownMenu> = {
     component: ChatModelDropdownMenu,
     decorators: [VSCodeStoryDecorator],
     args: {
-        models: ModelProvider.get(ModelUsage.Chat, String(DOTCOM_URL)),
+        models: ModelProvider.getProviders(ModelUsage.Chat, String(DOTCOM_URL)),
         disabled: false,
     },
     parameters: {

From 967dcdd4e0c57adf8b9a2a3d872dc33a12dfab74 Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 23:31:01 -0700
Subject: [PATCH 10/14] early return

---
 vscode/src/models/utils.ts | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/vscode/src/models/utils.ts b/vscode/src/models/utils.ts
index b3425662659..45f3891ef75 100644
--- a/vscode/src/models/utils.ts
+++ b/vscode/src/models/utils.ts
@@ -13,7 +13,9 @@ import * as vscode from 'vscode'
 export function setModelProviders(authStatus: AuthStatus): void {
     if (authStatus.endpoint && isDotCom(authStatus.endpoint)) {
         ModelProvider.setProviders(DEFAULT_DOT_COM_MODELS)
+        return
     }
+
     // In enterprise mode, we let the sg instance dictate the token limits and allow users to
     // overwrite it locally (for debugging purposes).
     //
@@ -32,6 +34,5 @@ export function setModelProviders(authStatus: AuthStatus): void {
                 tokenLimit
             ),
         ])
-        return
     }
 }

From f8cb7cfddafdd2b5dab67bc2fabd5178a71fb1af Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 23:51:34 -0700
Subject: [PATCH 11/14] clean up

---
 lib/shared/src/models/index.ts | 79 +++++++++++++++-------------------
 1 file changed, 34 insertions(+), 45 deletions(-)

diff --git a/lib/shared/src/models/index.ts b/lib/shared/src/models/index.ts
index 05091ba2e09..c13a947f036 100644
--- a/lib/shared/src/models/index.ts
+++ b/lib/shared/src/models/index.ts
@@ -35,7 +35,6 @@ export class ModelProvider {
         const { provider, title } = getModelInfo(model)
         this.provider = provider
         this.title = title
-        this.default = true
     }
 
     /**
@@ -45,52 +44,10 @@ export class ModelProvider {
     /**
      * Providers available from local ollama instances
      */
-    private static ollamaProvidersEnabled = false
     private static ollamaProviders: ModelProvider[] = []
-
-    public static onConfigChange(enableOllamaModels: boolean): void {
-        ModelProvider.ollamaProvidersEnabled = enableOllamaModels
-        ModelProvider.ollamaProviders = []
-        if (enableOllamaModels) {
-            ModelProvider.getLocalOllamaModels()
-        }
-    }
-
-    /**
-     * Fetches available Ollama models from the local Ollama server
-     * and adds them to the list of ollama providers.
-     */
-    public static getLocalOllamaModels(): void {
-        const isAgentTesting = process.env.CODY_SHIM_TESTING === 'true'
+    public static async onConfigChange(enableOllamaModels: boolean): Promise<void> {
         // Only fetch local models if user has enabled the config
-        if (isAgentTesting || !ModelProvider.ollamaProvidersEnabled) {
-            return
-        }
-        // TODO (bee) watch file change to determine if a new model is added
-        // to eliminate the needs of restarting the extension to get the new models
-        fetch(new URL('/api/tags', OLLAMA_DEFAULT_URL).href)
-            .then(response => response.json())
-            .then(
-                data => {
-                    const models = new Set<ModelProvider>()
-                    for (const model of data.models) {
-                        const name = `ollama/${model.model}`
-                        const newModel = new ModelProvider(
-                            name,
-                            [ModelUsage.Chat, ModelUsage.Edit],
-                            DEFAULT_FAST_MODEL_CHARS_LIMIT
-                        )
-                        models.add(newModel)
-                    }
-                    ModelProvider.ollamaProviders = Array.from(models)
-                },
-                error => {
-                    const fetchFailedErrors = ['Failed to fetch', 'fetch failed']
-                    const isFetchFailed = fetchFailedErrors.some(err => error.toString().includes(err))
-                    const serverErrorMsg = 'Please make sure the Ollama server is up & running.'
-                    logError('getLocalOllamaModels: failed ', isFetchFailed ? serverErrorMsg : error)
-                }
-            )
+        ModelProvider.ollamaProviders = enableOllamaModels ? await fetchLocalOllamaModels() : []
     }
 
     /**
@@ -131,3 +88,35 @@ export class ModelProvider {
         return tokensToChars(model?.maxToken || DEFAULT_FAST_MODEL_TOKEN_LIMIT)
     }
 }
+
+/**
+ * Fetches available Ollama models from the local Ollama server
+ * and adds them to the list of ollama providers.
+ */
+export async function fetchLocalOllamaModels(): Promise<ModelProvider[]> {
+    if (process.env.CODY_SHIM_TESTING === 'true') {
+        return []
+    }
+    // TODO (bee) watch file change to determine if a new model is added
+    // to eliminate the needs of restarting the extension to get the new models
+    return await fetch(new URL('/api/tags', OLLAMA_DEFAULT_URL).href)
+        .then(response => response.json())
+        .then(
+            data =>
+                data?.models?.map(
+                    (m: { model: string }) =>
+                        new ModelProvider(
+                            `ollama/${m.model}`,
+                            [ModelUsage.Chat, ModelUsage.Edit],
+                            DEFAULT_FAST_MODEL_CHARS_LIMIT
+                        )
+                ),
+            error => {
+                const fetchFailedErrors = ['Failed to fetch', 'fetch failed']
+                const isFetchFailed = fetchFailedErrors.some(err => error.toString().includes(err))
+                const serverErrorMsg = 'Please make sure the Ollama server is up & running.'
+                logError('getLocalOllamaModels: failed ', isFetchFailed ? serverErrorMsg : error)
+                return []
+            }
+        )
+}

From 8c3e6851694c4577a56d6730dd95bf00b9a965fe Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Thu, 21 Mar 2024 23:55:02 -0700
Subject: [PATCH 12/14] update bindings

---
 .../com/sourcegraph/cody/protocol_generated/ModelProvider.kt     | 1 -
 1 file changed, 1 deletion(-)

diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
index 7e448a9e25d..f120da78a04 100644
--- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
+++ b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
@@ -7,7 +7,6 @@ data class ModelProvider(
   val provider: String? = null,
   val title: String? = null,
   val primaryProviders: List<ModelProvider>? = null,
-  val ollamaProvidersEnabled: Boolean? = null,
   val ollamaProviders: List<ModelProvider>? = null,
 )
 

From 9da8bfe766b19187ba1704046d61b6fbcd72773d Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Fri, 22 Mar 2024 10:09:35 -0700
Subject: [PATCH 13/14] move sync to top level

---
 lib/shared/src/models/index.ts           | 79 +++++++-----------------
 lib/shared/src/models/utils.ts           | 33 ++++++++++
 vscode/src/chat/chat-view/ChatManager.ts |  3 -
 vscode/src/main.ts                       |  2 +
 vscode/src/models/utils.ts               |  4 +-
 5 files changed, 59 insertions(+), 62 deletions(-)

diff --git a/lib/shared/src/models/index.ts b/lib/shared/src/models/index.ts
index c13a947f036..2dfee9b31c0 100644
--- a/lib/shared/src/models/index.ts
+++ b/lib/shared/src/models/index.ts
@@ -1,13 +1,7 @@
-import { logError } from '../logger'
-import { OLLAMA_DEFAULT_URL } from '../ollama'
-import {
-    DEFAULT_FAST_MODEL_CHARS_LIMIT,
-    DEFAULT_FAST_MODEL_TOKEN_LIMIT,
-    tokensToChars,
-} from '../prompt/constants'
+import { DEFAULT_FAST_MODEL_TOKEN_LIMIT, tokensToChars } from '../prompt/constants'
 import { DEFAULT_DOT_COM_MODELS } from './dotcom'
-import { ModelUsage } from './types'
-import { getModelInfo } from './utils'
+import type { ModelUsage } from './types'
+import { fetchLocalOllamaModels, getModelInfo } from './utils'
 
 /**
  * ModelProvider manages available chat and edit models.
@@ -15,7 +9,9 @@ import { getModelInfo } from './utils'
  * retrieve and select between them.
  */
 export class ModelProvider {
+    // Whether the model is the default model
     public default = false
+    // Whether the model is only available to Pro users
     public codyProOnly = false
     // The name of the provider of the model, e.g. "Anthropic"
     public provider: string
@@ -38,16 +34,23 @@ export class ModelProvider {
     }
 
     /**
-     * Providers available on the user's instance
+     * Get all the providers currently available to the user
+     */
+    private static get providers(): ModelProvider[] {
+        return ModelProvider.primaryProviders.concat(ModelProvider.localProviders)
+    }
+    /**
+     * Providers available on the user's Sourcegraph instance
      */
     private static primaryProviders: ModelProvider[] = DEFAULT_DOT_COM_MODELS
     /**
-     * Providers available from local ollama instances
+     * Providers available from user's local instances, e.g. Ollama
      */
-    private static ollamaProviders: ModelProvider[] = []
+    private static localProviders: ModelProvider[] = []
+
     public static async onConfigChange(enableOllamaModels: boolean): Promise<void> {
         // Only fetch local models if user has enabled the config
-        ModelProvider.ollamaProviders = enableOllamaModels ? await fetchLocalOllamaModels() : []
+        ModelProvider.localProviders = enableOllamaModels ? await fetchLocalOllamaModels() : []
     }
 
     /**
@@ -63,17 +66,13 @@ export class ModelProvider {
      * If currentModel is provided, sets it as the default model.
      */
     public static getProviders(type: ModelUsage, currentModel?: string): ModelProvider[] {
-        const models = ModelProvider.primaryProviders
-            .concat(ModelProvider.ollamaProviders)
-            .filter(model => model.usage.includes(type))
-
-        // Set the current model as default
-        return models?.map(model => {
-            return {
+        return ModelProvider.providers
+            .filter(m => m.usage.includes(type))
+            ?.map(model => ({
                 ...model,
+                // Set the current model as default
                 default: model.model === currentModel,
-            }
-        })
+            }))
     }
 
     /**
@@ -82,41 +81,7 @@ export class ModelProvider {
      * E.g. 7000 tokens * 4 characters/token = 28000 characters
      */
     public static getMaxCharsByModel(modelID: string): number {
-        const model = ModelProvider.primaryProviders
-            .concat(ModelProvider.ollamaProviders)
-            .find(m => m.model === modelID)
+        const model = ModelProvider.providers.find(m => m.model === modelID)
         return tokensToChars(model?.maxToken || DEFAULT_FAST_MODEL_TOKEN_LIMIT)
     }
 }
-
-/**
- * Fetches available Ollama models from the local Ollama server
- * and adds them to the list of ollama providers.
- */
-export async function fetchLocalOllamaModels(): Promise<ModelProvider[]> {
-    if (process.env.CODY_SHIM_TESTING === 'true') {
-        return []
-    }
-    // TODO (bee) watch file change to determine if a new model is added
-    // to eliminate the needs of restarting the extension to get the new models
-    return await fetch(new URL('/api/tags', OLLAMA_DEFAULT_URL).href)
-        .then(response => response.json())
-        .then(
-            data =>
-                data?.models?.map(
-                    (m: { model: string }) =>
-                        new ModelProvider(
-                            `ollama/${m.model}`,
-                            [ModelUsage.Chat, ModelUsage.Edit],
-                            DEFAULT_FAST_MODEL_CHARS_LIMIT
-                        )
-                ),
-            error => {
-                const fetchFailedErrors = ['Failed to fetch', 'fetch failed']
-                const isFetchFailed = fetchFailedErrors.some(err => error.toString().includes(err))
-                const serverErrorMsg = 'Please make sure the Ollama server is up & running.'
-                logError('getLocalOllamaModels: failed ', isFetchFailed ? serverErrorMsg : error)
-                return []
-            }
-        )
-}
diff --git a/lib/shared/src/models/utils.ts b/lib/shared/src/models/utils.ts
index e1099eb62cd..1f32850b92b 100644
--- a/lib/shared/src/models/utils.ts
+++ b/lib/shared/src/models/utils.ts
@@ -1,3 +1,8 @@
+import { ModelProvider } from '.'
+import { logError } from '../logger'
+import { OLLAMA_DEFAULT_URL } from '../ollama'
+import { DEFAULT_FAST_MODEL_CHARS_LIMIT } from '../prompt/constants'
+import { ModelUsage } from './types'
 export function getProviderName(name: string): string {
     const providerName = name.toLowerCase()
     switch (providerName) {
@@ -28,3 +33,31 @@ export function getModelInfo(modelID: string): {
     const title = (rest.at(-1) || '').replace(/-/g, ' ')
     return { provider, title }
 }
+
+/**
+ * Fetches available Ollama models from the Ollama server.
+ */
+export async function fetchLocalOllamaModels(): Promise<ModelProvider[]> {
+    // TODO (bee) watch file change to determine if a new model is added
+    // to eliminate the needs of restarting the extension to get the new models
+    return await fetch(new URL('/api/tags', OLLAMA_DEFAULT_URL).href)
+        .then(response => response.json())
+        .then(
+            data =>
+                data?.models?.map(
+                    (m: { model: string }) =>
+                        new ModelProvider(
+                            `ollama/${m.model}`,
+                            [ModelUsage.Chat, ModelUsage.Edit],
+                            DEFAULT_FAST_MODEL_CHARS_LIMIT
+                        )
+                ),
+            error => {
+                const fetchFailedErrors = ['Failed to fetch', 'fetch failed']
+                const isFetchFailed = fetchFailedErrors.some(err => error.toString().includes(err))
+                const serverErrorMsg = 'Please make sure the Ollama server is up & running.'
+                logError('getLocalOllamaModels: failed ', isFetchFailed ? serverErrorMsg : error)
+                return []
+            }
+        )
+}
diff --git a/vscode/src/chat/chat-view/ChatManager.ts b/vscode/src/chat/chat-view/ChatManager.ts
index 002dfb9a932..b6453ae9243 100644
--- a/vscode/src/chat/chat-view/ChatManager.ts
+++ b/vscode/src/chat/chat-view/ChatManager.ts
@@ -21,7 +21,6 @@ import { telemetryRecorder } from '../../services/telemetry-v2'
 import type { ExecuteChatArguments } from '../../commands/execute/ask'
 import type { EnterpriseContextFactory } from '../../context/enterprise-context-factory'
 import type { ContextRankingController } from '../../local-context/context-ranking'
-import { setModelProviders } from '../../models/utils'
 import { ChatPanelsManager } from './ChatPanelsManager'
 import { SidebarViewController, type SidebarViewOptions } from './SidebarViewController'
 import type { ChatSession, SimpleChatPanelProvider } from './SimpleChatPanelProvider'
@@ -99,8 +98,6 @@ export class ChatManager implements vscode.Disposable {
     }
 
     public async syncAuthStatus(authStatus: AuthStatus): Promise<void> {
-        setModelProviders(authStatus)
-
         await this.chatPanelsManager.syncAuthStatus(authStatus)
     }
 
diff --git a/vscode/src/main.ts b/vscode/src/main.ts
index e2d349bcfaa..81073200214 100644
--- a/vscode/src/main.ts
+++ b/vscode/src/main.ts
@@ -43,6 +43,7 @@ import { VSCodeEditor } from './editor/vscode-editor'
 import type { PlatformContext } from './extension.common'
 import { configureExternalServices } from './external-services'
 import { logDebug, logError } from './log'
+import { syncModelProviders } from './models/utils'
 import type { FixupTask } from './non-stop/FixupTask'
 import { CodyProExpirationNotifications } from './notifications/cody-pro-expiration'
 import { showSetupNotification } from './notifications/setup-notification'
@@ -279,6 +280,7 @@ const register = async (
 
     // Adds a change listener to the auth provider that syncs the auth status
     authProvider.addChangeListener(async (authStatus: AuthStatus) => {
+        syncModelProviders(authStatus)
         // Chat Manager uses Simple Context Provider
         await chatManager.syncAuthStatus(authStatus)
         editorManager.syncAuthStatus(authStatus)
diff --git a/vscode/src/models/utils.ts b/vscode/src/models/utils.ts
index 45f3891ef75..74a556ee6fb 100644
--- a/vscode/src/models/utils.ts
+++ b/vscode/src/models/utils.ts
@@ -10,7 +10,7 @@ import * as vscode from 'vscode'
  * The token limit for the provider will use the configured limit,
  * or fallback to the limit from the authentication status if not configured.
  */
-export function setModelProviders(authStatus: AuthStatus): void {
+export function syncModelProviders(authStatus: AuthStatus): void {
     if (authStatus.endpoint && isDotCom(authStatus.endpoint)) {
         ModelProvider.setProviders(DEFAULT_DOT_COM_MODELS)
         return
@@ -24,7 +24,7 @@ export function setModelProviders(authStatus: AuthStatus): void {
     // the client.
     if (authStatus?.configOverwrites?.chatModel) {
         const codyConfig = vscode.workspace.getConfiguration('cody')
-        const tokenLimitConfig = codyConfig.get<number>('provider.limit.prompt')
+        const tokenLimitConfig = codyConfig?.get<number>('provider.limit.prompt')
         const tokenLimit = tokenLimitConfig ?? authStatus.configOverwrites?.chatModelMaxTokens
         ModelProvider.setProviders([
             new ModelProvider(

From 93386b45791d7a49d691c6c1da303a9662ba99ab Mon Sep 17 00:00:00 2001
From: Beatrix <beatrix@sourcegraph.com>
Date: Fri, 22 Mar 2024 10:16:43 -0700
Subject: [PATCH 14/14] Update agent binding

---
 .../com/sourcegraph/cody/protocol_generated/ModelProvider.kt    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
index f120da78a04..a7ecee298f4 100644
--- a/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
+++ b/agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ModelProvider.kt
@@ -7,6 +7,6 @@ data class ModelProvider(
   val provider: String? = null,
   val title: String? = null,
   val primaryProviders: List<ModelProvider>? = null,
-  val ollamaProviders: List<ModelProvider>? = null,
+  val localProviders: List<ModelProvider>? = null,
 )