fix: proper context window loading for LMStudio (fixes #5075)

Author: pwilkin

src/api/providers/fetchers/lm-studio.ts

@@ -0,0 +1,35 @@
+import axios from "axios"
+import { ModelRecord } from "../../../shared/api"
+import { openAiModelInfoSaneDefaults } from "@roo-code/types"
+
+export async function getLmStudioModels(baseUrl = "http://localhost:1234"): Promise<ModelRecord> {
+	try {
+		if (!URL.canParse(baseUrl)) {
+			return {}
+		}
+
+		const response = await axios.get(`${baseUrl}/api/v0/models`)
+		return response.data?.data?.reduce((acc: ModelRecord, model: any) => {
+			acc[model.id] = {
+				maxTokens:
+					model.loaded_context_length ||
+					model.max_context_length ||
+					openAiModelInfoSaneDefaults.contextWindow,
+				contextWindow:
+					model.loaded_context_length ||
+					model.max_context_length ||
+					openAiModelInfoSaneDefaults.contextWindow,
+				supportsImages: false,
+				supportsPromptCache: false,
+				supportsComputerUse: false,
+				inputPrice: 0,
+				outputPrice: 0,
+				cacheWritesPrice: 0,
+				cacheReadsPrice: 0,
+			}
+			return acc
+		}, {})
+	} catch (error) {
+		return {}
+	}
+}

src/api/providers/fetchers/modelCache.ts

@@ -14,6 +14,7 @@ import { getRequestyModels } from "./requesty"
 import { getGlamaModels } from "./glama"
 import { getUnboundModels } from "./unbound"
 import { getLiteLLMModels } from "./litellm"
+import { getLmStudioModels } from "./lm-studio"
 import { GetModelsOptions } from "../../../shared/api"
 import { getOllamaModels } from "./ollama"
 import { getLMStudioModels } from "./lmstudio"

src/api/providers/lm-studio.ts

@@ -1,8 +1,7 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import OpenAI from "openai"
-import axios from "axios"
 
-import { type ModelInfo, openAiModelInfoSaneDefaults, LMSTUDIO_DEFAULT_TEMPERATURE } from "@roo-code/types"
+import { LMSTUDIO_DEFAULT_TEMPERATURE, type ModelInfo, openAiModelInfoSaneDefaults } from "@roo-code/types"
 
 import type { ApiHandlerOptions } from "../../shared/api"
 
@@ -11,22 +10,32 @@ import { XmlMatcher } from "../../utils/xml-matcher"
 import { convertToOpenAiMessages } from "../transform/openai-format"
 import { ApiStream } from "../transform/stream"
 
+import type { ApiHandlerCreateMessageMetadata, SingleCompletionHandler } from "../index"
 import { BaseProvider } from "./base-provider"
-import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
+import { flushModels, getModels } from "./fetchers/modelCache"
 
 export class LmStudioHandler extends BaseProvider implements SingleCompletionHandler {
 	protected options: ApiHandlerOptions
 	private client: OpenAI
+	private cachedModelInfo: ModelInfo = openAiModelInfoSaneDefaults
 
 	constructor(options: ApiHandlerOptions) {
 		super()
 		this.options = options
 		this.client = new OpenAI({
-			baseURL: (this.options.lmStudioBaseUrl || "http://localhost:1234") + "/v1",
+			baseURL: this.getBaseUrl() + "/v1",
 			apiKey: "noop",
 		})
 	}
 
+	private getBaseUrl(): string {
+		if (this.options.lmStudioBaseUrl && this.options.lmStudioBaseUrl.trim() !== "") {
+			return this.options.lmStudioBaseUrl.trim()
+		} else {
+			return "http://localhost:1234"
+		}
+	}
+
 	override async *createMessage(
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
@@ -118,6 +127,23 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
 				outputTokens = 0
 			}
 
+			if (this.cachedModelInfo === openAiModelInfoSaneDefaults) {
+				// We need to fetch the model info every time we open a new session
+				// to ensure we have the latest context window and other details
+				// since LM Studio models can chance their context windows on reload
+				await flushModels("lmstudio")
+				const models = await getModels({ provider: "lmstudio", baseUrl: this.getBaseUrl() })
+				if (models && models[this.getModel().id]) {
+					this.cachedModelInfo = models[this.getModel().id]
+				} else {
+					// If model info is not found, use sane defaults
+					this.cachedModelInfo = {
+						...openAiModelInfoSaneDefaults,
+						description: "Fake description to avoid recache",
+					}
+				}
+			}
+
 			yield {
 				type: "usage",
 				inputTokens,
@@ -133,7 +159,7 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
 	override getModel(): { id: string; info: ModelInfo } {
 		return {
 			id: this.options.lmStudioModelId || "",
-			info: openAiModelInfoSaneDefaults,
+			info: this.cachedModelInfo,
 		}
 	}
 
@@ -161,17 +187,3 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan
 		}
 	}
 }
-
-export async function getLmStudioModels(baseUrl = "http://localhost:1234") {
-	try {
-		if (!URL.canParse(baseUrl)) {
-			return []
-		}
-
-		const response = await axios.get(`${baseUrl}/v1/models`)
-		const modelsArray = response.data?.data?.map((model: any) => model.id) || []
-		return [...new Set<string>(modelsArray)]
-	} catch (error) {
-		return []
-	}
-}

src/core/webview/webviewMessageHandler.ts

@@ -1,58 +1,56 @@
-import { safeWriteJson } from "../../utils/safeWriteJson"
-import * as path from "path"
-import * as os from "os"
 import * as fs from "fs/promises"
+import * as os from "os"
 import pWaitFor from "p-wait-for"
+import * as path from "path"
 import * as vscode from "vscode"
-import * as yaml from "yaml"
+import { safeWriteJson } from "../../utils/safeWriteJson"
 
+import { CloudService } from "@roo-code/cloud"
+import { TelemetryService } from "@roo-code/telemetry"
 import {
+	type ClineMessage,
+	type GlobalState,
 	type Language,
 	type ProviderSettings,
-	type GlobalState,
-	type ClineMessage,
 	TelemetryEventName,
 } from "@roo-code/types"
-import { CloudService } from "@roo-code/cloud"
-import { TelemetryService } from "@roo-code/telemetry"
 import { type ApiMessage } from "../task-persistence/apiMessages"
 
-import { ClineProvider } from "./ClineProvider"
 import { changeLanguage, t } from "../../i18n"
+import { ModelRecord, RouterName, toRouterName } from "../../shared/api"
 import { Package } from "../../shared/package"
-import { RouterName, toRouterName, ModelRecord } from "../../shared/api"
 import { supportPrompt } from "../../shared/support-prompt"
+import { ClineProvider } from "./ClineProvider"
 
-import { checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, WebviewMessage } from "../../shared/WebviewMessage"
-import { checkExistKey } from "../../shared/checkExistApiConfig"
-import { experimentDefault } from "../../shared/experiments"
-import { Terminal } from "../../integrations/terminal/Terminal"
-import { openFile } from "../../integrations/misc/open-file"
+import { flushModels, getModels } from "../../api/providers/fetchers/modelCache"
+import { getOpenAiModels } from "../../api/providers/openai"
+import { getVsCodeLmModels } from "../../api/providers/vscode-lm"
 import { openImage, saveImage } from "../../integrations/misc/image-handler"
+import { openFile } from "../../integrations/misc/open-file"
 import { selectImages } from "../../integrations/misc/process-images"
+import { Terminal } from "../../integrations/terminal/Terminal"
 import { getTheme } from "../../integrations/theme/getTheme"
 import { discoverChromeHostUrl, tryChromeHostUrl } from "../../services/browser/browserDiscovery"
 import { searchWorkspaceFiles } from "../../services/search/file-search"
+import { TelemetrySetting } from "../../shared/TelemetrySetting"
+import { checkoutDiffPayloadSchema, checkoutRestorePayloadSchema, WebviewMessage } from "../../shared/WebviewMessage"
+import { GetModelsOptions } from "../../shared/api"
+import { checkExistKey } from "../../shared/checkExistApiConfig"
+import { experimentDefault } from "../../shared/experiments"
+import { defaultModeSlug, Mode } from "../../shared/modes"
+import { getCommand } from "../../utils/commands"
 import { fileExistsAtPath } from "../../utils/fs"
-import { playTts, setTtsEnabled, setTtsSpeed, stopTts } from "../../utils/tts"
-import { singleCompletionHandler } from "../../utils/single-completion-handler"
 import { searchCommits } from "../../utils/git"
+import { getWorkspacePath } from "../../utils/path"
+import { singleCompletionHandler } from "../../utils/single-completion-handler"
+import { playTts, setTtsEnabled, setTtsSpeed, stopTts } from "../../utils/tts"
 import { exportSettings, importSettingsWithFeedback } from "../config/importExport"
-import { getOpenAiModels } from "../../api/providers/openai"
-import { getVsCodeLmModels } from "../../api/providers/vscode-lm"
 import { openMention } from "../mentions"
-import { TelemetrySetting } from "../../shared/TelemetrySetting"
-import { getWorkspacePath } from "../../utils/path"
-import { ensureSettingsDirectoryExists } from "../../utils/globalContext"
-import { Mode, defaultModeSlug } from "../../shared/modes"
-import { getModels, flushModels } from "../../api/providers/fetchers/modelCache"
-import { GetModelsOptions } from "../../shared/api"
 import { generateSystemPrompt } from "./generateSystemPrompt"
-import { getCommand } from "../../utils/commands"
 
 const ALLOWED_VSCODE_SETTINGS = new Set(["terminal.integrated.inheritEnv"])
 
-import { MarketplaceManager, MarketplaceItemType } from "../../services/marketplace"
+import { MarketplaceItemType, MarketplaceManager } from "../../services/marketplace"
 import { setPendingTodoList } from "../tools/updateTodoListTool"
 
 export const webviewMessageHandler = async (
@@ -555,6 +553,12 @@ export const webviewMessageHandler = async (
 				})
 			}
 
+			const lmStudioBaseUrl = apiConfiguration.lmStudioBaseUrl || message?.values?.lmStudioBaseUrl
+			modelFetchPromises.push({
+				key: "lmstudio",
+				options: { provider: "lmstudio", baseUrl: lmStudioBaseUrl },
+			})
+
 			const results = await Promise.allSettled(
 				modelFetchPromises.map(async ({ key, options }) => {
 					const models = await safeGetModels(options)
@@ -633,30 +637,6 @@ export const webviewMessageHandler = async (
 			}
 			break
 		}
-		case "requestLmStudioModels": {
-			// Specific handler for LM Studio models only
-			const { apiConfiguration: lmStudioApiConfig } = await provider.getState()
-			try {
-				// Flush cache first to ensure fresh models
-				await flushModels("lmstudio")
-
-				const lmStudioModels = await getModels({
-					provider: "lmstudio",
-					baseUrl: lmStudioApiConfig.lmStudioBaseUrl,
-				})
-
-				if (Object.keys(lmStudioModels).length > 0) {
-					provider.postMessageToWebview({
-						type: "lmStudioModels",
-						lmStudioModels: Object.keys(lmStudioModels),
-					})
-				}
-			} catch (error) {
-				// Silently fail - user hasn't configured LM Studio yet
-				console.debug("LM Studio models fetch failed:", error)
-			}
-			break
-		}
 		case "requestOpenAiModels":
 			if (message?.values?.baseUrl && message?.values?.apiKey) {
 				const openAiModels = await getOpenAiModels(

src/shared/WebviewMessage.ts

@@ -65,7 +65,6 @@ export interface WebviewMessage {
 		| "requestRouterModels"
 		| "requestOpenAiModels"
 		| "requestOllamaModels"
-		| "requestLmStudioModels"
 		| "requestVsCodeLmModels"
 		| "openImage"
 		| "saveImage"

webview-ui/src/components/settings/ApiOptions.tsx

@@ -193,12 +193,10 @@ const ApiOptions = ({
 					},
 				})
 			} else if (selectedProvider === "ollama") {
-				vscode.postMessage({ type: "requestOllamaModels" })
-			} else if (selectedProvider === "lmstudio") {
-				vscode.postMessage({ type: "requestLmStudioModels" })
+				vscode.postMessage({ type: "requestOllamaModels", text: apiConfiguration?.ollamaBaseUrl })
 			} else if (selectedProvider === "vscode-lm") {
 				vscode.postMessage({ type: "requestVsCodeLmModels" })
-			} else if (selectedProvider === "litellm") {
+			} else if (selectedProvider === "litellm" || selectedProvider === "lmstudio") {
 				vscode.postMessage({ type: "requestRouterModels" })
 			}
 		},

webview-ui/src/components/settings/providers/LMStudio.tsx

@@ -39,9 +39,9 @@ export const LMStudio = ({ apiConfiguration, setApiConfigurationField }: LMStudi
 		const message: ExtensionMessage = event.data
 
 		switch (message.type) {
-			case "lmStudioModels":
+			case "routerModels":
 				{
-					const newModels = message.lmStudioModels ?? []
+					const newModels = Object.keys(message.routerModels?.lmstudio || {})
 					setLmStudioModels(newModels)
 				}
 				break
@@ -53,7 +53,7 @@ export const LMStudio = ({ apiConfiguration, setApiConfigurationField }: LMStudi
 	// Refresh models on mount
 	useEffect(() => {
 		// Request fresh models - the handler now flushes cache automatically
-		vscode.postMessage({ type: "requestLmStudioModels" })
+		vscode.postMessage({ type: "requestRouterModels" })
 	}, [])
 
 	// Check if the selected model exists in the fetched models